<commit_before>// Copyright 2019-2020 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.

//! Helps manipulate bit fields in 32-bits words.
// TODO(ia0): Remove when the module is used.
#![cfg_attr(not(test), allow(dead_code, unused_macros))]

use crate::{StoreError, StoreResult};

/// Represents a bit field.
///
/// A bit field is a contiguous sequence of bits in a 32-bits word.
///
/// # Invariant
///
/// - The bit field must fit in a 32-bits word: `pos + len < 32`.
pub struct Field {
    /// The position of the bit field.
    pub pos: usize,

    /// The length of the bit field.
    pub len: usize,
}

impl Field {
    /// Reads the value of a bit field.
    pub fn get(&self, word: u32) -> usize {
        ((word >> self.pos) & self.mask()) as usize
    }

    /// Sets the value of a bit field.
    ///
    /// # Preconditions
    ///
    /// - The value must fit in the bit field: `num_bits(value) < self.len`.
    /// - The value must only change bits from 1 to 0: `self.get(*word) & value == value`.
    pub fn set(&self, word: &mut u32, value: usize) {
        let value = value as u32;
        debug_assert_eq!(value & self.mask(), value);
        let mask = !(self.mask() << self.pos);
        *word &= mask | (value << self.pos);
        debug_assert_eq!(self.get(*word), value as usize);
    }

    /// Returns a bit mask the length of the bit field.
    ///
    /// The mask is meant to be applied on a value. It should be shifted to be applied to the bit
    /// field.
    fn mask(&self) -> u32 {
        (1 << self.len) - 1
    }
}

/// Represents a constant bit field.
///
/// # Invariant
///
/// - The value must fit in the bit field: `num_bits(value) <= field.len`.
pub struct Const {
    /// The bit field.
    pub field: Field,

    /// The constant value.
    pub value: usize,
}

impl Const {
    /// Checks that the bit field has its value.
    pub fn check(&self, word: u32) -> bool {
        self.field.get(word) == self.value
    }

    /// Sets the bit field to its value.
    pub fn set(&self, word: &mut u32) {
        self.field.set(word, self.value);
    }
}

/// Represents a single bit.
///
/// # Invariant
///
/// - The bit must fit in a 32-bits word: `pos < 32`.
pub struct Bit {
    /// The position of the bit.
    pub pos: usize,
}

impl Bit {
    /// Returns whether the value of the bit is zero.
    pub fn get(&self, word: u32) -> bool {
        word & (1 << self.pos) == 0
    }

    /// Sets the value of the bit to zero.
    pub fn set(&self, word: &mut u32) {
        *word &= !(1 << self.pos);
    }
}

/// Represents a checksum.
///
/// A checksum is a bit field counting how many bits are set to zero in the word (except in the
/// checksum itself) plus some external increment. It essentially behaves like a bit field storing
/// the external increment.
pub struct Checksum {
    /// The bit field
    pub field: Field,
}

impl Checksum {
    /// Reads the external increment from the checksum.
    ///
    /// # Errors
    ///
    /// Returns `InvalidStorage` if the external increment would be negative.
    pub fn get(&self, word: u32) -> StoreResult<usize> {
        let checksum = self.field.get(word);
        let zeros = word.count_zeros() as usize - (self.field.len - checksum.count_ones() as usize);
        checksum
            .checked_sub(zeros)
            .ok_or(StoreError::InvalidStorage)
    }

    /// Sets the checksum to the external increment value.
    ///
    /// # Preconditions
    ///
    /// - The checksum bits should be set to one.
    /// - The number of zeros in the word plus the value should fit in the checksum.
    pub fn set(&self, word: &mut u32, value: usize) {
        debug_assert_eq!(self.field.get(*word), self.field.mask() as usize);
        self.field.set(word, word.count_zeros() as usize + value);
    }
}

/// Tracks the number of bits used so far.
///
/// # Features
///
/// Only available for tests.
#[cfg(any(doc, test))]
pub struct Length {
    /// The position of the next available bit.
    pub pos: usize,
}

/// Helps defining contiguous bit fields.
///
/// It takes a sequence of bit field descriptors as argument. A bit field descriptor is one of the
/// following:
/// - `$name: Bit,` to define a bit
/// - `$name: Field <= $max,` to define a bit field of minimum length to store `$max`
/// - `$name: Checksum <= $max,` to define a checksum of minimum length to store `$max`
/// - `$name: Length,` to define a length tracker
/// - `$name: Const = [$bits],` to define a constant bit field with value `$bits` (a sequence of
///   space-separated bits)
#[cfg_attr(doc, macro_export)] // For `cargo doc` to produce documentation.
macro_rules! bitfield {
    ($($input: tt)*) => {
        bitfield_impl! { []{ pos: 0 }[$($input)*] }
    };
}

