<commit_before>/// Uncompressed, absolute domain names.

use std::{cmp, fmt, hash, ops};
use std::ascii::AsciiExt;
use bytes::{BufMut, Bytes};
use ::bits::compose::Composable;
use ::bits::parse::{Parseable, Parser};
use super::error::{DnameError, IndexError, RootNameError};
use super::label::Label;
use super::relname::{RelativeDname, DnameIter};
use super::traits::{ToLabelIter, ToDname, ToRelativeDname};

//------------ Dname ---------------------------------------------------------

/// An uncompressed, absolute domain name.
///
/// The type wraps a [`Bytes`] value and guarantees that it always contains
/// a correctly encoded, absolute domain name. It derefs to [`Bytes`] and
/// therefore to `[u8]` allowing you direct access to the underlying byte
/// slice. It does overide all applicable methods providing access to parts
/// of the byte slice, though, returning either `Dname` or [`RelativeDname`]s
/// instead.
///
/// You can construct a domain name from a string via the `FromStr` trait or
/// manually via a [`DnameBuilder`]. In addition, you can also parse it from
/// a message. This will, however, require the name to be uncompressed.
///
/// [`Bytes`]: ../../../bytes/struct.Bytes.html
/// [`DnameBuilder`]: struct.DnameBuilder.html
#[derive(Clone)]
pub struct Dname {
    bytes: Bytes
}

/// # Creation and Conversion
///
impl Dname {
    /// Creates a domain name from the underlying bytes without any check.
    ///
    /// Since this will allow to actually construct an incorrectly encoded
    /// domain name value, the function is unsafe.
    pub(super) unsafe fn from_bytes_unchecked(bytes: Bytes) -> Self {
        Dname { bytes }
    }

    /// Creates a domain name representing the root.
    ///
    /// The resulting domain name will consist of the root label only.
    pub fn root() -> Self {
        unsafe { Self::from_bytes_unchecked(Bytes::from_static(b"\0")) }
    }

    /// Creates a domain name from a bytes value.
    ///
    /// This will only succeed if `bytes` contains a properly encoded
    /// absolute domain name. Because the function checks, this will take
    /// a wee bit of time.
    pub fn from_bytes(_bytes: Bytes) -> Result<Self, DnameError> {
        unimplemented!()
    }

    /// Returns a reference to the underlying bytes value.
    pub fn as_bytes(&self) -> &Bytes {
        &self.bytes
    }

    /// Returns a reference to the underlying byte slice.
    pub fn as_slice(&self) -> &[u8] {
        self.bytes.as_ref()
    }

    /// Converts the domain name into its underlying bytes slice.
    pub fn into_bytes(self) -> Bytes {
        self.bytes
    }

    /// Converts the name into a relative name by dropping the root label.
    pub fn into_relative(mut self) -> RelativeDname {
        let len = self.bytes.len() - 1;
        self.bytes.truncate(len);
        unsafe { RelativeDname::from_bytes_unchecked(self.bytes) }
    }
}

/// # Working with Labels
///
impl Dname {
    /// Returns an iterator over the labels of the domain name.
    pub fn iter(&self) -> DnameIter {
        DnameIter::new(self.bytes.as_ref())
    }

    /// Returns the number of labels in the domain name.
    pub fn label_count(&self) -> usize {
        self.iter().count()
    }

    /// Returns a reference to the first label.
    pub fn first(&self) -> &Label {
        self.iter().next().unwrap()
    }

    /// Returns a reference to the last label.
    pub fn last(&self) -> &Label {
        self.iter().next_back().unwrap()
    }

    /// Determines whether `base` is a prefix of `self`.
    ///
    /// As this methods accepts only relative domain names, it will only
    /// allow checking for a ‘strict’ prefix.
    pub fn starts_with<N: ToRelativeDname>(&self, base: &N) -> bool {
        <Self as ToLabelIter>::starts_with(self, base)
    }

    /// Determines whether `base` is a suffix of `self`.
    pub fn ends_with<N: ToDname>(&self, base: &N) -> bool {
        <Self as ToLabelIter>::ends_with(self, base)
    }

    /// Returns the part of the name indicated by start and end positions.
    ///
    /// The returned name will start at position `begin` and end right before
    /// position `end`. Both positions must point to the begining of a label
    /// or an error will be returned.
    ///
    /// Because the returned domain is a relative name, the method will also
    /// return an error if the end equal to the length of the name. If you
    /// want to slice the entire end of the name including the final root
    /// label, you can use [`slice_from()`] instead.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    ///
    /// [`slice_from()`]: #method.slice_from
    pub fn slice(&self, begin: usize, end: usize)
                 -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, begin)?;
        IndexError::check(&self.bytes, end)?;
        if end == self.len() {
            return Err(IndexError)
        }
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.slice(begin, end))
        })
    }

    /// Returns the part of the name starting at the given position.
    ///
    /// This will fail if the position isn’t the start of a label.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    pub fn slice_from(&self, begin: usize) -> Result<Self, IndexError> {
        IndexError::check(&self.bytes, begin)?;
        Ok(unsafe {
            Self::from_bytes_unchecked(self.bytes.slice_from(begin))
        })
    }

    /// Returns the part of the name ending at the given position.
    ///
    /// This will fail if the position isn’t the start of a label.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    pub fn slice_to(&self, end: usize) -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, end)?;
        if end == self.len() {
            return Err(IndexError)
        }
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.slice_to(end))
        })
    }

    // XXX No `split_off()` since that would require `self` to mysteriously
    //     change into a `RelativeDname`. Would could make this move `self`,
    //     but then you would loose it upon an error which is not nice,
    //     either.

    /// Splits the name into two at the given position.
    ///
    /// Afterwards, `self` will contain the name starting at the position
    /// while the name ending right before it will be returned. The method
    /// will fail if `mid` is not the start of a new label.
    ///
    /// # Panics
    ///
    /// The method will panic if `mid` is greater than the name’s length.
    pub fn split_to(&mut self, mid: usize)
                    -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, mid)?;
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.split_to(mid))
        })
    }

    // XXX No `truncate()` either.

    /// Splits off the first label.
    ///
    /// If this name is longer than just the root label, returns the first
    /// label as a relative name and removes it from the name itself. If the
    /// name is only the root label, returns an error and does nothing.
    pub fn split_first(&mut self) -> Result<RelativeDname, RootNameError> {
        if self.len() == 1 {
            return Err(RootNameError)
        }
        let end = self.iter().next().unwrap().len() + 1;
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.split_to(end))
        })
    }

    /// Reduces the name to the parent of the current name.
    ///
    /// This will fail if the name consists of the root label only.
    pub fn parent(&mut self) -> Result<(), RootNameError> {
        self.split_first().map(|_| ())
    }

    // XXX And no `strip_suffix()`.
}

//--- Parseable and Composable

impl Parseable for Dname {
    type Err = DnameError;

    fn parse(parser: &mut Parser) -> Result<Self, DnameError> {
        let len = {
            let mut tmp = parser.peek();
            loop {
                if tmp.is_empty() {
                    return Err(DnameError::ShortData)
                }
                let (label, tail) = Label::split_from(tmp)?;
                tmp = tail;
                if label.is_root() {
                    break;
                }
            }
            parser.remaining() - tmp.len()
        };
        if len > 255 {
            return Err(DnameError::LongName);
        }
        Ok(unsafe {
            Self::from_bytes_unchecked(parser.parse_bytes(len).unwrap())
        })
    }
}

impl Composable for Dname {
    fn compose_len(&self) -> usize {
        self.bytes.len()
    }

    fn compose<B: BufMut>(&self, buf: &mut B) {
        buf.put_slice(self.as_ref())
    }
}

//--- ToLabelIter and ToDname

impl<'a> ToLabelIter<'a> for Dname {
    type LabelIter = DnameIter<'a>;

    fn iter_labels(&'a self) -> Self::LabelIter {
        self.iter()
    }
}

impl ToDname for Dname {
    fn to_name(&self) -> Dname {
        self.clone()
    }
}

//--- Deref and AsRef

impl ops::Deref for Dname {
    type Target = Bytes;

    fn deref(&self) -> &Bytes {
        self.as_ref()
    }
}

impl AsRef<Bytes> for Dname {
    fn as_ref(&self) -> &Bytes {
        &self.bytes
    }
}

impl AsRef<[u8]> for Dname {
    fn as_ref(&self) -> &[u8] {
        self.bytes.as_ref()
    }
}

//--- IntoIterator

impl<'a> IntoIterator for &'a Dname {
    type Item = &'a Label;
    type IntoIter = DnameIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

//--- PartialEq and Eq

impl PartialEq for Dname {
    fn eq(&self, other: &Self) -> bool {
        self.as_slice().eq_ignore_ascii_case(other.as_slice())
    }
}

impl Eq for Dname { }

//--- PartialOrd and Ord

impl PartialOrd for Dname {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.iter().partial_cmp(other.iter())
    }
}

impl Ord for Dname {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.iter().cmp(other.iter())
    }
}

//--- Hash

impl hash::Hash for Dname {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        for item in self.iter() {
            item.hash(state)
        }
    }
}

//--- Display and Debug

impl fmt::Display for Dname {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut iter = self.iter();
        match iter.next() {
            Some(label) => label.fmt(f)?,
            None => return Ok(())
        }
        while let Some(label) = iter.next() {
            f.write_str(".")?;
            label.fmt(f)?;
        }
        Ok(())
    }
}

impl fmt::Debug for Dname {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Dname({})", self)
    }
}

<commit_msg>Implement Dname::from_bytes.<commit_after>/// Uncompressed, absolute domain names.

use std::{cmp, fmt, hash, ops};
use std::ascii::AsciiExt;
use bytes::{BufMut, Bytes};
use ::bits::compose::Composable;
use ::bits::parse::{Parseable, Parser};
use super::error::{DnameError, IndexError, RootNameError};
use super::label::Label;
use super::relname::{RelativeDname, DnameIter};
use super::traits::{ToLabelIter, ToDname, ToRelativeDname};

//------------ Dname ---------------------------------------------------------

/// An uncompressed, absolute domain name.
///
/// The type wraps a [`Bytes`] value and guarantees that it always contains
/// a correctly encoded, absolute domain name. It derefs to [`Bytes`] and
/// therefore to `[u8]` allowing you direct access to the underlying byte
/// slice. It does overide all applicable methods providing access to parts
/// of the byte slice, though, returning either `Dname` or [`RelativeDname`]s
/// instead.
///
/// You can construct a domain name from a string via the `FromStr` trait or
/// manually via a [`DnameBuilder`]. In addition, you can also parse it from
/// a message. This will, however, require the name to be uncompressed.
///
/// [`Bytes`]: ../../../bytes/struct.Bytes.html
/// [`DnameBuilder`]: struct.DnameBuilder.html
#[derive(Clone)]
pub struct Dname {
    bytes: Bytes
}

/// # Creation and Conversion
///
impl Dname {
    /// Creates a domain name from the underlying bytes without any check.
    ///
    /// Since this will allow to actually construct an incorrectly encoded
    /// domain name value, the function is unsafe.
    pub(super) unsafe fn from_bytes_unchecked(bytes: Bytes) -> Self {
        Dname { bytes }
    }

    /// Creates a domain name representing the root.
    ///
    /// The resulting domain name will consist of the root label only.
    pub fn root() -> Self {
        unsafe { Self::from_bytes_unchecked(Bytes::from_static(b"\0")) }
    }

    /// Creates a domain name from a bytes value.
    ///
    /// This will only succeed if `bytes` contains a properly encoded
    /// absolute domain name. Because the function checks, this will take
    /// a wee bit of time.
    pub fn from_bytes(bytes: Bytes) -> Result<Self, DnameError> {
        if bytes.len() > 255 {
            return Err(DnameError::LongName);
        }
        {
            let mut tmp = bytes.as_ref();
            loop {
                let (label, tail) = Label::split_from(tmp)?;
                if label.is_root() {
                    if tail.is_empty() {
                        break;
                    }
                    else {
                        return Err(DnameError::TrailingData)
                    }
                }
                if tail.is_empty() {
                    return Err(DnameError::RelativeName)
                }
                tmp = tail;
            }
        }
        Ok(unsafe { Dname::from_bytes_unchecked(bytes) })
    }

    /// Returns a reference to the underlying bytes value.
    pub fn as_bytes(&self) -> &Bytes {
        &self.bytes
    }

    /// Returns a reference to the underlying byte slice.
    pub fn as_slice(&self) -> &[u8] {
        self.bytes.as_ref()
    }

    /// Converts the domain name into its underlying bytes slice.
    pub fn into_bytes(self) -> Bytes {
        self.bytes
    }

    /// Converts the name into a relative name by dropping the root label.
    pub fn into_relative(mut self) -> RelativeDname {
        let len = self.bytes.len() - 1;
        self.bytes.truncate(len);
        unsafe { RelativeDname::from_bytes_unchecked(self.bytes) }
    }
}

/// # Working with Labels
///
impl Dname {
    /// Returns an iterator over the labels of the domain name.
    pub fn iter(&self) -> DnameIter {
        DnameIter::new(self.bytes.as_ref())
    }

    /// Returns the number of labels in the domain name.
    pub fn label_count(&self) -> usize {
        self.iter().count()
    }

    /// Returns a reference to the first label.
    pub fn first(&self) -> &Label {
        self.iter().next().unwrap()
    }

    /// Returns a reference to the last label.
    pub fn last(&self) -> &Label {
        self.iter().next_back().unwrap()
    }

    /// Determines whether `base` is a prefix of `self`.
    ///
    /// As this methods accepts only relative domain names, it will only
    /// allow checking for a ‘strict’ prefix.
    pub fn starts_with<N: ToRelativeDname>(&self, base: &N) -> bool {
        <Self as ToLabelIter>::starts_with(self, base)
    }

    /// Determines whether `base` is a suffix of `self`.
    pub fn ends_with<N: ToDname>(&self, base: &N) -> bool {
        <Self as ToLabelIter>::ends_with(self, base)
    }

    /// Returns the part of the name indicated by start and end positions.
    ///
    /// The returned name will start at position `begin` and end right before
    /// position `end`. Both positions must point to the begining of a label
    /// or an error will be returned.
    ///
    /// Because the returned domain is a relative name, the method will also
    /// return an error if the end equal to the length of the name. If you
    /// want to slice the entire end of the name including the final root
    /// label, you can use [`slice_from()`] instead.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    ///
    /// [`slice_from()`]: #method.slice_from
    pub fn slice(&self, begin: usize, end: usize)
                 -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, begin)?;
        IndexError::check(&self.bytes, end)?;
        if end == self.len() {
            return Err(IndexError)
        }
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.slice(begin, end))
        })
    }

    /// Returns the part of the name starting at the given position.
    ///
    /// This will fail if the position isn’t the start of a label.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    pub fn slice_from(&self, begin: usize) -> Result<Self, IndexError> {
        IndexError::check(&self.bytes, begin)?;
        Ok(unsafe {
            Self::from_bytes_unchecked(self.bytes.slice_from(begin))
        })
    }

    /// Returns the part of the name ending at the given position.
    ///
    /// This will fail if the position isn’t the start of a label.
    ///
    /// # Panics
    ///
    /// The method panics if either position points beyond the end of the
    /// name.
    pub fn slice_to(&self, end: usize) -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, end)?;
        if end == self.len() {
            return Err(IndexError)
        }
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.slice_to(end))
        })
    }

    // XXX No `split_off()` since that would require `self` to mysteriously
    //     change into a `RelativeDname`. Would could make this move `self`,
    //     but then you would loose it upon an error which is not nice,
    //     either.

    /// Splits the name into two at the given position.
    ///
    /// Afterwards, `self` will contain the name starting at the position
    /// while the name ending right before it will be returned. The method
    /// will fail if `mid` is not the start of a new label.
    ///
    /// # Panics
    ///
    /// The method will panic if `mid` is greater than the name’s length.
    pub fn split_to(&mut self, mid: usize)
                    -> Result<RelativeDname, IndexError> {
        IndexError::check(&self.bytes, mid)?;
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.split_to(mid))
        })
    }

    // XXX No `truncate()` either.

    /// Splits off the first label.
    ///
    /// If this name is longer than just the root label, returns the first
    /// label as a relative name and removes it from the name itself. If the
    /// name is only the root label, returns an error and does nothing.
    pub fn split_first(&mut self) -> Result<RelativeDname, RootNameError> {
        if self.len() == 1 {
            return Err(RootNameError)
        }
        let end = self.iter().next().unwrap().len() + 1;
        Ok(unsafe {
            RelativeDname::from_bytes_unchecked(self.bytes.split_to(end))
        })
    }

    /// Reduces the name to the parent of the current name.
    ///
    /// This will fail if the name consists of the root label only.
    pub fn parent(&mut self) -> Result<(), RootNameError> {
        self.split_first().map(|_| ())
    }

    // XXX And no `strip_suffix()`.
}