macro_rules! bitfield_impl {
    // Main rules:
    // - Input are bit field descriptors
    // - Position is the number of bits used by prior bit fields
    // - Output are the bit field definitions
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Bit, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Bit = Bit { pos: $pos };]
            { pos: $pos + 1 }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Field <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Field = Field { pos: $pos, len: num_bits($max) };]
            { pos: $pos + $name.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$name: ident: Checksum <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Checksum = Checksum {
                field: Field { pos: $pos, len: num_bits($max) }
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$(#[$meta: meta])* $name: ident: Length, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* $(#[$meta])* const $name: Length = Length { pos: $pos };]
            { pos: $pos }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Const = $bits: tt, $($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name []$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[]) => { $($output)* };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits
    // - Output is the reversed sequence of bits
    (Reverse $name: ident [$($output_bits: tt)*] [$bit: tt $($input_bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name [$bit $($output_bits)*][$($input_bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    (Reverse $name: ident $bits: tt []
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: 0, val: 0 }$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits in reversed order
    // - Output is the constant bit field definition with the sequence of bits as value
    (Const $name: ident { len: $len: expr, val: $val: expr }[]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Const = Const {
                field: Field { pos: $pos, len: $len },
                value: $val,
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    (Const $name: ident { len: $len: expr, val: $val: expr }[$bit: tt $($bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: $len + 1, val: $val * 2 + $bit }[$($bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
}

/// Counts the number of bits equal to zero in a byte slice.
pub fn count_zeros(slice: &[u8]) -> usize {
    slice.iter().map(|&x| x.count_zeros() as usize).sum()
}

/// Returns the number of bits necessary to represent a number.
pub const fn num_bits(x: usize) -> usize {
    8 * core::mem::size_of::<usize>() - x.leading_zeros() as usize
}

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

    #[test]
    fn field_ok() {
        let field = Field { pos: 3, len: 5 };
        assert_eq!(field.get(0x00000000), 0);
        assert_eq!(field.get(0x00000007), 0);
        assert_eq!(field.get(0x00000008), 1);
        assert_eq!(field.get(0x000000f8), 0x1f);
        assert_eq!(field.get(0x0000ff37), 6);
        let mut word = 0xffffffff;
        field.set(&mut word, 3);
        assert_eq!(word, 0xffffff1f);
    }

    #[test]
    fn const_ok() {
        let field = Const {
            field: Field { pos: 3, len: 5 },
            value: 9,
        };
        assert!(!field.check(0x00000000));
        assert!(!field.check(0x0000ffff));
        assert!(field.check(0x00000048));
        assert!(field.check(0x0000ff4f));
        let mut word = 0xffffffff;
        field.set(&mut word);
        assert_eq!(word, 0xffffff4f);
    }

    #[test]
    fn bit_ok() {
        let bit = Bit { pos: 3 };
        assert!(bit.get(0x00000000));
        assert!(bit.get(0xfffffff7));
        assert!(!bit.get(0x00000008));
        assert!(!bit.get(0xffffffff));
        let mut word = 0xffffffff;
        bit.set(&mut word);
        assert_eq!(word, 0xfffffff7);
    }

    #[test]
    fn checksum_ok() {
        let field = Checksum {
            field: Field { pos: 3, len: 5 },
        };
        assert_eq!(field.get(0x00000000), Err(StoreError::InvalidStorage));
        assert_eq!(field.get(0xffffffff), Ok(31));
        assert_eq!(field.get(0xffffff07), Ok(0));
        assert_eq!(field.get(0xffffff0f), Ok(1));
        assert_eq!(field.get(0x00ffff67), Ok(4));
        assert_eq!(field.get(0x7fffff07), Err(StoreError::InvalidStorage));
        let mut word = 0x0fffffff;
        field.set(&mut word, 4);
        assert_eq!(word, 0x0fffff47);
    }

    #[test]
    fn bitfield_ok() {
        bitfield! {
            FIELD: Field <= 127,
            CONST: Const = [0 1 0 1],
            BIT: Bit,
            CHECKSUM: Checksum <= 58,
            LENGTH: Length,
        }
        assert_eq!(FIELD.pos, 0);
        assert_eq!(FIELD.len, 7);
        assert_eq!(CONST.field.pos, 7);
        assert_eq!(CONST.field.len, 4);
        assert_eq!(CONST.value, 10);
        assert_eq!(BIT.pos, 11);
        assert_eq!(CHECKSUM.field.pos, 12);
        assert_eq!(CHECKSUM.field.len, 6);
        assert_eq!(LENGTH.pos, 18);
    }

    #[test]
    fn count_zeros_ok() {
        assert_eq!(count_zeros(&[0xff, 0xff]), 0);
        assert_eq!(count_zeros(&[0xff, 0xfe]), 1);
        assert_eq!(count_zeros(&[0x7f, 0xff]), 1);
        assert_eq!(count_zeros(&[0x12, 0x48]), 12);
        assert_eq!(count_zeros(&[0x00, 0x00]), 16);
    }

    #[test]
    fn num_bits_ok() {
        assert_eq!(num_bits(0), 0);
        assert_eq!(num_bits(1), 1);
        assert_eq!(num_bits(2), 2);
        assert_eq!(num_bits(3), 2);
        assert_eq!(num_bits(4), 3);
        assert_eq!(num_bits(5), 3);
        assert_eq!(num_bits(8), 4);
        assert_eq!(num_bits(9), 4);
        assert_eq!(num_bits(16), 5);
    }
}
<commit_msg>Rename Const to ConstField<commit_after>// Copyright 2019-2020 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.

//! Helps manipulate bit fields in 32-bits words.
// TODO(ia0): Remove when the module is used.
#![cfg_attr(not(test), allow(dead_code, unused_macros))]

use crate::{StoreError, StoreResult};

/// Represents a bit field.
///
/// A bit field is a contiguous sequence of bits in a 32-bits word.
///
/// # Invariant
///
/// - The bit field must fit in a 32-bits word: `pos + len < 32`.
pub struct Field {
    /// The position of the bit field.
    pub pos: usize,

    /// The length of the bit field.
    pub len: usize,
}

impl Field {
    /// Reads the value of a bit field.
    pub fn get(&self, word: u32) -> usize {
        ((word >> self.pos) & self.mask()) as usize
    }

    /// Sets the value of a bit field.
    ///
    /// # Preconditions
    ///
    /// - The value must fit in the bit field: `num_bits(value) < self.len`.
    /// - The value must only change bits from 1 to 0: `self.get(*word) & value == value`.
    pub fn set(&self, word: &mut u32, value: usize) {
        let value = value as u32;
        debug_assert_eq!(value & self.mask(), value);
        let mask = !(self.mask() << self.pos);
        *word &= mask | (value << self.pos);
        debug_assert_eq!(self.get(*word), value as usize);
    }

    /// Returns a bit mask the length of the bit field.
    ///
    /// The mask is meant to be applied on a value. It should be shifted to be applied to the bit
    /// field.
    fn mask(&self) -> u32 {
        (1 << self.len) - 1
    }
}

/// Represents a constant bit field.
///
/// # Invariant
///
/// - The value must fit in the bit field: `num_bits(value) <= field.len`.
pub struct ConstField {
    /// The bit field.
    pub field: Field,

    /// The constant value.
    pub value: usize,
}

impl ConstField {
    /// Checks that the bit field has its value.
    pub fn check(&self, word: u32) -> bool {
        self.field.get(word) == self.value
    }

    /// Sets the bit field to its value.
    pub fn set(&self, word: &mut u32) {
        self.field.set(word, self.value);
    }
}

/// Represents a single bit.
///
/// # Invariant
///
/// - The bit must fit in a 32-bits word: `pos < 32`.
pub struct Bit {
    /// The position of the bit.
    pub pos: usize,
}

impl Bit {
    /// Returns whether the value of the bit is zero.
    pub fn get(&self, word: u32) -> bool {
        word & (1 << self.pos) == 0
    }

    /// Sets the value of the bit to zero.
    pub fn set(&self, word: &mut u32) {
        *word &= !(1 << self.pos);
    }
}

/// Represents a checksum.
///
/// A checksum is a bit field counting how many bits are set to zero in the word (except in the
/// checksum itself) plus some external increment. It essentially behaves like a bit field storing
/// the external increment.
pub struct Checksum {
    /// The bit field
    pub field: Field,
}

impl Checksum {
    /// Reads the external increment from the checksum.
    ///
    /// # Errors
    ///
    /// Returns `InvalidStorage` if the external increment would be negative.
    pub fn get(&self, word: u32) -> StoreResult<usize> {
        let checksum = self.field.get(word);
        let zeros = word.count_zeros() as usize - (self.field.len - checksum.count_ones() as usize);
        checksum
            .checked_sub(zeros)
            .ok_or(StoreError::InvalidStorage)
    }

    /// Sets the checksum to the external increment value.
    ///
    /// # Preconditions
    ///
    /// - The bits of the checksum bit field should be set to one: `self.field.get(*word) ==
    ///   self.field.mask()`.
    /// - The checksum value should fit in the checksum bit field: `num_bits(word.count_zeros() +
    ///   value) < self.field.len`.
    pub fn set(&self, word: &mut u32, value: usize) {
        debug_assert_eq!(self.field.get(*word), self.field.mask() as usize);
        self.field.set(word, word.count_zeros() as usize + value);
    }
}

/// Tracks the number of bits used so far.
///
/// # Features
///
/// Only available for tests.
#[cfg(any(doc, test))]
pub struct Length {
    /// The position of the next available bit.
    pub pos: usize,
}

/// Helps defining contiguous bit fields.
///
/// It takes a sequence of bit field descriptors as argument. A bit field descriptor is one of the
/// following:
/// - `$name: Bit,` to define a bit
/// - `$name: Field <= $max,` to define a bit field of minimum length to store `$max`
/// - `$name: Checksum <= $max,` to define a checksum of minimum length to store `$max`
/// - `$name: Length,` to define a length tracker
/// - `$name: ConstField = [$bits],` to define a constant bit field with value `$bits` (a sequence
///   of space-separated bits)
#[cfg_attr(doc, macro_export)] // For `cargo doc` to produce documentation.
macro_rules! bitfield {
    ($($input: tt)*) => {
        bitfield_impl! { []{ pos: 0 }[$($input)*] }
    };
}

macro_rules! bitfield_impl {
    // Main rules:
    // - Input are bit field descriptors
    // - Position is the number of bits used by prior bit fields
    // - Output are the bit field definitions
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Bit, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Bit = Bit { pos: $pos };]
            { pos: $pos + 1 }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Field <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Field = Field { pos: $pos, len: num_bits($max) };]
            { pos: $pos + $name.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$name: ident: Checksum <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Checksum = Checksum {
                field: Field { pos: $pos, len: num_bits($max) }
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$(#[$meta: meta])* $name: ident: Length, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* $(#[$meta])* const $name: Length = Length { pos: $pos };]
            { pos: $pos }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$name: ident: ConstField = $bits: tt, $($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name []$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[]) => { $($output)* };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits
    // - Output is the reversed sequence of bits
    (Reverse $name: ident [$($output_bits: tt)*] [$bit: tt $($input_bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name [$bit $($output_bits)*][$($input_bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    (Reverse $name: ident $bits: tt []
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            ConstField $name { len: 0, val: 0 }$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits in reversed order
    // - Output is the constant bit field definition with the sequence of bits as value
    (ConstField $name: ident { len: $len: expr, val: $val: expr }[]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: ConstField = ConstField {
                field: Field { pos: $pos, len: $len },
                value: $val,
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    (ConstField $name: ident { len: $len: expr, val: $val: expr }[$bit: tt $($bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            ConstField $name { len: $len + 1, val: $val * 2 + $bit }[$($bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
}

/// Counts the number of bits equal to zero in a byte slice.
pub fn count_zeros(slice: &[u8]) -> usize {
    slice.iter().map(|&x| x.count_zeros() as usize).sum()
}

/// Returns the number of bits necessary to represent a number.
pub const fn num_bits(x: usize) -> usize {
    8 * core::mem::size_of::<usize>() - x.leading_zeros() as usize
}

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

    #[test]
    fn field_ok() {
        let field = Field { pos: 3, len: 5 };
        assert_eq!(field.get(0x00000000), 0);
        assert_eq!(field.get(0x00000007), 0);
        assert_eq!(field.get(0x00000008), 1);
        assert_eq!(field.get(0x000000f8), 0x1f);
        assert_eq!(field.get(0x0000ff37), 6);
        let mut word = 0xffffffff;
        field.set(&mut word, 3);
        assert_eq!(word, 0xffffff1f);
    }

    #[test]
    fn const_field_ok() {
        let field = ConstField {
            field: Field { pos: 3, len: 5 },
            value: 9,
        };
        assert!(!field.check(0x00000000));
        assert!(!field.check(0x0000ffff));
        assert!(field.check(0x00000048));
        assert!(field.check(0x0000ff4f));
        let mut word = 0xffffffff;
        field.set(&mut word);
        assert_eq!(word, 0xffffff4f);
    }

    #[test]
    fn bit_ok() {
        let bit = Bit { pos: 3 };
        assert!(bit.get(0x00000000));
        assert!(bit.get(0xfffffff7));
        assert!(!bit.get(0x00000008));
        assert!(!bit.get(0xffffffff));
        let mut word = 0xffffffff;
        bit.set(&mut word);
        assert_eq!(word, 0xfffffff7);
    }

    #[test]
    fn checksum_ok() {
        let field = Checksum {
            field: Field { pos: 3, len: 5 },
        };
        assert_eq!(field.get(0x00000000), Err(StoreError::InvalidStorage));
        assert_eq!(field.get(0xffffffff), Ok(31));
        assert_eq!(field.get(0xffffff07), Ok(0));
        assert_eq!(field.get(0xffffff0f), Ok(1));
        assert_eq!(field.get(0x00ffff67), Ok(4));
        assert_eq!(field.get(0x7fffff07), Err(StoreError::InvalidStorage));
        let mut word = 0x0fffffff;
        field.set(&mut word, 4);
        assert_eq!(word, 0x0fffff47);
    }

    #[test]
    fn bitfield_ok() {
        bitfield! {
            FIELD: Field <= 127,
            CONST_FIELD: ConstField = [0 1 0 1],
            BIT: Bit,
            CHECKSUM: Checksum <= 58,
            LENGTH: Length,
        }
        assert_eq!(FIELD.pos, 0);
        assert_eq!(FIELD.len, 7);
        assert_eq!(CONST_FIELD.field.pos, 7);
        assert_eq!(CONST_FIELD.field.len, 4);
        assert_eq!(CONST_FIELD.value, 10);
        assert_eq!(BIT.pos, 11);
        assert_eq!(CHECKSUM.field.pos, 12);
        assert_eq!(CHECKSUM.field.len, 6);
        assert_eq!(LENGTH.pos, 18);
    }

    #[test]
    fn count_zeros_ok() {
        assert_eq!(count_zeros(&[0xff, 0xff]), 0);
        assert_eq!(count_zeros(&[0xff, 0xfe]), 1);
        assert_eq!(count_zeros(&[0x7f, 0xff]), 1);
        assert_eq!(count_zeros(&[0x12, 0x48]), 12);
        assert_eq!(count_zeros(&[0x00, 0x00]), 16);
    }

    #[test]
    fn num_bits_ok() {
        assert_eq!(num_bits(0), 0);
        assert_eq!(num_bits(1), 1);
        assert_eq!(num_bits(2), 2);
        assert_eq!(num_bits(3), 2);
        assert_eq!(num_bits(4), 3);
        assert_eq!(num_bits(5), 3);
        assert_eq!(num_bits(8), 4);
        assert_eq!(num_bits(9), 4);
        assert_eq!(num_bits(16), 5);
    }
}
<|endoftext|>


<commit_before>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::*;
use command::{Command, parse_command};
use message::Message;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;
    let welcome_msg = "Welcome to Smazy\nPlease enter a username:\n".to_string();
    stream.write(&welcome_msg.into_bytes()).expect("Error writing to stream");
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
    let mut username = "".to_string();
    let (sender_to_client, receiver_from_router) = channel();

    loop {
        match reader.read_line(&mut username) {
            Ok(_) => {
                username = username.trim().to_string();

                {
                    let mut guard = chat_map.lock().expect("Error locking chatmap");

                    if is_valid_username(&*guard, username.to_string()) {

                        let client_info = ClientInfo{
                            partner: None,
                            sender_to_client: sender_to_client,
                            blocked_users: Vec::new(),
                        };
                        guard.insert(username.to_string(), client_info);

                        println!("New user: {}", username.to_string());
                        break;
                    }
                    else {
                        let invalid_msg = "Invalid username. Please try again.\n".to_string();
                        stream.write(&invalid_msg.into_bytes()).expect("Error writing to stream");
                        username = "".to_string();
                    }
                }
            },
            Err(e) => println!("{}", e)
        }
    }

    let partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
        username.to_string(), chat_map);

    chat(stream, username, partner, sender_to_router, receiver_from_router, chat_map);
}

fn display_available(stream: TcpStream, chat_map: &Arc<Mutex<ChatMap>>, username: String) {
    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).expect("Error writing to stream");

    let avail;

    {
        let guard = chat_map.lock().expect("Error locking chatmap");
        avail = available_users(&*guard, username);
    }

    for user in avail {
        stream.write(&(user + "\n").to_string().into_bytes()).expect("Error writing to stream");
    }
}

fn choose_chat_partner(stream: TcpStream,
                       username: String,
                       chat_map: &Arc<Mutex<ChatMap>>) -> String {
    let no_partner = Arc::new(Mutex::new(true));

    // Allow you to start a chat
    {
        let chat_map = chat_map.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");
        let no_partner = no_partner.clone();

        thread::spawn(move|| {

            loop {
                {
                    let guard = no_partner.lock().unwrap();
                    if !*guard {
                        break;
                    }
                }
                display_available(stream.try_clone().expect("Error cloning stream"),
                    &chat_map, username.to_string());

                let select_msg = "Select who you want to chat with:\n".to_string();
                stream.write(&select_msg.into_bytes()).expect("Error writing to stream");

                let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
                let mut partner = "".to_string();

                match reader.read_line(&mut partner) {
                    Ok(_) => {
                        partner = partner.trim().to_string();

                        if try_select_partner(&chat_map, username.to_string(), partner.to_string()) {
                            return partner;
                        }
                    },
                    Err(e) => println!("{}", e)
                }
            }
            return "Garbage".to_string();
        });
    }

    // Check if someone else has started a chat with you
    {
        let chat_map = chat_map.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");
        let no_partner = no_partner.clone();

        let h = thread::spawn(move|| {
            loop {
                let guard = chat_map.lock().expect("Error locking chatmap");

                match &guard.get(&username).unwrap().partner {
                    &Some(ref p) => {
                        println!("{} chatting with {}", &username, p);
                        let chatting_msg = "Now chatting with ".to_string() + p + ". Press enter to start chatting.\n";
                        stream.write(&chatting_msg.into_bytes()).expect("Error writing to stream");
                        {
                            let mut guard = no_partner.lock().unwrap();
                            *guard = false;
                        }
                        return p.to_string();
                    }
                    &None => continue
                }
            }
        });

        return h.join().unwrap();
    }
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    let success: bool;
    {
        let mut guard = chat_map.lock().expect("Error locking chatmap");

        {
            let my_client_info = guard.get(&username).unwrap();
            if my_client_info.blocked_users.contains(&partner) {
                return false;
            }
        }

        match guard.get_mut(&partner) {
            Some(clientinfo) => {
                if &partner[..] != &username[..]
                    && !clientinfo.blocked_users.contains(&username.to_string())
                    && (clientinfo.partner == None
                        || clientinfo.partner == Some(username.to_string())) {
                    clientinfo.partner = Some(username.to_string());
                    success = true;
                } else {
                    success = false;
                }
            },
            None => success = false
        }
    }

    if success {
        let mut guard = chat_map.lock().expect("Error locking chatmap");
        println!("Assigning partner {} to {}", partner, username);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
    }
    success
}

fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>,
        chat_map: &Arc<Mutex<ChatMap>>) {

    let mut partner = partner.to_string();

    // Send messages
    {
        let username = username.clone();
        let chat_map = chat_map.clone();
        let stream = stream.try_clone().expect("Error cloning stream");
        let reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());
                if line.len() > 0 {
                    if &line[0..1] == "/" {
                        match parse_command(msg.content().to_string()) {
                            Command::Quit => {
                                println!("Quit command");
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                }

                                let quit_msg = "I have quit out of the conversation. Type /q to quit.";

                                let msg = Message::new(time::now().asctime().to_string(),
                                                           username.to_string(),
                                                           partner.to_string(),
                                                           quit_msg.to_string());

                                sender_to_router.send(msg).expect("Error sending message");

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },
                            Command::DisplayAvailable => {
                                display_available(stream.try_clone().expect("Error cloning stream"),
                                    &chat_map, username.to_string());
                            },
                            Command::Block => {
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                    let mut client_info = guard.get_mut(&username).unwrap();
                                    client_info.blocked_users.push(partner.to_string());
                                }

                                let quit_msg = "I have blocked you and quit out of the conversation. Type /q to quit.";

                                let msg = Message::new(time::now().asctime().to_string(),
                                                           username.to_string(),
                                                           partner.to_string(),
                                                           quit_msg.to_string());

                                sender_to_router.send(msg).expect("Error sending message");

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },
                            Command::Unrecognized => println!("Unrecognized command")
                        }
                    } else {
                        sender_to_router.send(msg).expect("Error sending message");
                    }
                }
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().expect("Error cloning stream");
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).expect("Error writing to stream");
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};

    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_success_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_success_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));
    }

    #[test]
    fn try_select_partner_test_succes_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: b can choose a when they are already partners
        //   (This allows someone to confirm a chat request)
        assert!(try_select_partner(&cm, "b".to_string(), "a".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a,
            blocked_users: Vec::new()
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b,
            blocked_users: Vec::new()
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c,
            blocked_users: Vec::new()
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d,
            blocked_users: Vec::new()
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e,
            blocked_users: Vec::new()
        });
        cm
    }
}<commit_msg>A few comments<commit_after>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::*;
use command::{Command, parse_command};
use message::Message;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;
    let welcome_msg = "Welcome to Smazy\nPlease enter a username:\n".to_string();
    stream.write(&welcome_msg.into_bytes()).expect("Error writing to stream");
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
    let mut username = "".to_string();
    let (sender_to_client, receiver_from_router) = channel();

    loop {
        match reader.read_line(&mut username) {
            Ok(_) => {
                username = username.trim().to_string();

                {
                    let mut guard = chat_map.lock().expect("Error locking chatmap");

                    if is_valid_username(&*guard, username.to_string()) {

                        let client_info = ClientInfo{
                            partner: None,
                            sender_to_client: sender_to_client,
                            blocked_users: Vec::new(),
                        };

                        guard.insert(username.to_string(), client_info);

                        println!("New user: {}", username.to_string());
                        break;
                    }
                    else {
                        let invalid_msg = "Invalid username. Please try again.\n".to_string();
                        stream.write(&invalid_msg.into_bytes()).expect("Error writing to stream");
                        username = "".to_string();
                    }
                }
            },
            Err(e) => println!("{}", e)
        }
    }

    let partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
        username.to_string(), chat_map);

    chat(stream, username, partner, sender_to_router, receiver_from_router, chat_map);
}

fn display_available(stream: TcpStream, chat_map: &Arc<Mutex<ChatMap>>, username: String) {
    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).expect("Error writing to stream");

    let avail;

    {
        let guard = chat_map.lock().expect("Error locking chatmap");
        avail = available_users(&*guard, username);
    }

    for user in avail {
        stream.write(&(user + "\n").to_string().into_bytes()).expect("Error writing to stream");
    }
}

fn choose_chat_partner(stream: TcpStream,
                       username: String,
                       chat_map: &Arc<Mutex<ChatMap>>) -> String {

    let no_partner = Arc::new(Mutex::new(true));

    // Allow you to start a chat
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        thread::spawn(move|| {

            loop {
                {
                    let guard = no_partner.lock().unwrap();
                    if !*guard {
                        break;
                    }
                }
                display_available(stream.try_clone().expect("Error cloning stream"),
                    &chat_map, username.to_string());

                let select_msg = "Select who you want to chat with:\n".to_string();
                stream.write(&select_msg.into_bytes()).expect("Error writing to stream");

                let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
                let mut partner = "".to_string();

                match reader.read_line(&mut partner) {
                    Ok(_) => {
                        partner = partner.trim().to_string();

                        if try_select_partner(&chat_map, username.to_string(), partner.to_string()) {
                            return partner;
                        }
                    },
                    Err(e) => println!("{}", e)
                }
            }

            // Value returned from this thread is unused
            // Rust compiler complains if we don't return a String
            return "Unused".to_string();
        });
    }

    // Check if someone else has started a chat with you
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        let h = thread::spawn(move|| {
            loop {
                let guard = chat_map.lock().expect("Error locking chatmap");

                match &guard.get(&username).unwrap().partner {
                    &Some(ref p) => {
                        println!("{} chatting with {}", &username, p);
                        let chatting_msg = "Now chatting with ".to_string() + p + ". Press enter to start chatting.\n";
                        stream.write(&chatting_msg.into_bytes()).expect("Error writing to stream");
                        {
                            let mut guard = no_partner.lock().unwrap();
                            *guard = false;
                        }
                        return p.to_string();
                    }
                    &None => continue
                }
            }
        });

        // Return partner
        return h.join().unwrap();
    }
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    let success: bool;
    {
        let mut guard = chat_map.lock().expect("Error locking chatmap");

        {
            // Not allowed to select someone you have blocked
            let my_client_info = guard.get(&username).unwrap();
            if my_client_info.blocked_users.contains(&partner) {
                return false;
            }
        }

        match guard.get_mut(&partner) {
            Some(clientinfo) => {
                if &partner[..] != &username[..]
                    && !clientinfo.blocked_users.contains(&username.to_string())
                    && (clientinfo.partner == None
                        || clientinfo.partner == Some(username.to_string())) {

                    // Set their partner to you
                    clientinfo.partner = Some(username.to_string());
                    success = true;
                } else {
                    success = false;
                }
            },
            None => success = false
        }
    }

    if success {
        // Set your partner to them
        let mut guard = chat_map.lock().expect("Error locking chatmap");
        println!("Assigning partner {} to {}", partner, username);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
    }
    success
}

fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>,
        chat_map: &Arc<Mutex<ChatMap>>) {

    let mut partner = partner.to_string();

    // Send messages
    {
        let username = username.clone();
        let chat_map = chat_map.clone();
        let stream = stream.try_clone().expect("Error cloning stream");
        let reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));

        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());
                if line.len() > 0 {
                    if &line[0..1] == "/" {
                        match parse_command(msg.content().to_string()) {
                            Command::Quit => {
                                println!("Quit command");
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                }

                                let quit_msg = "I have quit out of the conversation. Type /q to quit.";

                                let msg = Message::new(time::now().asctime().to_string(),
                                                           username.to_string(),
                                                           partner.to_string(),
                                                           quit_msg.to_string());

                                sender_to_router.send(msg).expect("Error sending message");

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },
                            Command::DisplayAvailable => {
                                display_available(stream.try_clone().expect("Error cloning stream"),
                                    &chat_map, username.to_string());
                            },
                            Command::Block => {
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                    let mut client_info = guard.get_mut(&username).unwrap();
                                    client_info.blocked_users.push(partner.to_string());
                                }

                                let quit_msg = "I have blocked you and quit out of the conversation. Type /q to quit.";

                                let msg = Message::new(time::now().asctime().to_string(),
                                                           username.to_string(),
                                                           partner.to_string(),
                                                           quit_msg.to_string());

                                sender_to_router.send(msg).expect("Error sending message");

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },
                            Command::Unrecognized => println!("Unrecognized command")
                        }
                    } else {
                        sender_to_router.send(msg).expect("Error sending message");
                    }
                }
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().expect("Error cloning stream");
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).expect("Error writing to stream");
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};

    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_success_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_success_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));
    }

    #[test]
    fn try_select_partner_test_succes_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: b can choose a when they are already partners
        //   (This allows someone to confirm a chat request)
        assert!(try_select_partner(&cm, "b".to_string(), "a".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a,
            blocked_users: Vec::new()
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b,
            blocked_users: Vec::new()
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c,
            blocked_users: Vec::new()
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d,
            blocked_users: Vec::new()
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e,
            blocked_users: Vec::new()
        });
        cm
    }
}<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    completion::on_completion,
    goto_definition::on_goto_definition,
    hover::on_hover,
    lsp::{
        set_not_started_status, CompletionOptions, Connection, GotoDefinition, HoverRequest,
        InitializeParams, Message, ServerCapabilities, ServerResponse, TextDocumentSyncCapability,
        TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    references::on_references,
    shutdown::{on_exit, on_shutdown},
    text_documents::on_did_save_text_document,
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::info;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{Completion, References, Shutdown},
};
use relay_compiler::config::Config;
use std::sync::Arc;