//--- Parseable and Composable

impl Parseable for Dname {
    type Err = DnameError;

    fn parse(parser: &mut Parser) -> Result<Self, DnameError> {
        let len = {
            let mut tmp = parser.peek();
            loop {
                if tmp.is_empty() {
                    return Err(DnameError::ShortData)
                }
                let (label, tail) = Label::split_from(tmp)?;
                tmp = tail;
                if label.is_root() {
                    break;
                }
            }
            parser.remaining() - tmp.len()
        };
        if len > 255 {
            return Err(DnameError::LongName);
        }
        Ok(unsafe {
            Self::from_bytes_unchecked(parser.parse_bytes(len).unwrap())
        })
    }
}

impl Composable for Dname {
    fn compose_len(&self) -> usize {
        self.bytes.len()
    }

    fn compose<B: BufMut>(&self, buf: &mut B) {
        buf.put_slice(self.as_ref())
    }
}

//--- ToLabelIter and ToDname

impl<'a> ToLabelIter<'a> for Dname {
    type LabelIter = DnameIter<'a>;

    fn iter_labels(&'a self) -> Self::LabelIter {
        self.iter()
    }
}

impl ToDname for Dname {
    fn to_name(&self) -> Dname {
        self.clone()
    }
}

//--- Deref and AsRef

impl ops::Deref for Dname {
    type Target = Bytes;

    fn deref(&self) -> &Bytes {
        self.as_ref()
    }
}

impl AsRef<Bytes> for Dname {
    fn as_ref(&self) -> &Bytes {
        &self.bytes
    }
}

impl AsRef<[u8]> for Dname {
    fn as_ref(&self) -> &[u8] {
        self.bytes.as_ref()
    }
}

//--- IntoIterator

impl<'a> IntoIterator for &'a Dname {
    type Item = &'a Label;
    type IntoIter = DnameIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

//--- PartialEq and Eq

impl PartialEq for Dname {
    fn eq(&self, other: &Self) -> bool {
        self.as_slice().eq_ignore_ascii_case(other.as_slice())
    }
}

impl Eq for Dname { }

//--- PartialOrd and Ord

impl PartialOrd for Dname {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.iter().partial_cmp(other.iter())
    }
}

impl Ord for Dname {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.iter().cmp(other.iter())
    }
}

//--- Hash

impl hash::Hash for Dname {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        for item in self.iter() {
            item.hash(state)
        }
    }
}

//--- Display and Debug

impl fmt::Display for Dname {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut iter = self.iter();
        match iter.next() {
            Some(label) => label.fmt(f)?,
            None => return Ok(())
        }
        while let Some(label) = iter.next() {
            f.write_str(".")?;
            label.fmt(f)?;
        }
        Ok(())
    }
}

impl fmt::Debug for Dname {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Dname({})", self)
    }
}

<|endoftext|>


<commit_before>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![desc = "A fast Bloom filter implementation."]
#![license = "BSD"]
#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]

extern crate collections;

use std::cmp;
use std::f64;
use std::hash::Hash;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use std::num::Float;
use std::rand;
use collections::bitv;

#[cfg(test)]
use std::rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::with_capacity(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * Float::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> uint {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * Float::ln(2.0f64)).ceil() as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>Remove unused attributes<commit_after>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]

extern crate collections;

use std::cmp;
use std::f64;
use std::hash::Hash;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use std::num::Float;
use std::rand;
use collections::bitv;

#[cfg(test)]
use std::rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::with_capacity(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * Float::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> uint {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * Float::ln(2.0f64)).ceil() as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>


<commit_before>//! Defines an Oren-Nayar BRDF that describes diffuse reflection from rough surfaces
//! See [Oren-Nayar reflectance model](https://en.wikipedia.org/wiki/Oren%E2%80%93Nayar_reflectance_model)

use std::f32;
use std::num::Float;
use collect::enum_set::EnumSet;

use linalg::Vector;
use film::Colorf;
use bxdf;
use bxdf::{BxDF, BxDFType};

/// Oren-Nayar BRDf that implements the Oren-Nayar reflectance model
#[derive(Copy, Show)]
pub struct OrenNayar {
    /// Color of the diffuse material
    reflectance: Colorf,
    /// Precomputed and stored value of the A constant
    a: f32,
    /// Precomputed and stored value of the B constant
    b: f32,
}

impl OrenNayar {
    /// Create a new Oren-Nayar BRDF with the desired color and roughness
    /// `roughness` should be the variance of the Gaussian describing the
    /// microfacet distribution
    pub fn new(c: &Colorf, roughness: f32) -> OrenNayar {
        OrenNayar { reflectance: *c,
                    a: 1.0 - roughness / (2.0 * (roughness + 0.33)),
                    b: 0.45 * roughness / (roughness + 0.09),
        }
    }
}

impl BxDF for OrenNayar {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Diffuse);
        e.insert(BxDFType::Reflection);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        let sin_theta_o = bxdf::sin_theta(w_o);
        let sin_theta_i = bxdf::sin_theta(w_i);
        let max_cos =
            if sin_theta_i > 1e-4 && sin_theta_o > 1e-4 {
                Float::max(0.0, bxdf::cos_phi(w_i) * bxdf::cos_phi(w_o)
                           + bxdf::sin_phi(w_i) * bxdf::sin_phi(w_o))
            } else {
                0.0
            };
        let (sin_alpha, tan_beta) =
            if Float::abs(bxdf::cos_theta(w_i)) > Float::abs(bxdf::cos_theta(w_o)) {
                (sin_theta_o, sin_theta_i / Float::abs(bxdf::cos_theta(w_i)))
            } else {
                (sin_theta_i, sin_theta_o / Float::abs(bxdf::cos_theta(w_o)))
            };
        self.reflectance * f32::consts::FRAC_1_PI * (self.a + self.b * max_cos * sin_alpha * tan_beta)
    }
}

<commit_msg>Fix Oren-Nayar BRDF. I incorrectly though we could skip finding sigma<commit_after>//! Defines an Oren-Nayar BRDF that describes diffuse reflection from rough surfaces
//! See [Oren-Nayar reflectance model](https://en.wikipedia.org/wiki/Oren%E2%80%93Nayar_reflectance_model)

use std::f32;
use std::num::Float;
use collect::enum_set::EnumSet;

use linalg::Vector;
use film::Colorf;
use bxdf;
use bxdf::{BxDF, BxDFType};

/// Oren-Nayar BRDf that implements the Oren-Nayar reflectance model
#[derive(Copy, Show)]
pub struct OrenNayar {
    /// Color of the diffuse material
    albedo: Colorf,
    /// Precomputed and stored value of the A constant
    a: f32,
    /// Precomputed and stored value of the B constant
    b: f32,
}

impl OrenNayar {
    /// Create a new Oren-Nayar BRDF with the desired color and roughness
    /// `roughness` should be the variance of the Gaussian describing the
    /// microfacet distribution
    pub fn new(c: &Colorf, roughness: f32) -> OrenNayar {
        let mut sigma = Float::to_radians(roughness);
        sigma *= sigma;
        OrenNayar { albedo: *c,
                    a: 1.0 - 0.5 * sigma / (sigma + 0.33),
                    b: 0.45 * sigma / (sigma + 0.09),
        }
    }
}

impl BxDF for OrenNayar {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Diffuse);
        e.insert(BxDFType::Reflection);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        let sin_theta_o = bxdf::sin_theta(w_o);
        let sin_theta_i = bxdf::sin_theta(w_i);
        let max_cos =
            if sin_theta_i > 1e-4 && sin_theta_o > 1e-4 {
                Float::max(0.0, bxdf::cos_phi(w_i) * bxdf::cos_phi(w_o)
                           + bxdf::sin_phi(w_i) * bxdf::sin_phi(w_o))
            } else {
                0.0
            };
        let (sin_alpha, tan_beta) =
            if Float::abs(bxdf::cos_theta(w_i)) > Float::abs(bxdf::cos_theta(w_o)) {
                (sin_theta_o, sin_theta_i / Float::abs(bxdf::cos_theta(w_i)))
            } else {
                (sin_theta_i, sin_theta_o / Float::abs(bxdf::cos_theta(w_o)))
            };
        self.albedo * f32::consts::FRAC_1_PI * (self.a + self.b * max_cos * sin_alpha * tan_beta)
    }
}