mod lsp_request_dispatch;
use lsp_request_dispatch::LSPRequestDispatch;
mod lsp_notification_dispatch;
use lsp_notification_dispatch::LSPNotificationDispatch;
mod lsp_state;
pub(crate) use lsp_state::{ExtraDataProvider, LSPState};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: Some(vec!["(".into(), "{".into()]),
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });

    server_capabilities.hover_provider = Some(true);
    server_capabilities.definition_provider = Some(true);
    server_capabilities.references_provider = Some(true);

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    config: Config,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: ExtraDataProvider,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    info!(
        "Running language server with config root {:?}",
        config.root_dir
    );

    let mut lsp_state = LSPState::new(
        config,
        &connection.sender,
        Arc::clone(&perf_logger),
        extra_data_provider,
    );

    set_not_started_status(&connection.sender);

    for msg in connection.receiver {
        info!("LSP message received {:?}", msg);
        match msg {
            Message::Request(req) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Message::Notification(notification) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            _ => {
                // Ignore responses for now
            }
        }
    }

    Ok(())
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            lsp_request_event.number("lsp_error_code", error.code as usize);
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    // The only possible error (for now) is not handling the notification.
    // N.B. is_ok is correct here.
    if notification_result.is_ok() {
        lsp_notification_event.string("lsp_outcome", "error".to_string());
        lsp_notification_event.number("lsp_error_code", ErrorCode::MethodNotFound as usize);
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), ()> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    info!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<commit_msg>Trigger autocomplete on new line and ","<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    completion::on_completion,
    goto_definition::on_goto_definition,
    hover::on_hover,
    lsp::{
        set_not_started_status, CompletionOptions, Connection, GotoDefinition, HoverRequest,
        InitializeParams, Message, ServerCapabilities, ServerResponse, TextDocumentSyncCapability,
        TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    references::on_references,
    shutdown::{on_exit, on_shutdown},
    text_documents::on_did_save_text_document,
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::info;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{Completion, References, Shutdown},
};
use relay_compiler::config::Config;
use std::sync::Arc;