<|endoftext|>


<commit_before>use mli::*;
pub use ndarray::{azip, s, Array, Array2, ArrayView2};

fn convolve<'a>(signal: ArrayView2<'a, f32>, filter: ArrayView2<'a, f32>) -> Array2<f32> {
    let filter_dims = filter.raw_dim();
    let output_dims = (
        signal.shape()[0] + 1 - filter_dims[0],
        signal.shape()[1] + 1 - filter_dims[1],
    );
    Array::from_shape_vec(
        output_dims,
        signal
            .windows(filter_dims)
            .into_iter()
            .map(|view| (view.to_owned() * filter).sum())
            .collect(),
    )
    .expect("convolution produced incorrectly sized output")
}

pub struct Conv2(pub Array2<f32>);

impl Forward<Array2<f32>> for Conv2 {
    type Internal = ();
    type Output = Array2<f32>;

    fn forward(&self, input: &Array2<f32>) -> ((), Self::Output) {
        let Self(filter) = self;
        ((), convolve(input.view(), filter.view()))
    }
}

impl Backward<Array2<f32>, Array2<f32>> for Conv2 {
    type InputDelta = Array2<f32>;
    type TrainDelta = Array2<f32>;

    fn backward(
        &self,
        input: &Array2<f32>,
        internal: &Self::Internal,
        output_delta: &Array2<f32>,
    ) -> (Self::InputDelta, Self::TrainDelta) {
        // To compute the `input_delta`, we need to convolve the filter flipped across all its axes with
        // the zero-padded version of the `output_delta`. The amount of zero-padding that needs to be added
        // is `(l - 1) * 2` where `l` is the length of the filter along a given dimension.
        let Self(filter) = self;
        let filter_dims = filter.raw_dim();
        // If the filter is 1 in either dimension, the below code wont work.
        if filter_dims[0] == 1 || filter_dims[1] == 1 {
            unimplemented!("mli-conv: filter dimensions of 1 not implemented yet");
        }
        let padding = (2 * (filter_dims[0] - 1), 2 * (filter_dims[1] - 1));
        let pad_dims = (input.shape()[0] + padding.0, input.shape()[1] + padding.1);
        let mut pad = Array::zeros(pad_dims);
        pad.slice_mut(s![
            padding.0 as i32..-(padding.0 as i32),
            padding.1 as i32..-(padding.1 as i32)
        ])
        .assign(output_delta);
        let input_delta = convolve(pad.view(), filter.slice(s![..;-1,..;-1]));
        (
            input_delta.clone(),
            unimplemented!("train data not yet computed"),
        )
    }
}
<commit_msg>finish Backward impl for Conv2<commit_after>use mli::*;
pub use ndarray::{azip, s, Array, Array2, ArrayView2};

fn convolve<'a>(signal: ArrayView2<'a, f32>, filter: ArrayView2<'a, f32>) -> Array2<f32> {
    let filter_dims = filter.raw_dim();
    let output_dims = (
        signal.shape()[0] + 1 - filter_dims[0],
        signal.shape()[1] + 1 - filter_dims[1],
    );
    Array::from_shape_vec(
        output_dims,
        signal
            .windows(filter_dims)
            .into_iter()
            .map(|view| (view.to_owned() * filter).sum())
            .collect(),
    )
    .expect("convolution produced incorrectly sized output")
}

pub struct Conv2(pub Array2<f32>);

impl Forward<Array2<f32>> for Conv2 {
    type Internal = ();
    type Output = Array2<f32>;

    fn forward(&self, input: &Array2<f32>) -> ((), Self::Output) {
        let Self(filter) = self;
        ((), convolve(input.view(), filter.view()))
    }
}

impl Backward<Array2<f32>, Array2<f32>> for Conv2 {
    type InputDelta = Array2<f32>;
    type TrainDelta = Array2<f32>;

    fn backward(
        &self,
        input: &Array2<f32>,
        _: &Self::Internal,
        output_delta: &Array2<f32>,
    ) -> (Self::InputDelta, Self::TrainDelta) {
        // To compute the `input_delta`, we need to convolve the filter flipped across all its axes with
        // the zero-padded version of the `output_delta`. The amount of zero-padding that needs to be added
        // is `(l - 1) * 2` where `l` is the length of the filter along a given dimension.
        let Self(filter) = self;
        let filter_dims = filter.raw_dim();
        // If the filter is 1 in either dimension, the below code wont work.
        if filter_dims[0] == 1 || filter_dims[1] == 1 {
            unimplemented!("mli-conv: filter dimensions of 1 not implemented yet");
        }
        let padding = (2 * (filter_dims[0] - 1), 2 * (filter_dims[1] - 1));
        let pad_dims = (input.shape()[0] + padding.0, input.shape()[1] + padding.1);
        let mut pad = Array::zeros(pad_dims);
        pad.slice_mut(s![
            padding.0 as i32 / 2..-(padding.0 as i32) / 2,
            padding.1 as i32 / 2..-(padding.1 as i32) / 2
        ])
        .assign(output_delta);
        let input_delta = convolve(pad.view(), filter.slice(s![..;-1,..;-1]));

        let train_delta = input
            .windows(filter_dims)
            .into_iter()
            .zip(output_delta.iter())
            .map(|(view, &delta)| (view.to_owned() * delta))
            .fold(Array2::zeros(filter_dims), |acc, item| acc + item);
        (input_delta, train_delta)
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![cfg_attr(feature = "benchmarks", feature(test))]

#![cfg_attr(feature = "cargo-clippy", allow(
    if_same_then_else,
    needless_bool,
    needless_pass_by_value,
    ptr_arg,
))]

extern crate xi_trace;

#[cfg(any(feature = "chrome_trace_event"))]
extern crate serde;

#[macro_use]
extern crate serde_derive;

#[cfg(feature = "chrome_trace_event")]
#[macro_use]
extern crate serde_json;

#[cfg(all(test, feature = "benchmarks"))]
extern crate test;

#[cfg(feature = "chrome_trace_event")]
pub mod chrome_trace;

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "benchmarks")]
    use test::Bencher;
    #[cfg(feature = "dict_payload")]
    use xi_trace::{StrCow, TracePayloadT};

    #[cfg(all(not(feature = "dict_payload"), not(feature = "json_payload")))]
    fn to_payload(value: &'static str) -> &'static str {
        value
    }

    #[cfg(feature = "dict_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        let mut d = TracePayloadT::with_capacity(1);
        d.insert(StrCow::from("test"), StrCow::from(value));
        d
    }

    #[cfg(feature = "json_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        json!({"test": value})
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_serialization() {
        use xi_trace::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"],|| {
            let _guard = trace.block("sample5", &["test,chrome"]);
        }, to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();

        let result = chrome_trace::serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let decoded_result : Vec<serde_json::Value> = serde_json::from_slice(&serialized).unwrap();
        assert_eq!(decoded_result.len(), 8);
        assert_eq!(decoded_result[0]["name"].as_str().unwrap(), "process_name");
        assert_eq!(decoded_result[1]["name"].as_str().unwrap(), "thread_name");
        for i in 2..5 {
            assert_eq!(decoded_result[i]["name"].as_str().unwrap(), samples[i].name);
            assert_eq!(decoded_result[i]["cat"].as_str().unwrap(), "test,chrome");
            assert_eq!(decoded_result[i]["ph"].as_str().unwrap(), "i");
            assert_eq!(decoded_result[i]["ts"], samples[i].timestamp_us);
            let nth_sample = &samples[i];
            let mut nth_args = nth_sample.args.as_ref().unwrap();
            assert_eq!(decoded_result[i]["args"]["xi_payload"], json!(nth_args.payload.as_ref()));
        }
        assert_eq!(decoded_result[5]["ph"], "B");
        assert_eq!(decoded_result[6]["ph"], "E");
        assert_eq!(decoded_result[7]["ph"], "X");
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_deserialization() {
        use xi_trace::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"],|| (), to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();
        let result = chrome_trace::serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let deserialized_samples = chrome_trace::deserialize(serialized.as_slice()).unwrap();
        assert_eq!(deserialized_samples, samples);
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_one_element(b: &mut Bencher) {
        use chrome_trace::*;

        let mut serialized = Vec::<u8>::new();
        let samples = [xi_trace::Sample::new_instant("trace1", &["benchmark", "test"], None)];
        b.iter(|| {
            serialized.clear();
            serialize(samples.iter(), &mut serialized).unwrap();
        });
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_multiple_elements(b: &mut Bencher) {
        use chrome_trace::*;
        use xi_trace::*;

        let mut serialized = Vec::<u8>::new();
        let mut samples = [
            Sample::new_instant("trace1", &["benchmark", "test"], None),
            Sample::new_instant("trace2", &["benchmark"], None),
            Sample::new_duration("trace3", &["benchmark"], Some(to_payload("some payload"), 0)),
            Sample::new_instant("trace4", &["benchmark"], None)];

        b.iter(|| {
            serialized.clear();
            serialize(samples.iter(), &mut serialized).unwrap();
        });
    }
}
<commit_msg>Silence trace_dump compiler warning<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![cfg_attr(feature = "benchmarks", feature(test))]

#![cfg_attr(feature = "cargo-clippy", allow(
    if_same_then_else,
    needless_bool,
    needless_pass_by_value,
    ptr_arg,
))]

extern crate xi_trace;

#[cfg(any(feature = "chrome_trace_event"))]
extern crate serde;

#[macro_use]
extern crate serde_derive;

#[cfg(all(not(test), feature = "chrome_trace_event"))]
extern crate serde_json;

#[cfg(test)]
#[macro_use] extern crate serde_json;

#[cfg(all(test, feature = "benchmarks"))]
extern crate test;

#[cfg(feature = "chrome_trace_event")]
pub mod chrome_trace;

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "benchmarks")]
    use test::Bencher;
    #[cfg(feature = "dict_payload")]
    use xi_trace::{StrCow, TracePayloadT};

    #[cfg(all(not(feature = "dict_payload"), not(feature = "json_payload")))]
    fn to_payload(value: &'static str) -> &'static str {
        value
    }

    #[cfg(feature = "dict_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        let mut d = TracePayloadT::with_capacity(1);
        d.insert(StrCow::from("test"), StrCow::from(value));
        d
    }

    #[cfg(feature = "json_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        json!({"test": value})
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_serialization() {
        use xi_trace::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"],|| {
            let _guard = trace.block("sample5", &["test,chrome"]);
        }, to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();

        let result = chrome_trace::serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let decoded_result : Vec<serde_json::Value> = serde_json::from_slice(&serialized).unwrap();
        assert_eq!(decoded_result.len(), 8);
        assert_eq!(decoded_result[0]["name"].as_str().unwrap(), "process_name");
        assert_eq!(decoded_result[1]["name"].as_str().unwrap(), "thread_name");
        for i in 2..5 {
            assert_eq!(decoded_result[i]["name"].as_str().unwrap(), samples[i].name);
            assert_eq!(decoded_result[i]["cat"].as_str().unwrap(), "test,chrome");
            assert_eq!(decoded_result[i]["ph"].as_str().unwrap(), "i");
            assert_eq!(decoded_result[i]["ts"], samples[i].timestamp_us);
            let nth_sample = &samples[i];
            let mut nth_args = nth_sample.args.as_ref().unwrap();
            assert_eq!(decoded_result[i]["args"]["xi_payload"], json!(nth_args.payload.as_ref()));
        }
        assert_eq!(decoded_result[5]["ph"], "B");
        assert_eq!(decoded_result[6]["ph"], "E");
        assert_eq!(decoded_result[7]["ph"], "X");
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_deserialization() {
        use xi_trace::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"],|| (), to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();
        let result = chrome_trace::serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let deserialized_samples = chrome_trace::deserialize(serialized.as_slice()).unwrap();
        assert_eq!(deserialized_samples, samples);
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_one_element(b: &mut Bencher) {
        use chrome_trace::*;

        let mut serialized = Vec::<u8>::new();
        let samples = [xi_trace::Sample::new_instant("trace1", &["benchmark", "test"], None)];
        b.iter(|| {
            serialized.clear();
            serialize(samples.iter(), &mut serialized).unwrap();
        });
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_multiple_elements(b: &mut Bencher) {
        use chrome_trace::*;
        use xi_trace::*;

        let mut serialized = Vec::<u8>::new();
        let mut samples = [
            Sample::new_instant("trace1", &["benchmark", "test"], None),
            Sample::new_instant("trace2", &["benchmark"], None),
            Sample::new_duration("trace3", &["benchmark"], Some(to_payload("some payload"), 0)),
            Sample::new_instant("trace4", &["benchmark"], None)];

        b.iter(|| {
            serialized.clear();
            serialize(samples.iter(), &mut serialized).unwrap();
        });
    }
}
<|endoftext|>


<commit_before>//! See [`patch_json_for_outdated_configs`]
use serde_json::{json, Value};

/// This function patches the json config to the new expected keys.
/// That is we try to load old known config keys here and convert them to the new ones.
/// See https://github.com/rust-lang/rust-analyzer/pull/12010
pub(super) fn patch_json_for_outdated_configs(json: &mut Value) {
    let copy = json.clone();

    macro_rules! patch {
        ($(
            $($src:ident).+ -> $($dst:ident).+ ;
        )+) => { $(
            if let Some(it) = copy.pointer(concat!($("/", stringify!($src)),+)).cloned() {
                let mut last = it;
                for segment in [$(stringify!($dst)),+].into_iter().rev() {
                    last = Value::Object(serde_json::Map::from_iter(std::iter::once((segment.to_string(), last))));
                }

                merge(json, last);
            }
        )+ };
    }

    patch! {
        assist.allowMergingIntoGlobImports -> imports.merge.glob;
        assist.exprFillDefault -> assist.expressionFillDefault;
        assist.importEnforceGranularity -> imports.granularity.enforce;
        assist.importGranularity -> imports.granularity.group;
        assist.importMergeBehavior -> imports.granularity.group;
        assist.importMergeBehaviour -> imports.granularity.group;
        assist.importGroup -> imports.group.enable;
        assist.importPrefix -> imports.prefix;
        cache.warmup -> primeCaches.enable;
        cargo.loadOutDirsFromCheck -> cargo.buildScripts.enable;
        cargo.runBuildScripts -> cargo.runBuildScripts.overrideCommand;
        cargo.runBuildScriptsCommand -> cargo.runBuildScripts.overrideCommand;
        cargo.useRustcWrapperForBuildScripts -> cargo.runBuildScripts.useRustcWrapper;
        completion.snippets -> completion.snippets.custom;
        diagnostics.enableExperimental -> diagnostics.experimental.enable;
        experimental.procAttrMacros -> procMacro.attributes.enable;
        highlighting.strings -> semanticHighlighting.strings.enable;
        highlightRelated.breakPoints -> semanticHighlighting.breakPoints.enable;
        highlightRelated.exitPoints -> semanticHighlighting.exitPoints.enable;
        highlightRelated.yieldPoints -> semanticHighlighting.yieldPoints.enable;
        highlightRelated.references -> semanticHighlighting.references.enable;
        hover.documentation -> hover.documentation.enable;
        hover.linksInHover -> hover.links.enable;
        hoverActions.linksInHover -> hover.links.enable;
        hoverActions.debug -> hoverActions.debug.enable;
        hoverActions.enable -> hoverActions.enable.enable;
        hoverActions.gotoTypeDef -> hoverActions.gotoTypeDef.enable;
        hoverActions.implementations -> hoverActions.implementations.enable;
        hoverActions.references -> hoverActions.references.enable;
        hoverActions.run -> hoverActions.run.enable;
        inlayHints.chainingHints -> inlayHints.chainingHints.enable;
        inlayHints.closureReturnTypeHints -> inlayHints.closureReturnTypeHints.enable;
        inlayHints.hideNamedConstructorHints -> inlayHints.typeHints.hideNamedConstructorHints;
        inlayHints.parameterHints -> inlayHints.parameterHints.enable;
        inlayHints.reborrowHints -> inlayHints.reborrowHints.enable;
        inlayHints.typeHints -> inlayHints.typeHints.enable;
        lruCapacity -> lru.capacity;
        runnables.cargoExtraArgs -> runnables.extraArgs ;
        runnables.overrideCargo -> runnables.command ;
        rustcSource -> rustc.source;
        rustfmt.enableRangeFormatting -> rustfmt.rangeFormatting.enable;
    }

    // callInfo_full -> signatureInfo_detail, signatureInfo_documentation_enable
    if let Some(Value::Bool(b)) = copy.pointer("/callInfo/full") {
        let sig_info = match b {
            true => json!({ "signatureInfo": {
                "documentation": {"enable": true}},
                "detail": "full"
            }),
            false => json!({ "signatureInfo": {
                "documentation": {"enable": false}},
                "detail": "parameters"
            }),
        };
        merge(json, sig_info);
    }

    // cargo_allFeatures, cargo_features -> cargo_features
    if let Some(Value::Bool(true)) = copy.pointer("/cargo/allFeatures") {
        merge(json, json!({ "cargo": { "features": "all" } }));
    }

    // checkOnSave_allFeatures, checkOnSave_features -> checkOnSave_features
    if let Some(Value::Bool(true)) = copy.pointer("/checkOnSave/allFeatures") {
        merge(json, json!({ "checkOnSave": { "features": "all" } }));
    }

    // completion_addCallArgumentSnippets completion_addCallParenthesis -> completion_callable_snippets
    let res = match (
        copy.pointer("/completion/addCallArgumentSnippets"),
        copy.pointer("/completion/addCallParenthesis"),
    ) {
        (Some(Value::Bool(true)), Some(Value::Bool(true))) => json!("fill_arguments"),
        (Some(Value::Bool(true)), _) => json!("add_parentheses"),
        (_, _) => json!(null),
    };
    merge(json, json!({ "completion": { "callable": {"snippets": res }} }));
}

fn merge(dst: &mut Value, src: Value) {
    match (dst, src) {
        (Value::Object(dst), Value::Object(src)) => {
            for (k, v) in src {
                merge(dst.entry(k).or_insert(v.clone()), v)
            }
        }
        (dst, src) => *dst = src,
    }
}
<commit_msg>fix: Fix config patching failing when appending suffixes<commit_after>//! See [`patch_json_for_outdated_configs`]
use serde_json::{json, Value};

/// This function patches the json config to the new expected keys.
/// That is we try to load old known config keys here and convert them to the new ones.
/// See https://github.com/rust-lang/rust-analyzer/pull/12010
pub(super) fn patch_json_for_outdated_configs(json: &mut Value) {
    let copy = json.clone();

    macro_rules! patch {
        ($(
            $($src:ident).+ -> $($dst:ident).+ ;
        )+) => { $(
            match copy.pointer(concat!($("/", stringify!($src)),+)).cloned() {
                Some(Value::Object(_)) | None => (),
                Some(it) => {
                    let mut last = it;
                    for segment in [$(stringify!($dst)),+].into_iter().rev() {
                        last = Value::Object(serde_json::Map::from_iter(std::iter::once((segment.to_string(), last))));
                    }

                    merge(json, last);
                },
            }
        )+ };
    }

    patch! {
        assist.allowMergingIntoGlobImports -> imports.merge.glob;
        assist.exprFillDefault -> assist.expressionFillDefault;
        assist.importEnforceGranularity -> imports.granularity.enforce;
        assist.importGranularity -> imports.granularity.group;
        assist.importMergeBehavior -> imports.granularity.group;
        assist.importMergeBehaviour -> imports.granularity.group;
        assist.importGroup -> imports.group.enable;
        assist.importPrefix -> imports.prefix;
        cache.warmup -> primeCaches.enable;
        cargo.loadOutDirsFromCheck -> cargo.buildScripts.enable;
        cargo.runBuildScripts -> cargo.runBuildScripts.overrideCommand;
        cargo.runBuildScriptsCommand -> cargo.runBuildScripts.overrideCommand;
        cargo.useRustcWrapperForBuildScripts -> cargo.runBuildScripts.useRustcWrapper;
        diagnostics.enableExperimental -> diagnostics.experimental.enable;
        experimental.procAttrMacros -> procMacro.attributes.enable;
        highlighting.strings -> semanticHighlighting.strings.enable;
        highlightRelated.breakPoints -> semanticHighlighting.breakPoints.enable;
        highlightRelated.exitPoints -> semanticHighlighting.exitPoints.enable;
        highlightRelated.yieldPoints -> semanticHighlighting.yieldPoints.enable;
        highlightRelated.references -> semanticHighlighting.references.enable;
        hover.documentation -> hover.documentation.enable;
        hover.linksInHover -> hover.links.enable;
        hoverActions.linksInHover -> hover.links.enable;
        hoverActions.debug -> hoverActions.debug.enable;
        hoverActions.enable -> hoverActions.enable.enable;
        hoverActions.gotoTypeDef -> hoverActions.gotoTypeDef.enable;
        hoverActions.implementations -> hoverActions.implementations.enable;
        hoverActions.references -> hoverActions.references.enable;
        hoverActions.run -> hoverActions.run.enable;
        inlayHints.chainingHints -> inlayHints.chainingHints.enable;
        inlayHints.closureReturnTypeHints -> inlayHints.closureReturnTypeHints.enable;
        inlayHints.hideNamedConstructorHints -> inlayHints.typeHints.hideNamedConstructorHints;
        inlayHints.parameterHints -> inlayHints.parameterHints.enable;
        inlayHints.reborrowHints -> inlayHints.reborrowHints.enable;
        inlayHints.typeHints -> inlayHints.typeHints.enable;
        lruCapacity -> lru.capacity;
        runnables.cargoExtraArgs -> runnables.extraArgs ;
        runnables.overrideCargo -> runnables.command ;
        rustcSource -> rustc.source;
        rustfmt.enableRangeFormatting -> rustfmt.rangeFormatting.enable;
    }

    // completion.snippets -> completion.snippets.custom;
    if let Some(Value::Object(obj)) = copy.pointer("/completion/snippets").cloned() {
        if obj.len() != 1 || obj.get("custom").is_none() {
            merge(
                json,
                json! {{
                    "completion": {
                        "snippets": {
                            "custom": obj
                        },
                    },
                }},
            );
        }
    }

    // callInfo_full -> signatureInfo_detail, signatureInfo_documentation_enable
    if let Some(Value::Bool(b)) = copy.pointer("/callInfo/full") {
        let sig_info = match b {
            true => json!({ "signatureInfo": {
                "documentation": {"enable": true}},
                "detail": "full"
            }),
            false => json!({ "signatureInfo": {
                "documentation": {"enable": false}},
                "detail": "parameters"
            }),
        };
        merge(json, sig_info);
    }

    // cargo_allFeatures, cargo_features -> cargo_features
    if let Some(Value::Bool(true)) = copy.pointer("/cargo/allFeatures") {
        merge(json, json!({ "cargo": { "features": "all" } }));
    }

    // checkOnSave_allFeatures, checkOnSave_features -> checkOnSave_features
    if let Some(Value::Bool(true)) = copy.pointer("/checkOnSave/allFeatures") {
        merge(json, json!({ "checkOnSave": { "features": "all" } }));
    }

    // completion_addCallArgumentSnippets completion_addCallParenthesis -> completion_callable_snippets
    let res = match (
        copy.pointer("/completion/addCallArgumentSnippets"),
        copy.pointer("/completion/addCallParenthesis"),
    ) {
        (Some(Value::Bool(true)), Some(Value::Bool(true))) => json!("fill_arguments"),
        (Some(Value::Bool(true)), _) => json!("add_parentheses"),
        (_, _) => json!(null),
    };
    merge(json, json!({ "completion": { "callable": {"snippets": res }} }));
}

fn merge(dst: &mut Value, src: Value) {
    match (dst, src) {
        (Value::Object(dst), Value::Object(src)) => {
            for (k, v) in src {
                merge(dst.entry(k).or_insert(v.clone()), v)
            }
        }
        (dst, src) => *dst = src,
    }
}
<|endoftext|>


<commit_before>use std::cmp::Reverse;
use std::collections::hash_map::Entry;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::fmt::Display;
use std::io::Read;
use std::mem::swap;

use crate::common::LineIter;

type Item<const S: usize> = (u32, State<S>);
type Todo<const S: usize> = BinaryHeap<Reverse<Item<S>>>;
type Visited<const S: usize> = HashMap<State<S>, u32>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Hash)]
enum Pod {
    A,
    B,
    C,
    D,
}