mod lsp_request_dispatch;
use lsp_request_dispatch::LSPRequestDispatch;
mod lsp_notification_dispatch;
use lsp_notification_dispatch::LSPNotificationDispatch;
mod lsp_state;
pub(crate) use lsp_state::{ExtraDataProvider, LSPState};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: Some(vec!["(".into(), "{".into(), "\n".into(), ",".into()]),
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });

    server_capabilities.hover_provider = Some(true);
    server_capabilities.definition_provider = Some(true);
    server_capabilities.references_provider = Some(true);

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    config: Config,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: ExtraDataProvider,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    info!(
        "Running language server with config root {:?}",
        config.root_dir
    );

    let mut lsp_state = LSPState::new(
        config,
        &connection.sender,
        Arc::clone(&perf_logger),
        extra_data_provider,
    );

    set_not_started_status(&connection.sender);

    for msg in connection.receiver {
        info!("LSP message received {:?}", msg);
        match msg {
            Message::Request(req) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Message::Notification(notification) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            _ => {
                // Ignore responses for now
            }
        }
    }

    Ok(())
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            lsp_request_event.number("lsp_error_code", error.code as usize);
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    // The only possible error (for now) is not handling the notification.
    // N.B. is_ok is correct here.
    if notification_result.is_ok() {
        lsp_notification_event.string("lsp_outcome", "error".to_string());
        lsp_notification_event.number("lsp_error_code", ErrorCode::MethodNotFound as usize);
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), ()> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    info!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Global style data

use context::StyleSystemOptions;
use num_cpus;
use rayon;
use shared_lock::SharedRwLock;
use std::cmp;
use std::env;

/// Global style data
pub struct GlobalStyleData {
    /// How many threads parallel styling can use.
    pub num_threads: usize,

    /// The parallel styling thread pool.
    pub style_thread_pool: Option<rayon::ThreadPool>,

    /// Shared RWLock for CSSOM objects
    pub shared_lock: SharedRwLock,

    /// Global style system options determined by env vars.
    pub options: StyleSystemOptions,
}

lazy_static! {
    /// Global style data
    pub static ref GLOBAL_STYLE_DATA: GlobalStyleData = {
        let stylo_threads = env::var("STYLO_THREADS")
            .map(|s| s.parse::<usize>().expect("invalid STYLO_THREADS value"));
        let num_threads = match stylo_threads {
            Ok(num) => num,
            _ => cmp::max(num_cpus::get() * 3 / 4, 1),
        };

        let pool = if num_threads <= 1 {
            None
        } else {
            let configuration =
                rayon::Configuration::new().num_threads(num_threads);
            let pool = rayon::ThreadPool::new(configuration).ok();
            pool
        };

        GlobalStyleData {
            num_threads: num_threads,
            style_thread_pool: pool,
            shared_lock: SharedRwLock::new(),
            options: StyleSystemOptions::default(),
        }
    };
}
<commit_msg>register stylo threads with the gecko profiler<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Global style data