impl Pod {
    pub fn cost(self) -> u32 {
        match self {
            Pod::A => 1,
            Pod::B => 10,
            Pod::C => 100,
            Pod::D => 1000,
        }
    }
}

impl TryFrom<char> for Pod {
    type Error = &'static str;

    fn try_from(c: char) -> Result<Self, Self::Error> {
        match c {
            'A' => Ok(Pod::A),
            'B' => Ok(Pod::B),
            'C' => Ok(Pod::C),
            'D' => Ok(Pod::D),
            _ => Err("Invalid pod"),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
struct State<const S: usize> {
    hallway: [Option<Pod>; 11],
    rooms: [[Option<Pod>; S]; 4],
}

fn room_hallway_pos(room: usize) -> usize {
    room * 2 + 2
}

fn abs_delta(a: usize, b: usize) -> usize {
    if a < b {
        b - a
    } else {
        a - b
    }
}

impl<const S: usize> State<S> {
    const VALID_HALLWAY_POS: [usize; 7] = [0, 1, 3, 5, 7, 9, 10];

    pub fn is_done(&self) -> bool {
        self == &State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); S],
                [Some(Pod::B); S],
                [Some(Pod::C); S],
                [Some(Pod::D); S],
            ],
        }
    }

    fn add_to_queue(self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        let entry = visited.entry(self.clone());

        if matches!(&entry, Entry::Occupied(entry) if *entry.get() <= cost) {
            // Already got a better one
            return;
        }

        // nightly only :'(
        // entry.insert(cost);
        *entry.or_default() = cost;

        todo.push(Reverse((cost + self.estimate(), self)))
    }

    fn estimate(&self) -> u32 {
        // A* estimate. For every entry that is not already "at rest", the cost is the cost
        // required to get it to the top of its intended room.

        // Cost to enter the hole for all pods that still need to
        let enter_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(index, room)| {
                let pod = match index {
                    0 => Pod::A,
                    1 => Pod::B,
                    2 => Pod::C,
                    3 => Pod::D,
                    _ => unreachable!(),
                };

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| entry == Some(pod))
                    .map(|(index, _)| index as u32 + 1)
                    .sum::<u32>()
                    * pod.cost()
            })
            .sum();

        // Cost for all of the hallway to move to above their intended rooms
        let hallway_estimate: u32 = self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, &pod)| {
                let pod = pod?;

                let destination_pos = room_hallway_pos(pod as usize);

                Some(abs_delta(pos, destination_pos) as u32 * pod.cost())
            })
            .sum();

        // Cost to move out of the room and above the correct rooms
        let rooms_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(room_index, room)| {
                let hallway_pos = room_hallway_pos(room_index);

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| {
                        entry.map(|pod| pod as usize == room_index).unwrap_or(false)
                    })
                    .filter_map(|(room_pos, &pod)| {
                        let pod = pod?;

                        let destination_pos = room_hallway_pos(pod as usize);

                        let steps = 1 + room_pos + abs_delta(hallway_pos, destination_pos).max(2);

                        Some(steps as u32 * pod.cost())
                    })
                    .sum::<u32>()
            })
            .sum();

        enter_estimate + hallway_estimate + rooms_estimate
    }

    pub fn generate_next(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        self.room_to_hallway(cost, todo, visited);
        self.hallway_to_room(cost, todo, visited);
    }

    fn room_to_hallway(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (index, room) in self.rooms.iter().enumerate() {
            // Check if we even want to move anything out of this room
            if room
                .iter()
                .all(|entry| entry.map(|pod| pod as usize == index).unwrap_or(true))
            {
                continue;
            }

            let (pos, pod) = room
                .iter()
                .enumerate()
                .find_map(|(pos, entry)| entry.map(|pod| (pos, pod)))
                .unwrap(); // Safe unwrap, we know it exists from above.

            let base_cost = 1 + pos;
            let hallway_pos = room_hallway_pos(index);

            let mut queue_new = |new_pos, new_cost| {
                let mut new_state = self.clone();
                swap(
                    &mut new_state.hallway[new_pos],
                    &mut new_state.rooms[index][pos],
                );

                new_state.add_to_queue(new_cost + cost, todo, visited)
            };

            // Check positions to the left
            for new_pos in (0..hallway_pos).rev() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + hallway_pos - new_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }

            // And to the right
            for new_pos in hallway_pos..self.hallway.len() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + new_pos - hallway_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }
        }
    }

    fn hallway_to_room(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (pos, pod) in self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, pod)| pod.map(|pod| (pos, pod)))
        {
            let room = pod as usize;
            let new_hallway_pos = room_hallway_pos(room);

            // Check if the path is free
            let in_between = if new_hallway_pos < pos {
                &self.hallway[(new_hallway_pos + 1)..pos]
            } else {
                &self.hallway[(pos + 1)..new_hallway_pos]
            };

            if in_between.iter().any(Option::is_some) {
                // Something's in the way
                continue;
            }

            // Check if we can move into the room
            if self.rooms[room]
                .iter()
                .copied()
                .flatten()
                .any(|other| other != pod)
            {
                // Scared of other pods
                continue;
            }

            let room_pos = if let Some(pos) = self.rooms[room].iter().rposition(Option::is_none) {
                pos
            } else {
                continue;
            };

            let new_cost = (abs_delta(pos, new_hallway_pos) + room_pos + 1) as u32 * pod.cost();
            let mut new_state = self.clone();
            swap(
                &mut new_state.hallway[pos],
                &mut new_state.rooms[room][room_pos],
            );
            new_state.add_to_queue(cost + new_cost, todo, visited);
        }
    }

    pub fn solve(&self) -> u32 {
        let mut todo = Todo::new();

        let mut visited = HashMap::new();
        visited.insert(self.clone(), 0);

        todo.push(Reverse((self.estimate(), self.clone())));

        while let Some(Reverse((_, state))) = todo.pop() {
            let cost = *visited.get(&state).unwrap_or(&0);

            if state.is_done() {
                return cost;
            }

            state.generate_next(cost, &mut todo, &mut visited);
        }

        panic!("No route found!")
    }
}