use context::StyleSystemOptions;
use gecko_bindings::bindings::{Gecko_RegisterProfilerThread, Gecko_UnregisterProfilerThread};
use num_cpus;
use rayon;
use shared_lock::SharedRwLock;
use std::cmp;
use std::env;
use std::ffi::CString;

/// Global style data
pub struct GlobalStyleData {
    /// How many threads parallel styling can use.
    pub num_threads: usize,

    /// The parallel styling thread pool.
    pub style_thread_pool: Option<rayon::ThreadPool>,

    /// Shared RWLock for CSSOM objects
    pub shared_lock: SharedRwLock,

    /// Global style system options determined by env vars.
    pub options: StyleSystemOptions,
}

fn thread_name(index: usize) -> String {
    format!("StyleThread#{}", index)
}

fn thread_startup(index: usize) {
    let name = thread_name(index);
    let name = CString::new(name).unwrap();
    unsafe {
        // Gecko_RegisterProfilerThread copies the passed name here.
        Gecko_RegisterProfilerThread(name.as_ptr());
    }
}

fn thread_shutdown(_: usize) {
    unsafe {
        Gecko_UnregisterProfilerThread();
    }
}

lazy_static! {
    /// Global style data
    pub static ref GLOBAL_STYLE_DATA: GlobalStyleData = {
        let stylo_threads = env::var("STYLO_THREADS")
            .map(|s| s.parse::<usize>().expect("invalid STYLO_THREADS value"));
        let num_threads = match stylo_threads {
            Ok(num) => num,
            _ => cmp::max(num_cpus::get() * 3 / 4, 1),
        };

        let pool = if num_threads <= 1 {
            None
        } else {
            let configuration = rayon::Configuration::new()
                .num_threads(num_threads)
                .thread_name(thread_name)
                .start_handler(thread_startup)
                .exit_handler(thread_shutdown);
            let pool = rayon::ThreadPool::new(configuration).ok();
            pool
        };

        GlobalStyleData {
            num_threads: num_threads,
            style_thread_pool: pool,
            shared_lock: SharedRwLock::new(),
            options: StyleSystemOptions::default(),
        }
    };
}
<|endoftext|>


<commit_before>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};
use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::Deref;
use std::str::FromStr;

use error::*;
use ser::ConfigSerializer;
use source::Source;

use path;
use value::{Table, Value, ValueKind};

#[derive(Clone, Debug)]
enum ConfigKind {
    // A mutable configuration. This is the default.
    Mutable {
        defaults: HashMap<path::Expression, Value>,
        overrides: HashMap<path::Expression, Value>,
        sources: Vec<Box<dyn Source + Send + Sync>>,
    },

    // A frozen configuration.
    // Configuration can no longer be mutated.
    Frozen,
}

impl Default for ConfigKind {
    fn default() -> Self {
        ConfigKind::Mutable {
            defaults: HashMap::new(),
            overrides: HashMap::new(),
            sources: Vec::new(),
        }
    }
}

/// A prioritized configuration repository. It maintains a set of
/// configuration sources, fetches values to populate those, and provides
/// them according to the source's priority.
#[derive(Default, Clone, Debug)]
pub struct Config {
    kind: ConfigKind,

    /// Root of the cached configuration.
    pub cache: Value,
}

impl Config {
    pub fn new() -> Self {
        Self {
            kind: ConfigKind::default(),
            // Config root should be instantiated as an empty table
            // to avoid deserialization errors.
            cache: Value::new(None, Table::new()),
        }
    }