impl<const S: usize> Display for State<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let helper = |opt_pod| match opt_pod {
            Some(Pod::A) => 'A',
            Some(Pod::B) => 'B',
            Some(Pod::C) => 'C',
            Some(Pod::D) => 'D',
            None => '.',
        };
        writeln!(f, "#############")?;
        write!(f, "#")?;

        for entry in self.hallway {
            write!(f, "{}", helper(entry))?;
        }
        writeln!(f, "#")?;

        for i in 0..S {
            writeln!(
                f,
                "  #{}#{}#{}#{}#",
                helper(self.rooms[0][i]),
                helper(self.rooms[1][i]),
                helper(self.rooms[2][i]),
                helper(self.rooms[3][i])
            )?;
        }

        write!(f, "  #########")
    }
}

fn read_input(input: &mut dyn Read) -> State<2> {
    let mut reader = LineIter::new(input);
    let mut state = State {
        hallway: Default::default(),
        rooms: Default::default(),
    };

    let _ = reader.next();
    let _ = reader.next();

    let mut helper = |idx: usize| {
        reader
            .next()
            .unwrap()
            .chars()
            .filter_map(|c| Pod::try_from(c).ok())
            .zip(&mut state.rooms)
            .for_each(|(pod, room)| room[idx] = Some(pod))
    };

    helper(0);
    helper(1);

    state
}

pub fn part1(input: &mut dyn Read) -> String {
    let state = read_input(input);

    state.solve().to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let state2 = read_input(input);

    let state4 = State {
        hallway: Default::default(),
        rooms: [
            [
                state2.rooms[0][0],
                Some(Pod::D),
                Some(Pod::D),
                state2.rooms[0][1],
            ],
            [
                state2.rooms[1][0],
                Some(Pod::C),
                Some(Pod::B),
                state2.rooms[1][1],
            ],
            [
                state2.rooms[2][0],
                Some(Pod::B),
                Some(Pod::A),
                state2.rooms[2][1],
            ],
            [
                state2.rooms[3][0],
                Some(Pod::A),
                Some(Pod::C),
                state2.rooms[3][1],
            ],
        ],
    };

    state4.solve().to_string()
}

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

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/23.txt");

    #[test]
    fn test_is_done() {
        let state = State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); 2],
                [Some(Pod::B); 2],
                [Some(Pod::C); 2],
                [Some(Pod::D); 2],
            ],
        };

        assert!(state.is_done());
    }

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 12521);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 44169);
    }
}
<commit_msg>Readability<commit_after>use std::cmp::Reverse;
use std::collections::hash_map::Entry;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::fmt::Display;
use std::io::Read;
use std::mem::swap;

use crate::common::LineIter;

type Item<const S: usize> = (u32, State<S>);
type Todo<const S: usize> = BinaryHeap<Reverse<Item<S>>>;
type Visited<const S: usize> = HashMap<State<S>, u32>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Hash)]
enum Pod {
    A,
    B,
    C,
    D,
}

impl Pod {
    pub fn cost(self) -> u32 {
        match self {
            Pod::A => 1,
            Pod::B => 10,
            Pod::C => 100,
            Pod::D => 1000,
        }
    }