    /// Merge in a configuration property source.
    pub fn merge<T>(&mut self, source: T) -> Result<&mut Config>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()
    }

    /// Refresh the configuration cache with fresh
    /// data from added sources.
    ///
    /// Configuration is automatically refreshed after a mutation
    /// operation (`set`, `merge`, `set_default`, etc.).
    pub fn refresh(&mut self) -> Result<&mut Config> {
        self.cache = match self.kind {
            // TODO: We need to actually merge in all the stuff
            ConfigKind::Mutable {
                ref overrides,
                ref sources,
                ref defaults,
            } => {
                let mut cache: Value = HashMap::<String, Value>::new().into();

                // Add defaults
                for (key, val) in defaults {
                    key.set(&mut cache, val.clone());
                }

                // Add sources
                sources.collect_to(&mut cache)?;

                // Add overrides
                for (key, val) in overrides {
                    key.set(&mut cache, val.clone());
                }

                cache
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        };

        Ok(self)
    }

    pub fn set_default<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut defaults, ..
            } => {
                defaults.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn set<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut overrides, ..
            } => {
                overrides.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn get<'de, T: Deserialize<'de>>(&self, key: &str) -> Result<T> {
        // Parse the key into a path expression
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        let value = expr.get(&self.cache).cloned();

        match value {
            Some(value) => {
                // Deserialize the received value into the requested type
                T::deserialize(value).map_err(|e| e.extend_with_key(key))
            }

            None => Err(ConfigError::NotFound(key.into())),
        }
    }

    pub fn get_str(&self, key: &str) -> Result<String> {
        self.get(key).and_then(Value::into_str)
    }

    pub fn get_int(&self, key: &str) -> Result<i64> {
        self.get(key).and_then(Value::into_int)
    }

    pub fn get_float(&self, key: &str) -> Result<f64> {
        self.get(key).and_then(Value::into_float)
    }

    pub fn get_bool(&self, key: &str) -> Result<bool> {
        self.get(key).and_then(Value::into_bool)
    }

    pub fn get_table(&self, key: &str) -> Result<HashMap<String, Value>> {
        self.get(key).and_then(Value::into_table)
    }

    pub fn get_array(&self, key: &str) -> Result<Vec<Value>> {
        self.get(key).and_then(Value::into_array)
    }

    /// Attempt to deserialize the entire configuration into the requested type.
    pub fn try_into<'de, T: Deserialize<'de>>(self) -> Result<T> {
        T::deserialize(self)
    }

    /// Attempt to serialize the entire configuration from the given type.
    pub fn try_from<T: Serialize>(from: &T) -> Result<Self> {
        let mut serializer = ConfigSerializer::default();
        from.serialize(&mut serializer)?;
        Ok(serializer.output)
    }

    #[deprecated(since = "0.7.0", note = "please use 'try_into' instead")]
    pub fn deserialize<'de, T: Deserialize<'de>>(self) -> Result<T> {
        self.try_into()
    }
}

impl Source for Config {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        self.cache.clone().into_table()
    }
}
<commit_msg>Add a Config::set_once() function to set a value once (and let Config::merge() overwrite it later<commit_after>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};
use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::Deref;
use std::str::FromStr;

use error::*;
use ser::ConfigSerializer;
use source::Source;

use path;
use value::{Table, Value, ValueKind};

#[derive(Clone, Debug)]
enum ConfigKind {
    // A mutable configuration. This is the default.
    Mutable {
        defaults: HashMap<path::Expression, Value>,
        overrides: HashMap<path::Expression, Value>,
        sources: Vec<Box<dyn Source + Send + Sync>>,
    },

    // A frozen configuration.
    // Configuration can no longer be mutated.
    Frozen,
}

impl Default for ConfigKind {
    fn default() -> Self {
        ConfigKind::Mutable {
            defaults: HashMap::new(),
            overrides: HashMap::new(),
            sources: Vec::new(),
        }
    }
}

/// A prioritized configuration repository. It maintains a set of
/// configuration sources, fetches values to populate those, and provides
/// them according to the source's priority.
#[derive(Default, Clone, Debug)]
pub struct Config {
    kind: ConfigKind,

    /// Root of the cached configuration.
    pub cache: Value,
}

impl Config {
    pub fn new() -> Self {
        Self {
            kind: ConfigKind::default(),
            // Config root should be instantiated as an empty table
            // to avoid deserialization errors.
            cache: Value::new(None, Table::new()),
        }
    }

    /// Merge in a configuration property source.
    pub fn merge<T>(&mut self, source: T) -> Result<&mut Config>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()
    }

    /// Refresh the configuration cache with fresh
    /// data from added sources.
    ///
    /// Configuration is automatically refreshed after a mutation
    /// operation (`set`, `merge`, `set_default`, etc.).
    pub fn refresh(&mut self) -> Result<&mut Config> {
        self.cache = match self.kind {
            // TODO: We need to actually merge in all the stuff
            ConfigKind::Mutable {
                ref overrides,
                ref sources,
                ref defaults,
            } => {
                let mut cache: Value = HashMap::<String, Value>::new().into();

                // Add defaults
                for (key, val) in defaults {
                    key.set(&mut cache, val.clone());
                }

                // Add sources
                sources.collect_to(&mut cache)?;

                // Add overrides
                for (key, val) in overrides {
                    key.set(&mut cache, val.clone());
                }

                cache
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        };

        Ok(self)
    }

    pub fn set_default<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut defaults, ..
            } => {
                defaults.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn set<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut overrides, ..
            } => {
                overrides.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn set_once(&mut self, key: &str, value: Value) -> Result<()> {
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        if let Some(ref mut val) = expr.get_mut(&mut self.cache) {
            **val = value;
        } else {
            expr.set(&mut self.cache, value);
        }
        Ok(())
    }

    pub fn get<'de, T: Deserialize<'de>>(&self, key: &str) -> Result<T> {
        // Parse the key into a path expression
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        let value = expr.get(&self.cache).cloned();

        match value {
            Some(value) => {
                // Deserialize the received value into the requested type
                T::deserialize(value).map_err(|e| e.extend_with_key(key))
            }

            None => Err(ConfigError::NotFound(key.into())),
        }
    }

    pub fn get_str(&self, key: &str) -> Result<String> {
        self.get(key).and_then(Value::into_str)
    }

    pub fn get_int(&self, key: &str) -> Result<i64> {
        self.get(key).and_then(Value::into_int)
    }

    pub fn get_float(&self, key: &str) -> Result<f64> {
        self.get(key).and_then(Value::into_float)
    }

    pub fn get_bool(&self, key: &str) -> Result<bool> {
        self.get(key).and_then(Value::into_bool)
    }

    pub fn get_table(&self, key: &str) -> Result<HashMap<String, Value>> {
        self.get(key).and_then(Value::into_table)
    }

    pub fn get_array(&self, key: &str) -> Result<Vec<Value>> {
        self.get(key).and_then(Value::into_array)
    }

    /// Attempt to deserialize the entire configuration into the requested type.
    pub fn try_into<'de, T: Deserialize<'de>>(self) -> Result<T> {
        T::deserialize(self)
    }

    /// Attempt to serialize the entire configuration from the given type.
    pub fn try_from<T: Serialize>(from: &T) -> Result<Self> {
        let mut serializer = ConfigSerializer::default();
        from.serialize(&mut serializer)?;
        Ok(serializer.output)
    }

    #[deprecated(since = "0.7.0", note = "please use 'try_into' instead")]
    pub fn deserialize<'de, T: Deserialize<'de>>(self) -> Result<T> {
        self.try_into()
    }
}

impl Source for Config {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        self.cache.clone().into_table()
    }
}
<|endoftext|>