    pub fn dest(self) -> usize {
        self as usize
    }
}

impl TryFrom<char> for Pod {
    type Error = &'static str;

    fn try_from(c: char) -> Result<Self, Self::Error> {
        match c {
            'A' => Ok(Pod::A),
            'B' => Ok(Pod::B),
            'C' => Ok(Pod::C),
            'D' => Ok(Pod::D),
            _ => Err("Invalid pod"),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
struct State<const S: usize> {
    hallway: [Option<Pod>; 11],
    rooms: [[Option<Pod>; S]; 4],
}

fn room_hallway_pos(room: usize) -> usize {
    room * 2 + 2
}

fn abs_delta(a: usize, b: usize) -> usize {
    if a < b {
        b - a
    } else {
        a - b
    }
}

impl<const S: usize> State<S> {
    const VALID_HALLWAY_POS: [usize; 7] = [0, 1, 3, 5, 7, 9, 10];

    pub fn is_done(&self) -> bool {
        self == &State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); S],
                [Some(Pod::B); S],
                [Some(Pod::C); S],
                [Some(Pod::D); S],
            ],
        }
    }

    fn add_to_queue(self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        let entry = visited.entry(self.clone());

        if matches!(&entry, Entry::Occupied(entry) if *entry.get() <= cost) {
            // Already got a better one
            return;
        }

        // nightly only :'(
        // entry.insert(cost);
        *entry.or_default() = cost;

        todo.push(Reverse((cost + self.estimate(), self)))
    }

    fn estimate(&self) -> u32 {
        // A* estimate. For every entry that is not already "at rest", the cost is the cost
        // required to get it to the top of its intended room.

        // Cost to enter the hole for all pods that still need to
        let enter_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(index, room)| {
                let pod = match index {
                    0 => Pod::A,
                    1 => Pod::B,
                    2 => Pod::C,
                    3 => Pod::D,
                    _ => unreachable!(),
                };

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| entry == Some(pod))
                    .map(|(index, _)| index as u32 + 1)
                    .sum::<u32>()
                    * pod.cost()
            })
            .sum();

        // Cost for all of the hallway to move to above their intended rooms
        let hallway_estimate: u32 = self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, &pod)| {
                let pod = pod?;

                let destination_pos = room_hallway_pos(pod.dest());

                Some(abs_delta(pos, destination_pos) as u32 * pod.cost())
            })
            .sum();

        // Cost to move out of the room and above the correct rooms
        let rooms_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(room_index, room)| {
                let hallway_pos = room_hallway_pos(room_index);

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| {
                        entry.map(|pod| pod.dest() == room_index).unwrap_or(false)
                    })
                    .filter_map(|(room_pos, &pod)| {
                        let pod = pod?;

                        let destination_pos = room_hallway_pos(pod.dest());

                        let steps = 1 + room_pos + abs_delta(hallway_pos, destination_pos).max(2);

                        Some(steps as u32 * pod.cost())
                    })
                    .sum::<u32>()
            })
            .sum();

        enter_estimate + hallway_estimate + rooms_estimate
    }

    pub fn generate_next(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        self.room_to_hallway(cost, todo, visited);
        self.hallway_to_room(cost, todo, visited);
    }

    fn room_to_hallway(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (index, room) in self.rooms.iter().enumerate() {
            // Check if we even want to move anything out of this room
            if room
                .iter()
                .all(|entry| entry.map(|pod| pod.dest() == index).unwrap_or(true))
            {
                continue;
            }

            let (pos, pod) = room
                .iter()
                .enumerate()
                .find_map(|(pos, entry)| entry.map(|pod| (pos, pod)))
                .unwrap(); // Safe unwrap, we know it exists from above.

            let base_cost = 1 + pos;
            let hallway_pos = room_hallway_pos(index);

            let mut queue_new = |new_pos, new_cost| {
                let mut new_state = self.clone();
                swap(
                    &mut new_state.hallway[new_pos],
                    &mut new_state.rooms[index][pos],
                );

                new_state.add_to_queue(new_cost + cost, todo, visited)
            };

            // Check positions to the left
            for new_pos in (0..hallway_pos).rev() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + hallway_pos - new_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }

            // And to the right
            for new_pos in hallway_pos..self.hallway.len() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + new_pos - hallway_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }
        }
    }

    fn hallway_to_room(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (pos, pod) in self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, pod)| pod.map(|pod| (pos, pod)))
        {
            let room = pod.dest();
            let new_hallway_pos = room_hallway_pos(room);

            // Check if the path is free
            let in_between = if new_hallway_pos < pos {
                &self.hallway[(new_hallway_pos + 1)..pos]
            } else {
                &self.hallway[(pos + 1)..new_hallway_pos]
            };

            if in_between.iter().any(Option::is_some) {
                // Something's in the way
                continue;
            }

            // Check if we can move into the room
            if self.rooms[room]
                .iter()
                .copied()
                .flatten()
                .any(|other| other != pod)
            {
                // Scared of other pods
                continue;
            }

            let room_pos = if let Some(pos) = self.rooms[room].iter().rposition(Option::is_none) {
                pos
            } else {
                continue;
            };

            let new_cost = (abs_delta(pos, new_hallway_pos) + room_pos + 1) as u32 * pod.cost();
            let mut new_state = self.clone();
            swap(
                &mut new_state.hallway[pos],
                &mut new_state.rooms[room][room_pos],
            );
            new_state.add_to_queue(cost + new_cost, todo, visited);
        }
    }

    pub fn solve(&self) -> u32 {
        let mut todo = Todo::new();

        let mut visited = HashMap::new();
        visited.insert(self.clone(), 0);

        todo.push(Reverse((self.estimate(), self.clone())));

        while let Some(Reverse((_, state))) = todo.pop() {
            let cost = *visited.get(&state).unwrap_or(&0);

            if state.is_done() {
                return cost;
            }

            state.generate_next(cost, &mut todo, &mut visited);
        }

        panic!("No route found!")
    }
}

impl<const S: usize> Display for State<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let helper = |opt_pod| match opt_pod {
            Some(Pod::A) => 'A',
            Some(Pod::B) => 'B',
            Some(Pod::C) => 'C',
            Some(Pod::D) => 'D',
            None => '.',
        };
        writeln!(f, "#############")?;
        write!(f, "#")?;

        for entry in self.hallway {
            write!(f, "{}", helper(entry))?;
        }
        writeln!(f, "#")?;

        for i in 0..S {
            writeln!(
                f,
                "  #{}#{}#{}#{}#",
                helper(self.rooms[0][i]),
                helper(self.rooms[1][i]),
                helper(self.rooms[2][i]),
                helper(self.rooms[3][i])
            )?;
        }

        write!(f, "  #########")
    }
}

fn read_input(input: &mut dyn Read) -> State<2> {
    let mut reader = LineIter::new(input);
    let mut state = State {
        hallway: Default::default(),
        rooms: Default::default(),
    };

    let _ = reader.next();
    let _ = reader.next();

    let mut helper = |idx: usize| {
        reader
            .next()
            .unwrap()
            .chars()
            .filter_map(|c| Pod::try_from(c).ok())
            .zip(&mut state.rooms)
            .for_each(|(pod, room)| room[idx] = Some(pod))
    };

    helper(0);
    helper(1);

    state
}

pub fn part1(input: &mut dyn Read) -> String {
    let state = read_input(input);

    state.solve().to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let state2 = read_input(input);

    let state4 = State {
        hallway: Default::default(),
        rooms: [
            [
                state2.rooms[0][0],
                Some(Pod::D),
                Some(Pod::D),
                state2.rooms[0][1],
            ],
            [
                state2.rooms[1][0],
                Some(Pod::C),
                Some(Pod::B),
                state2.rooms[1][1],
            ],
            [
                state2.rooms[2][0],
                Some(Pod::B),
                Some(Pod::A),
                state2.rooms[2][1],
            ],
            [
                state2.rooms[3][0],
                Some(Pod::A),
                Some(Pod::C),
                state2.rooms[3][1],
            ],
        ],
    };

    state4.solve().to_string()
}

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

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/23.txt");

    #[test]
    fn test_is_done() {
        let state = State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); 2],
                [Some(Pod::B); 2],
                [Some(Pod::C); 2],
                [Some(Pod::D); 2],
            ],
        };

        assert!(state.is_done());
    }

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 12521);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 44169);
    }
}
<|endoftext|>