<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    cast::AsPrimitive,
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: Int> {
    a: T,
    n: T,
}

impl<T: Int> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::Intermediate) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_intermediate()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_intermediate(x).wrapping_mul(a);
        let nm = (n.as_intermediate()) * (m.as_intermediate());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::Intermediate::one() << T::zero().count_zeros() as usize),
            T::Intermediate::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_intermediate(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: Int> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_()
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        debug_assert!(x < self.n.as_());
        let r = T::from_intermediate(
            ((T::Intermediate::from(x)) << T::zero().count_zeros() as usize)
                % self.n.as_intermediate(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n.as_intermediate()).as_()
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 =
                (((a_r as u128) + (b_r as u128)) % <T as AsPrimitive<u128>>::as_(self.n)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce(a.as_intermediate() * b.as_intermediate());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2: u64 = ((T::Intermediate::from(a_r) * T::Intermediate::from(b_r))
                % self.n.as_intermediate())
            .as_();
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
impl OverflowingAdd for u32 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u64 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u128 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}

pub(crate) trait Int:
    AsPrimitive<u64>
    + AsPrimitive<u128>
    + Display
    + Debug
    + PrimInt
    + OverflowingAdd
    + WrappingNeg
    + WrappingSub
    + WrappingMul
{
    type Intermediate: From<u64>
        + AsPrimitive<u64>
        + Display
        + Debug
        + PrimInt
        + OverflowingAdd
        + WrappingNeg
        + WrappingSub
        + WrappingMul;

    fn as_intermediate(self) -> Self::Intermediate;
    fn from_intermediate(n: Self::Intermediate) -> Self;
    fn from_u64(n: u64) -> Self;
}
impl Int for u64 {
    type Intermediate = u128;

    fn as_intermediate(self) -> u128 {
        self as _
    }
    fn from_intermediate(n: u128) -> u64 {
        n as _
    }
    fn from_u64(n: u64) -> u64 {
        n
    }
}
impl Int for u32 {
    type Intermediate = u64;

    fn as_intermediate(self) -> u64 {
        self as _
    }
    fn from_intermediate(n: u64) -> u32 {
        n as _
    }
    fn from_u64(n: u64) -> u32 {
        n as _
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }

    #[test]
    fn test_inverter_u32() {
        test_inverter::<u32>()
    }

    #[test]
    fn test_inverter_u64() {
        test_inverter::<u64>()
    }

    fn test_add<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_add_m32() {
        test_add::<Montgomery<u32>>()
    }

    #[test]
    fn test_add_m64() {
        test_add::<Montgomery<u64>>()
    }

    fn test_mult<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_mult_m32() {
        test_mult::<Montgomery<u32>>()
    }

    #[test]
    fn test_mult_m64() {
        test_mult::<Montgomery<u64>>()
    }

    fn test_roundtrip<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }

    #[test]
    fn test_roundtrip_m32() {
        test_roundtrip::<Montgomery<u32>>()
    }

    #[test]
    fn test_roundtrip_m64() {
        test_roundtrip::<Montgomery<u64>>()
    }
}
<commit_msg>factor::numeric: Split Int and DoubleInt traits<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    cast::{AsPrimitive, ToPrimitive},
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: DoubleInt> {
    a: T,
    n: T,
}

impl<T: DoubleInt> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::Double) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_double()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_double(x).wrapping_mul(a);
        let nm = (n.as_double()) * (m.as_double());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::Double::one() << T::zero().count_zeros() as usize),
            T::Double::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_double(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: DoubleInt> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_()
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        debug_assert!(x < self.n.as_());
        let r = T::from_double(
            ((T::Double::from(x)) << T::zero().count_zeros() as usize) % self.n.as_double(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n.as_double()).as_()
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 =
                (((a_r as u128) + (b_r as u128)) % <T as AsPrimitive<u128>>::as_(self.n)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce(a.as_double() * b.as_double());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2: u64 = ((T::Double::from(a_r) * T::Double::from(b_r)) % self.n.as_double())
                .to_u64()
                .unwrap();
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
impl OverflowingAdd for u32 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u64 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u128 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}

pub(crate) trait Int:
    AsPrimitive<u64>
    + AsPrimitive<u128>
    + Display
    + Debug
    + PrimInt
    + OverflowingAdd
    + WrappingNeg
    + WrappingSub
    + WrappingMul
{
}

pub(crate) trait DoubleInt: Int {
    type Double: Int + From<u64> + ToPrimitive;

    fn as_double(self) -> Self::Double;
    fn from_double(n: Self::Double) -> Self;
    fn from_u64(n: u64) -> Self;
}

impl Int for u32 {}
impl Int for u64 {}
impl Int for u128 {}

impl DoubleInt for u64 {
    type Double = u128;

    fn as_double(self) -> u128 {
        self as _
    }
    fn from_double(n: u128) -> u64 {
        n as _
    }
    fn from_u64(n: u64) -> u64 {
        n
    }
}
impl DoubleInt for u32 {
    type Double = u64;

    fn as_double(self) -> u64 {
        self as _
    }
    fn from_double(n: u64) -> u32 {
        n as _
    }
    fn from_u64(n: u64) -> u32 {
        n as _
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }

    #[test]
    fn test_inverter_u32() {
        test_inverter::<u32>()
    }

    #[test]
    fn test_inverter_u64() {
        test_inverter::<u64>()
    }

    fn test_add<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_add_m32() {
        test_add::<Montgomery<u32>>()
    }

    #[test]
    fn test_add_m64() {
        test_add::<Montgomery<u64>>()
    }

    fn test_mult<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_mult_m32() {
        test_mult::<Montgomery<u32>>()
    }

    #[test]
    fn test_mult_m64() {
        test_mult::<Montgomery<u64>>()
    }

    fn test_roundtrip<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }

    #[test]
    fn test_roundtrip_m32() {
        test_roundtrip::<Montgomery<u32>>()
    }

    #[test]
    fn test_roundtrip_m64() {
        test_roundtrip::<Montgomery<u64>>()
    }
}
<|endoftext|>


<commit_before>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.

use std::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    mem,
    sync::Arc,
};

use base_db::{
    salsa::{self, ParallelDatabase},
    SourceDatabaseExt, SourceRootId, Upcast,
};
use fst::{self, Streamer};
use hir::{
    db::HirDatabase,
    symbols::{FileSymbol, SymbolCollector},
    Crate, Module,
};
use rayon::prelude::*;
use rustc_hash::FxHashSet;

use crate::RootDatabase;

#[derive(Debug)]
pub struct Query {
    query: String,
    lowercased: String,
    only_types: bool,
    libs: bool,
    exact: bool,
    case_sensitive: bool,
    limit: usize,
}

impl Query {
    pub fn new(query: String) -> Query {
        let lowercased = query.to_lowercase();
        Query {
            query,
            lowercased,
            only_types: false,
            libs: false,
            exact: false,
            case_sensitive: false,
            limit: usize::max_value(),
        }
    }

    pub fn only_types(&mut self) {
        self.only_types = true;
    }

    pub fn libs(&mut self) {
        self.libs = true;
    }

    pub fn exact(&mut self) {
        self.exact = true;
    }

    pub fn case_sensitive(&mut self) {
        self.case_sensitive = true;
    }

    pub fn limit(&mut self, limit: usize) {
        self.limit = limit
    }
}

#[salsa::query_group(SymbolsDatabaseStorage)]
pub trait SymbolsDatabase: HirDatabase + SourceDatabaseExt + Upcast<dyn HirDatabase> {
    /// The symbol index for a given module. These modules should only be in source roots that
    /// are inside local_roots.
    fn module_symbols(&self, module: Module) -> Arc<SymbolIndex>;

    /// The symbol index for a given source root within library_roots.
    fn library_symbols(&self, source_root_id: SourceRootId) -> Arc<SymbolIndex>;

    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<FxHashSet<SourceRootId>>;

    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<FxHashSet<SourceRootId>>;
}

fn library_symbols(db: &dyn SymbolsDatabase, source_root_id: SourceRootId) -> Arc<SymbolIndex> {
    let _p = profile::span("library_symbols");

    // todo: this could be parallelized, once I figure out how to do that...
    let symbols = db
        .source_root_crates(source_root_id)
        .iter()
        .flat_map(|&krate| Crate::from(krate).modules(db.upcast()))
        // we specifically avoid calling SymbolsDatabase::module_symbols here, even they do the same thing,
        // as the index for a library is not going to really ever change, and we do not want to store each
        // module's index in salsa.
        .flat_map(|module| SymbolCollector::collect(db.upcast(), module))
        .collect();

    Arc::new(SymbolIndex::new(symbols))
}

fn module_symbols(db: &dyn SymbolsDatabase, module: Module) -> Arc<SymbolIndex> {
    let _p = profile::span("module_symbols");
    let symbols = SymbolCollector::collect(db.upcast(), module);
    Arc::new(SymbolIndex::new(symbols))
}

/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
struct Snap<DB>(DB);
impl<DB: ParallelDatabase> Snap<salsa::Snapshot<DB>> {
    fn new(db: &DB) -> Self {
        Self(db.snapshot())
    }
}
impl<DB: ParallelDatabase> Clone for Snap<salsa::Snapshot<DB>> {
    fn clone(&self) -> Snap<salsa::Snapshot<DB>> {
        Snap(self.0.snapshot())
    }
}
impl<DB> std::ops::Deref for Snap<DB> {
    type Target = DB;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

// Feature: Workspace Symbol
//
// Uses fuzzy-search to find types, modules and functions by name across your
// project and dependencies. This is **the** most useful feature, which improves code
// navigation tremendously. It mostly works on top of the built-in LSP
// functionality, however `#` and `*` symbols can be used to narrow down the
// search. Specifically,
//
// - `Foo` searches for `Foo` type in the current workspace
// - `foo#` searches for `foo` function in the current workspace
// - `Foo*` searches for `Foo` type among dependencies, including `stdlib`
// - `foo#*` searches for `foo` function among dependencies
//
// That is, `#` switches from "types" to all symbols, `*` switches from the current
// workspace to dependencies.
//
// Note that filtering does not currently work in VSCode due to the editor never
// sending the special symbols to the language server. Instead, you can configure
// the filtering via the `rust-analyzer.workspace.symbol.search.scope` and
// `rust-analyzer.workspace.symbol.search.kind` settings.
//
// |===
// | Editor  | Shortcut
//
// | VS Code | kbd:[Ctrl+T]
// |===
pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    let _p = profile::span("world_symbols").detail(|| query.query.clone());

    let indices: Vec<_> = if query.libs {
        db.library_roots()
            .par_iter()
            .map_with(Snap::new(db), |snap, &root| snap.library_symbols(root))
            .collect()
    } else {
        let mut modules = Vec::new();

        for &root in db.local_roots().iter() {
            let crates = db.source_root_crates(root);
            for &krate in crates.iter() {
                modules.extend(Crate::from(krate).modules(db));
            }
        }

        modules
            .par_iter()
            .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
            .collect()
    };

    query.search(&indices)
}

pub fn crate_symbols(db: &RootDatabase, krate: Crate, query: Query) -> Vec<FileSymbol> {
    let _p = profile::span("crate_symbols").detail(|| format!("{:?}", query));

    let modules = krate.modules(db);
    let indices: Vec<_> = modules
        .par_iter()
        .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
        .collect();

    query.search(&indices)
}

pub fn index_resolve(db: &RootDatabase, name: &str) -> Vec<FileSymbol> {
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    world_symbols(db, query)
}

#[derive(Default)]
pub struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map<Vec<u8>>,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(lhs: &FileSymbol, rhs: &FileSymbol) -> Ordering {
            let lhs_chars = lhs.name.chars().map(|c| c.to_ascii_lowercase());
            let rhs_chars = rhs.name.chars().map(|c| c.to_ascii_lowercase());
            lhs_chars.cmp(rhs_chars)
        }

        symbols.par_sort_by(cmp);

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if let Some(next_symbol) = symbols.get(idx + 1) {
                if cmp(&symbols[last_batch_start], next_symbol) == Ordering::Equal {
                    continue;
                }
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_ascii_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::new(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    pub fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) | end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let _p = profile::span("symbol_index::Query::search");
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !symbol.kind.is_type() {
                        continue;
                    }
                    if self.exact {
                        if symbol.name != self.query {
                            continue;
                        }
                    } else if self.case_sensitive {
                        if self.query.chars().any(|c| !symbol.name.contains(c)) {
                            continue;
                        }
                    }

                    res.push(symbol.clone());
                    if res.len() >= self.limit {
                        return res;
                    }
                }
            }
        }
        res
    }
}

#[cfg(test)]
mod tests {

    use base_db::fixture::WithFixture;
    use expect_test::expect_file;
    use hir::symbols::SymbolCollector;

    use super::*;

    #[test]
    fn test_symbol_index_collection() {
        let (db, _) = RootDatabase::with_many_files(
            r#"
//- /main.rs

macro_rules! macro_rules_macro {
    () => {}
};

macro_rules! define_struct {
    () => {
        struct StructFromMacro;
    }
};

define_struct!();

macro Macro { }

struct Struct;
enum Enum {
    A, B
}
union Union {}

impl Struct {
    fn impl_fn() {}
}

trait Trait {
    fn trait_fn(&self);
}

fn main() {
    struct StructInFn;
}

const CONST: u32 = 1;
static STATIC: &'static str = "2";
type Alias = Struct;

mod a_mod {
    struct StructInModA;
}

const _: () = {
    struct StructInUnnamedConst;

    ()
};

const CONST_WITH_INNER: () = {
    struct StructInNamedConst;

    ()
};

mod b_mod;

//- /b_mod.rs
struct StructInModB;
        "#,
        );

        let symbols: Vec<_> = Crate::from(db.test_crate())
            .modules(&db)
            .into_iter()
            .map(|module_id| {
                let mut symbols = SymbolCollector::collect(&db, module_id);
                symbols.sort_by_key(|it| it.name.clone());
                (module_id, symbols)
            })
            .collect();

        expect_file!["./test_data/test_symbol_index_collection.txt"].assert_debug_eq(&symbols);
    }
}
<commit_msg>dead code<commit_after>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.

use std::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    mem,
    sync::Arc,
};

use base_db::{
    salsa::{self, ParallelDatabase},
    SourceDatabaseExt, SourceRootId, Upcast,
};
use fst::{self, Streamer};
use hir::{
    db::HirDatabase,
    symbols::{FileSymbol, SymbolCollector},
    Crate, Module,
};
use rayon::prelude::*;
use rustc_hash::FxHashSet;

use crate::RootDatabase;

#[derive(Debug)]
pub struct Query {
    query: String,
    lowercased: String,
    only_types: bool,
    libs: bool,
    exact: bool,
    case_sensitive: bool,
    limit: usize,
}

impl Query {
    pub fn new(query: String) -> Query {
        let lowercased = query.to_lowercase();
        Query {
            query,
            lowercased,
            only_types: false,
            libs: false,
            exact: false,
            case_sensitive: false,
            limit: usize::max_value(),
        }
    }

    pub fn only_types(&mut self) {
        self.only_types = true;
    }

    pub fn libs(&mut self) {
        self.libs = true;
    }

    pub fn exact(&mut self) {
        self.exact = true;
    }

    pub fn case_sensitive(&mut self) {
        self.case_sensitive = true;
    }

    pub fn limit(&mut self, limit: usize) {
        self.limit = limit
    }
}

#[salsa::query_group(SymbolsDatabaseStorage)]
pub trait SymbolsDatabase: HirDatabase + SourceDatabaseExt + Upcast<dyn HirDatabase> {
    /// The symbol index for a given module. These modules should only be in source roots that
    /// are inside local_roots.
    fn module_symbols(&self, module: Module) -> Arc<SymbolIndex>;

    /// The symbol index for a given source root within library_roots.
    fn library_symbols(&self, source_root_id: SourceRootId) -> Arc<SymbolIndex>;

    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<FxHashSet<SourceRootId>>;

    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<FxHashSet<SourceRootId>>;
}

fn library_symbols(db: &dyn SymbolsDatabase, source_root_id: SourceRootId) -> Arc<SymbolIndex> {
    let _p = profile::span("library_symbols");

    // todo: this could be parallelized, once I figure out how to do that...
    let symbols = db
        .source_root_crates(source_root_id)
        .iter()
        .flat_map(|&krate| Crate::from(krate).modules(db.upcast()))
        // we specifically avoid calling SymbolsDatabase::module_symbols here, even they do the same thing,
        // as the index for a library is not going to really ever change, and we do not want to store each
        // module's index in salsa.
        .flat_map(|module| SymbolCollector::collect(db.upcast(), module))
        .collect();

    Arc::new(SymbolIndex::new(symbols))
}

fn module_symbols(db: &dyn SymbolsDatabase, module: Module) -> Arc<SymbolIndex> {
    let _p = profile::span("module_symbols");
    let symbols = SymbolCollector::collect(db.upcast(), module);
    Arc::new(SymbolIndex::new(symbols))
}

/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
struct Snap<DB>(DB);
impl<DB: ParallelDatabase> Snap<salsa::Snapshot<DB>> {
    fn new(db: &DB) -> Self {
        Self(db.snapshot())
    }
}
impl<DB: ParallelDatabase> Clone for Snap<salsa::Snapshot<DB>> {
    fn clone(&self) -> Snap<salsa::Snapshot<DB>> {
        Snap(self.0.snapshot())
    }
}
impl<DB> std::ops::Deref for Snap<DB> {
    type Target = DB;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

// Feature: Workspace Symbol
//
// Uses fuzzy-search to find types, modules and functions by name across your
// project and dependencies. This is **the** most useful feature, which improves code
// navigation tremendously. It mostly works on top of the built-in LSP
// functionality, however `#` and `*` symbols can be used to narrow down the
// search. Specifically,
//
// - `Foo` searches for `Foo` type in the current workspace
// - `foo#` searches for `foo` function in the current workspace
// - `Foo*` searches for `Foo` type among dependencies, including `stdlib`
// - `foo#*` searches for `foo` function among dependencies
//
// That is, `#` switches from "types" to all symbols, `*` switches from the current
// workspace to dependencies.
//
// Note that filtering does not currently work in VSCode due to the editor never
// sending the special symbols to the language server. Instead, you can configure
// the filtering via the `rust-analyzer.workspace.symbol.search.scope` and
// `rust-analyzer.workspace.symbol.search.kind` settings.
//
// |===
// | Editor  | Shortcut
//
// | VS Code | kbd:[Ctrl+T]
// |===
pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    let _p = profile::span("world_symbols").detail(|| query.query.clone());

    let indices: Vec<_> = if query.libs {
        db.library_roots()
            .par_iter()
            .map_with(Snap::new(db), |snap, &root| snap.library_symbols(root))
            .collect()
    } else {
        let mut modules = Vec::new();

        for &root in db.local_roots().iter() {
            let crates = db.source_root_crates(root);
            for &krate in crates.iter() {
                modules.extend(Crate::from(krate).modules(db));
            }
        }

        modules
            .par_iter()
            .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
            .collect()
    };

    query.search(&indices)
}

pub fn crate_symbols(db: &RootDatabase, krate: Crate, query: Query) -> Vec<FileSymbol> {
    let _p = profile::span("crate_symbols").detail(|| format!("{:?}", query));

    let modules = krate.modules(db);
    let indices: Vec<_> = modules
        .par_iter()
        .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
        .collect();

    query.search(&indices)
}

#[derive(Default)]
pub struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map<Vec<u8>>,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(lhs: &FileSymbol, rhs: &FileSymbol) -> Ordering {
            let lhs_chars = lhs.name.chars().map(|c| c.to_ascii_lowercase());
            let rhs_chars = rhs.name.chars().map(|c| c.to_ascii_lowercase());
            lhs_chars.cmp(rhs_chars)
        }

        symbols.par_sort_by(cmp);

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if let Some(next_symbol) = symbols.get(idx + 1) {
                if cmp(&symbols[last_batch_start], next_symbol) == Ordering::Equal {
                    continue;
                }
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_ascii_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::new(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    pub fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) | end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let _p = profile::span("symbol_index::Query::search");
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !symbol.kind.is_type() {
                        continue;
                    }
                    if self.exact {
                        if symbol.name != self.query {
                            continue;
                        }
                    } else if self.case_sensitive {
                        if self.query.chars().any(|c| !symbol.name.contains(c)) {
                            continue;
                        }
                    }

                    res.push(symbol.clone());
                    if res.len() >= self.limit {
                        return res;
                    }
                }
            }
        }
        res
    }
}

#[cfg(test)]
mod tests {

    use base_db::fixture::WithFixture;
    use expect_test::expect_file;
    use hir::symbols::SymbolCollector;

    use super::*;

    #[test]
    fn test_symbol_index_collection() {
        let (db, _) = RootDatabase::with_many_files(
            r#"
//- /main.rs

macro_rules! macro_rules_macro {
    () => {}
};

macro_rules! define_struct {
    () => {
        struct StructFromMacro;
    }
};

define_struct!();

macro Macro { }

struct Struct;
enum Enum {
    A, B
}
union Union {}

impl Struct {
    fn impl_fn() {}
}

trait Trait {
    fn trait_fn(&self);
}

fn main() {
    struct StructInFn;
}

const CONST: u32 = 1;
static STATIC: &'static str = "2";
type Alias = Struct;

mod a_mod {
    struct StructInModA;
}

const _: () = {
    struct StructInUnnamedConst;

    ()
};

const CONST_WITH_INNER: () = {
    struct StructInNamedConst;

    ()
};

mod b_mod;

//- /b_mod.rs
struct StructInModB;
        "#,
        );

        let symbols: Vec<_> = Crate::from(db.test_crate())
            .modules(&db)
            .into_iter()
            .map(|module_id| {
                let mut symbols = SymbolCollector::collect(&db, module_id);
                symbols.sort_by_key(|it| it.name.clone());
                (module_id, symbols)
            })
            .collect();

        expect_file!["./test_data/test_symbol_index_collection.txt"].assert_debug_eq(&symbols);
    }
}
<|endoftext|>


<commit_before>#![feature(collections)]

fn main() {
    // (all the type annotations are superfluous)
    // A reference to a string allocated in read only memory
    let pangram: &'static str = "the quick brown fox jumps over the lazy dog";
    println!("Pangram: {}", pangram);

    // Iterate over words in reverse, no new string is allocated
    println!("Words in reverse");
    for word in pangram.split_whitespace().rev() {
        println!("> {}", word);
    }

    // Copy chars into a vector, sort and remove duplicates
    let mut chars: Vec<char> = pangram.chars().collect();
    chars.sort();
    chars.dedup();

    // Create an empty and growable `String`
    let mut string = String::new();
    for c in chars {
        // Insert a char at the end of string
        string.push(c);
        // Insert a string at the end of string
        string.push_str(", ");
    }

    // The trimmed string is a slice to the original string, hence no new
    // allocation is performed
    let chars_to_trim: &[char] = &[' ', ','];
    let trimmed_str: &str = string.trim_matches(chars_to_trim);
    println!("Used characters: {}", trimmed_str);

    // Heap allocate a string
    let alice = String::from_str("I like dogs");
    // Allocate new memory and store the modified string there
    let bob: String = alice.replace("dog", "cat");

    println!("Alice says: {}", alice);
    println!("Bob says: {}", bob);
}
<commit_msg>Fix warning: deprecated item<commit_after>fn main() {
    // (all the type annotations are superfluous)
    // A reference to a string allocated in read only memory
    let pangram: &'static str = "the quick brown fox jumps over the lazy dog";
    println!("Pangram: {}", pangram);

    // Iterate over words in reverse, no new string is allocated
    println!("Words in reverse");
    for word in pangram.split_whitespace().rev() {
        println!("> {}", word);
    }

    // Copy chars into a vector, sort and remove duplicates
    let mut chars: Vec<char> = pangram.chars().collect();
    chars.sort();
    chars.dedup();

    // Create an empty and growable `String`
    let mut string = String::new();
    for c in chars {
        // Insert a char at the end of string
        string.push(c);
        // Insert a string at the end of string
        string.push_str(", ");
    }

    // The trimmed string is a slice to the original string, hence no new
    // allocation is performed
    let chars_to_trim: &[char] = &[' ', ','];
    let trimmed_str: &str = string.trim_matches(chars_to_trim);
    println!("Used characters: {}", trimmed_str);

    // Heap allocate a string
    let alice = String::from("I like dogs");
    // Allocate new memory and store the modified string there
    let bob: String = alice.replace("dog", "cat");

    println!("Alice says: {}", alice);
    println!("Bob says: {}", bob);
}
<|endoftext|>


<commit_before>use crate::grammar::ast::{BinaryExpression, BinaryOp, Expression};
use crate::grammar::model::Fragment;

/// Shunting yard algorithm for structuring binary expressions. Takes
pub fn shunting_yard<'s>(
    expressions: Vec<Expression<'s>>,
    ops: Vec<BinaryOp>,
    frag: Fragment<'s>,
) -> BinaryExpression<'s> {
    unimplemented!()
}
<commit_msg>Add binary operator parser<commit_after>use crate::grammar::ast::{BinaryExpression, BinaryOp, Expression};
use crate::grammar::model::Fragment;
use crate::grammar::parsers::expression::binary_expression::operators::OperatorInfo;
use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::combinator::{map, value};
use nom::IResult;

fn binary_operator<'s>(input: Fragment<'s>) -> IResult<Fragment<'s>, OperatorInfo> {
    use BinaryOp::*;
    map(
        alt((
            value(OrOr, tag("||")),
            value(AndAnd, tag("&&")),
            value(Or, tag("|")),
            value(Xor, tag("^")),
            value(And, tag("&")),
            value(EqEq, tag("==")),
            value(NotEq, tag("!=")),
            value(Le, tag("<=")),
            value(Ge, tag(">=")),
            value(Lt, tag("<")),
            value(Gt, tag(">")),
            value(DotDot, tag("..")),
            value(Add, tag("+")),
            value(Sub, tag("-")),
            value(Mul, tag("*")),
            value(Div, tag("/")),
            value(Mod, tag("%")),
        )),
        BinaryOp::get_info,
    )(input)
}

/// Shunting yard algorithm for structuring binary expressions. Takes
pub fn shunting_yard<'s>(
    expressions: Vec<Expression<'s>>,
    ops: Vec<BinaryOp>,
    frag: Fragment<'s>,
) -> BinaryExpression<'s> {
    unimplemented!()
}
<|endoftext|>


<commit_before>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! TODO: Some LSP clients support watching the disk, so this crate should to
//! support custom watcher events (related to
//! <https://github.com/rust-analyzer/rust-analyzer/issues/131>)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! which are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod roots;
mod io;

use std::{
    fmt, fs, mem,
    path::{Path, PathBuf},
    sync::Arc,
};

use crossbeam_channel::Receiver;
use ra_arena::{impl_arena_id, Arena, RawId, map::ArenaMap};
use relative_path::{RelativePath, RelativePathBuf};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{
    io::{TaskResult, Worker},
    roots::Roots,
};

pub use crate::{
    io::TaskResult as VfsTask,
    roots::VfsRoot,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VfsFile(pub RawId);
impl_arena_id!(VfsFile);

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    is_overlayed: bool,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arc<Roots>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: ArenaMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: Worker,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Vfs")
            .field("n_roots", &self.roots.len())
            .field("n_files", &self.files.len())
            .field("n_pending_changes", &self.pending_changes.len())
            .finish()
    }
}

impl Vfs {
    pub fn new(roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let roots = Arc::new(Roots::new(roots));
        let worker = io::start(Arc::clone(&roots));
        let mut root2files = ArenaMap::default();

        for root in roots.iter() {
            root2files.insert(root, Default::default());
            worker.sender().send(io::Task::AddRoot { root }).unwrap();
        }
        let res =
            Vfs { roots, files: Arena::default(), root2files, worker, pending_changes: Vec::new() };
        let vfs_roots = res.roots.iter().collect();
        (res, vfs_roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots.path(root).to_path_buf()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots.path(self.files[file].root);
        rel_path.to_path(root_path)
    }

    pub fn n_roots(&self) -> usize {
        self.roots.len()
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text), false);
                let change = VfsChange::AddFile { file, text, root, path: rel_path };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        if let Some(file) = file {
            self.do_change_file(file, text, true);
            Some(file)
        } else {
            self.do_add_file(root, rel_path, text, true)
        }
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            self.do_change_file(file, new_text, true);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        let file = file.expect("can't remove a file which wasn't added");
        let full_path = rel_path.to_path(&self.roots.path(root));
        if let Ok(text) = fs::read_to_string(&full_path) {
            self.do_change_file(file, text, true);
        } else {
            self.do_remove_file(root, rel_path, file, true);
        }
        Some(file)
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        self.worker.receiver()
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        match task {
            TaskResult::BulkLoadRoot { root, files } => {
                let mut cur_files = Vec::new();
                // While we were scanning the root in the background, a file might have
                // been open in the editor, so we need to account for that.
                let existing = self.root2files[root]
                    .iter()
                    .map(|&file| (self.files[file].path.clone(), file))
                    .collect::<FxHashMap<_, _>>();
                for (path, text) in files {
                    if let Some(&file) = existing.get(&path) {
                        let text = Arc::clone(&self.files[file].text);
                        cur_files.push((file, path, text));
                        continue;
                    }
                    let text = Arc::new(text);
                    let file = self.add_file(root, path.clone(), Arc::clone(&text), false);
                    cur_files.push((file, path, text));
                }

                let change = VfsChange::AddRoot { root, files: cur_files };
                self.pending_changes.push(change);
            }
            TaskResult::SingleFile { root, path, text } => {
                match (self.find_file(root, &path), text) {
                    (Some(file), None) => {
                        self.do_remove_file(root, path, file, false);
                    }
                    (None, Some(text)) => {
                        self.do_add_file(root, path, text, false);
                    }
                    (Some(file), Some(text)) => {
                        self.do_change_file(file, text, false);
                    }
                    (None, None) => (),
                }
            }
        }
    }

    fn do_add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
        is_overlay: bool,
    ) -> Option<VfsFile> {
        let text = Arc::new(text);
        let file = self.add_file(root, path.clone(), text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::AddFile { file, root, path, text });
        Some(file)
    }

    fn do_change_file(&mut self, file: VfsFile, text: String, is_overlay: bool) {
        if !is_overlay && self.files[file].is_overlayed {
            return;
        }
        let text = Arc::new(text);
        self.change_file(file, text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::ChangeFile { file, text });
    }

    fn do_remove_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        file: VfsFile,
        is_overlay: bool,
    ) {
        if !is_overlay && self.files[file].is_overlayed {
            return;
        }
        self.remove_file(file);
        self.pending_changes.push(VfsChange::RemoveFile { root, path, file });
    }

    fn add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: Arc<String>,
        is_overlayed: bool,
    ) -> VfsFile {
        let data = VfsFileData { root, path, text, is_overlayed };
        let file = self.files.alloc(data);
        self.root2files.get_mut(root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>, is_overlayed: bool) {
        let mut file_data = &mut self.files[file];
        file_data.text = new_text;
        file_data.is_overlayed = is_overlayed;
    }

    fn remove_file(&mut self, file: VfsFile) {
        // FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self.roots.find(&path)?;
        let file = self.find_file(root, &path);
        Some((root, path, file))
    }

    fn find_file(&self, root: VfsRoot, path: &RelativePath) -> Option<VfsFile> {
        self.root2files[root].iter().map(|&it| it).find(|&file| self.files[file].path == path)
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot { root: VfsRoot, files: Vec<(VfsFile, RelativePathBuf, Arc<String>)> },
    AddFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf, text: Arc<String> },
    RemoveFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf },
    ChangeFile { file: VfsFile, text: Arc<String> },
}
<commit_msg>simplify overlay handling<commit_after>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! TODO: Some LSP clients support watching the disk, so this crate should to
//! support custom watcher events (related to
//! <https://github.com/rust-analyzer/rust-analyzer/issues/131>)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! which are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod roots;
mod io;

use std::{
    fmt, fs, mem,
    path::{Path, PathBuf},
    sync::Arc,
};

use crossbeam_channel::Receiver;
use ra_arena::{impl_arena_id, Arena, RawId, map::ArenaMap};
use relative_path::{RelativePath, RelativePathBuf};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{
    io::{TaskResult, Worker},
    roots::Roots,
};

pub use crate::{
    io::TaskResult as VfsTask,
    roots::VfsRoot,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VfsFile(pub RawId);
impl_arena_id!(VfsFile);

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    is_overlayed: bool,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arc<Roots>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: ArenaMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: Worker,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Vfs")
            .field("n_roots", &self.roots.len())
            .field("n_files", &self.files.len())
            .field("n_pending_changes", &self.pending_changes.len())
            .finish()
    }
}

impl Vfs {
    pub fn new(roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let roots = Arc::new(Roots::new(roots));
        let worker = io::start(Arc::clone(&roots));
        let mut root2files = ArenaMap::default();

        for root in roots.iter() {
            root2files.insert(root, Default::default());
            worker.sender().send(io::Task::AddRoot { root }).unwrap();
        }
        let res =
            Vfs { roots, files: Arena::default(), root2files, worker, pending_changes: Vec::new() };
        let vfs_roots = res.roots.iter().collect();
        (res, vfs_roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots.path(root).to_path_buf()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots.path(self.files[file].root);
        rel_path.to_path(root_path)
    }

    pub fn n_roots(&self) -> usize {
        self.roots.len()
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text), false);
                let change = VfsChange::AddFile { file, text, root, path: rel_path };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        if let Some(file) = file {
            self.do_change_file(file, text, true);
            Some(file)
        } else {
            self.do_add_file(root, rel_path, text, true)
        }
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            self.do_change_file(file, new_text, true);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        let file = file.expect("can't remove a file which wasn't added");
        let full_path = rel_path.to_path(&self.roots.path(root));
        if let Ok(text) = fs::read_to_string(&full_path) {
            self.do_change_file(file, text, true);
        } else {
            self.do_remove_file(root, rel_path, file);
        }
        Some(file)
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        self.worker.receiver()
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        match task {
            TaskResult::BulkLoadRoot { root, files } => {
                let mut cur_files = Vec::new();
                // While we were scanning the root in the background, a file might have
                // been open in the editor, so we need to account for that.
                let existing = self.root2files[root]
                    .iter()
                    .map(|&file| (self.files[file].path.clone(), file))
                    .collect::<FxHashMap<_, _>>();
                for (path, text) in files {
                    if let Some(&file) = existing.get(&path) {
                        let text = Arc::clone(&self.files[file].text);
                        cur_files.push((file, path, text));
                        continue;
                    }
                    let text = Arc::new(text);
                    let file = self.add_file(root, path.clone(), Arc::clone(&text), false);
                    cur_files.push((file, path, text));
                }

                let change = VfsChange::AddRoot { root, files: cur_files };
                self.pending_changes.push(change);
            }
            TaskResult::SingleFile { root, path, text } => {
                let file = self.find_file(root, &path);
                if file.map(|file| self.files[file].is_overlayed) == Some(true) {
                    return;
                }
                match (file, text) {
                    (Some(file), None) => {
                        self.do_remove_file(root, path, file);
                    }
                    (None, Some(text)) => {
                        self.do_add_file(root, path, text, false);
                    }
                    (Some(file), Some(text)) => {
                        self.do_change_file(file, text, false);
                    }
                    (None, None) => (),
                }
            }
        }
    }

    fn do_add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
        is_overlay: bool,
    ) -> Option<VfsFile> {
        let text = Arc::new(text);
        let file = self.add_file(root, path.clone(), text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::AddFile { file, root, path, text });
        Some(file)
    }

    fn do_change_file(&mut self, file: VfsFile, text: String, is_overlay: bool) {
        let text = Arc::new(text);
        self.change_file(file, text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::ChangeFile { file, text });
    }

    fn do_remove_file(&mut self, root: VfsRoot, path: RelativePathBuf, file: VfsFile) {
        self.remove_file(file);
        self.pending_changes.push(VfsChange::RemoveFile { root, path, file });
    }

    fn add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: Arc<String>,
        is_overlayed: bool,
    ) -> VfsFile {
        let data = VfsFileData { root, path, text, is_overlayed };
        let file = self.files.alloc(data);
        self.root2files.get_mut(root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>, is_overlayed: bool) {
        let mut file_data = &mut self.files[file];
        file_data.text = new_text;
        file_data.is_overlayed = is_overlayed;
    }

    fn remove_file(&mut self, file: VfsFile) {
        // FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self.roots.find(&path)?;
        let file = self.find_file(root, &path);
        Some((root, path, file))
    }

    fn find_file(&self, root: VfsRoot, path: &RelativePath) -> Option<VfsFile> {
        self.root2files[root].iter().map(|&it| it).find(|&file| self.files[file].path == path)
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot { root: VfsRoot, files: Vec<(VfsFile, RelativePathBuf, Arc<String>)> },
    AddFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf, text: Arc<String> },
    RemoveFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf },
    ChangeFile { file: VfsFile, text: Arc<String> },
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Processes that are directly spawned by shared mempool runtime initialization

use crate::{
    core_mempool::{CoreMempool, TimelineState},
    counters,
    logging::{LogEntry, LogEvent, LogSchema},
    network::{MempoolNetworkEvents, MempoolSyncMsg},
    shared_mempool::{
        tasks,
        types::{notify_subscribers, SharedMempool, SharedMempoolNotification},
    },
    CommitNotification, ConsensusRequest, SubmissionStatus,
};
use ::network::protocols::network::Event;
use anyhow::Result;
use bounded_executor::BoundedExecutor;
use channel::diem_channel;
use diem_config::{config::PeerNetworkId, network_id::NodeNetworkId};
use diem_infallible::Mutex;
use diem_logger::prelude::*;
use diem_trace::prelude::*;
use diem_types::{on_chain_config::OnChainConfigPayload, transaction::SignedTransaction};
use futures::{
    channel::{mpsc, oneshot},
    stream::{select_all, FuturesUnordered},
    StreamExt,
};
use std::{
    ops::Deref,
    sync::Arc,
    time::{Duration, SystemTime},
};
use tokio::{runtime::Handle, time::interval};
use tokio_stream::wrappers::IntervalStream;
use vm_validator::vm_validator::TransactionValidation;

/// Coordinator that handles inbound network events and outbound txn broadcasts.
pub(crate) async fn coordinator<V>(
    mut smp: SharedMempool<V>,
    executor: Handle,
    network_events: Vec<(NodeNetworkId, MempoolNetworkEvents)>,
    mut client_events: mpsc::Receiver<(
        SignedTransaction,
        oneshot::Sender<Result<SubmissionStatus>>,
    )>,
    mut consensus_requests: mpsc::Receiver<ConsensusRequest>,
    mut state_sync_requests: mpsc::Receiver<CommitNotification>,
    mut mempool_reconfig_events: diem_channel::Receiver<(), OnChainConfigPayload>,
) where
    V: TransactionValidation,
{
    info!(LogSchema::event_log(
        LogEntry::CoordinatorRuntime,
        LogEvent::Start
    ));
    let smp_events: Vec<_> = network_events
        .into_iter()
        .map(|(network_id, events)| events.map(move |e| (network_id.clone(), e)))
        .collect();
    let mut events = select_all(smp_events).fuse();
    let mempool = smp.mempool.clone();
    let peer_manager = smp.peer_manager.clone();
    let subscribers = &mut smp.subscribers.clone();
    let mut scheduled_broadcasts = FuturesUnordered::new();

    // Use a BoundedExecutor to restrict only `workers_available` concurrent
    // worker tasks that can process incoming transactions.
    let workers_available = smp.config.shared_mempool_max_concurrent_inbound_syncs;
    let bounded_executor = BoundedExecutor::new(workers_available, executor.clone());

    loop {
        let _timer = counters::MAIN_LOOP.start_timer();
        ::futures::select! {
            (msg, callback) = client_events.select_next_some() => {
                trace_event!("mempool::client_event", {"txn", msg.sender(), msg.sequence_number()});
                // This timer measures how long it took for the bounded executor to *schedule* the
                // task.
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::CLIENT_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                // This timer measures how long it took for the task to go from scheduled to started.
                let task_start_timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::CLIENT_EVENT_LABEL, counters::START_LABEL])
                    .start_timer();
                bounded_executor
                .spawn(tasks::process_client_transaction_submission(
                    smp.clone(),
                    msg,
                    callback,
                    task_start_timer,
                ))
                .await;
            },
            msg = consensus_requests.select_next_some() => {
                tasks::process_consensus_request(&mempool, msg).await;
            }
            msg = state_sync_requests.select_next_some() => {
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::STATE_SYNC_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                tokio::spawn(tasks::process_state_sync_request(mempool.clone(), msg));
            }
            config_update = mempool_reconfig_events.select_next_some() => {
                info!(LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::Received));
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::RECONFIG_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                bounded_executor
                    .spawn(tasks::process_config_update(config_update, smp.validator.clone()))
                    .await;
            },
            (peer, backoff) = scheduled_broadcasts.select_next_some() => {
                tasks::execute_broadcast(peer, backoff, &mut smp, &mut scheduled_broadcasts, executor.clone());
            },
            (network_id, event) = events.select_next_some() => {
                match event {
                    Event::NewPeer(metadata) => {
                        counters::SHARED_MEMPOOL_EVENTS
                            .with_label_values(&["new_peer".to_string().deref()])
                            .inc();
                        let origin = metadata.origin;
                        let peer = PeerNetworkId(network_id, metadata.remote_peer_id);
                        let is_new_peer = peer_manager.add_peer(peer.clone(), origin);
                        let is_upstream_peer = peer_manager.is_upstream_peer(&peer, Some(origin));
                        debug!(LogSchema::new(LogEntry::NewPeer).peer(&peer).is_upstream_peer(is_upstream_peer));
                        notify_subscribers(SharedMempoolNotification::PeerStateChange, &subscribers);
                        if is_new_peer && is_upstream_peer {
                            tasks::execute_broadcast(peer, false, &mut smp, &mut scheduled_broadcasts, executor.clone());
                        }
                    }
                    Event::LostPeer(metadata) => {
                        counters::SHARED_MEMPOOL_EVENTS
                            .with_label_values(&["lost_peer".to_string().deref()])
                            .inc();
                        let peer = PeerNetworkId(network_id, metadata.remote_peer_id);
                        debug!(LogSchema::new(LogEntry::LostPeer)
                            .peer(&peer)
                            .is_upstream_peer(peer_manager.is_upstream_peer(&peer, Some(metadata.origin)))
                        );
                        peer_manager.disable_peer(peer);
                        notify_subscribers(SharedMempoolNotification::PeerStateChange, &subscribers);
                    }
                    Event::Message(peer_id, msg) => {
                        counters::SHARED_MEMPOOL_EVENTS
                            .with_label_values(&["message".to_string().deref()])
                            .inc();
                        match msg {
                            MempoolSyncMsg::BroadcastTransactionsRequest{request_id, transactions} => {
                                let smp_clone = smp.clone();
                                let peer = PeerNetworkId(network_id, peer_id);
                                let timeline_state = match peer_manager
                                    .is_upstream_peer(&peer, None)
                                {
                                    true => TimelineState::NonQualified,
                                    false => TimelineState::NotReady,
                                };
                                // This timer measures how long it took for the bounded executor to
                                // *schedule* the task.
                                let _timer = counters::TASK_SPAWN_LATENCY
                                    .with_label_values(&[counters::PEER_BROADCAST_EVENT_LABEL, counters::SPAWN_LABEL])
                                    .start_timer();
                                // This timer measures how long it took for the task to go from scheduled
                                // to started.
                                let task_start_timer = counters::TASK_SPAWN_LATENCY
                                    .with_label_values(&[counters::PEER_BROADCAST_EVENT_LABEL, counters::START_LABEL])
                                    .start_timer();
                                bounded_executor
                                    .spawn(tasks::process_transaction_broadcast(
                                        smp_clone,
                                        transactions,
                                        request_id,
                                        timeline_state,
                                        peer,
                                        task_start_timer
                                    ))
                                    .await;
                            }
                            MempoolSyncMsg::BroadcastTransactionsResponse{request_id, retry, backoff} => {
                                let ack_timestamp = SystemTime::now();
                                peer_manager.process_broadcast_ack(PeerNetworkId(network_id.clone(), peer_id), request_id, retry, backoff, ack_timestamp);
                            }
                        }
                    }
                    Event::RpcRequest(peer_id, _msg, _res_tx) => {
                        counters::UNEXPECTED_NETWORK_MSG_COUNT
                            .with_label_values(&[&network_id.network_id().to_string(), &peer_id.to_string()])
                            .inc();
                        sample!(
                            SampleRate::Duration(Duration::from_secs(60)),
                            warn!(
                                LogSchema::new(LogEntry::UnexpectedNetworkMsg)
                                    .peer(&PeerNetworkId(network_id, peer_id))
                            )
                        );
                    }
                }
            },
            complete => break,
        }
    }
    error!(LogSchema::event_log(
        LogEntry::CoordinatorRuntime,
        LogEvent::Terminated
    ));
}

/// Garbage collect all expired transactions by SystemTTL.
pub(crate) async fn gc_coordinator(mempool: Arc<Mutex<CoreMempool>>, gc_interval_ms: u64) {
    info!(LogSchema::event_log(LogEntry::GCRuntime, LogEvent::Start));
    let mut interval = IntervalStream::new(interval(Duration::from_millis(gc_interval_ms)));
    while let Some(_interval) = interval.next().await {
        sample!(
            SampleRate::Duration(Duration::from_secs(60)),
            info!(LogSchema::event_log(LogEntry::GCRuntime, LogEvent::Live))
        );
        mempool.lock().gc();
    }

    error!(LogSchema::event_log(
        LogEntry::GCRuntime,
        LogEvent::Terminated
    ));
}

/// Periodically logs a snapshot of transactions in core mempool.
/// In the future we may want an interactive way to directly query mempool's internal state.
/// For now, we will rely on this periodic snapshot to observe the internal state.
pub(crate) async fn snapshot_job(mempool: Arc<Mutex<CoreMempool>>, snapshot_interval_secs: u64) {
    let mut interval = IntervalStream::new(interval(Duration::from_secs(snapshot_interval_secs)));
    while let Some(_interval) = interval.next().await {
        let snapshot = mempool.lock().gen_snapshot();
        debug!(LogSchema::new(LogEntry::MempoolSnapshot).txns(snapshot));
    }
}
<commit_msg>[mempool] Move network message handling into separate function<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Processes that are directly spawned by shared mempool runtime initialization

use crate::{
    core_mempool::{CoreMempool, TimelineState},
    counters,
    logging::{LogEntry, LogEvent, LogSchema},
    network::{MempoolNetworkEvents, MempoolSyncMsg},
    shared_mempool::{
        tasks,
        types::{notify_subscribers, ScheduledBroadcast, SharedMempool, SharedMempoolNotification},
    },
    CommitNotification, ConsensusRequest, SubmissionStatus,
};
use ::network::protocols::network::Event;
use anyhow::Result;
use bounded_executor::BoundedExecutor;
use channel::diem_channel;
use diem_config::{config::PeerNetworkId, network_id::NodeNetworkId};
use diem_infallible::Mutex;
use diem_logger::prelude::*;
use diem_trace::prelude::*;
use diem_types::{on_chain_config::OnChainConfigPayload, transaction::SignedTransaction};
use futures::{
    channel::{mpsc, oneshot},
    stream::{select_all, FuturesUnordered},
    StreamExt,
};
use std::{
    ops::Deref,
    sync::Arc,
    time::{Duration, SystemTime},
};
use tokio::{runtime::Handle, time::interval};
use tokio_stream::wrappers::IntervalStream;
use vm_validator::vm_validator::TransactionValidation;

/// Coordinator that handles inbound network events and outbound txn broadcasts.
pub(crate) async fn coordinator<V>(
    mut smp: SharedMempool<V>,
    executor: Handle,
    network_events: Vec<(NodeNetworkId, MempoolNetworkEvents)>,
    mut client_events: mpsc::Receiver<(
        SignedTransaction,
        oneshot::Sender<Result<SubmissionStatus>>,
    )>,
    mut consensus_requests: mpsc::Receiver<ConsensusRequest>,
    mut state_sync_requests: mpsc::Receiver<CommitNotification>,
    mut mempool_reconfig_events: diem_channel::Receiver<(), OnChainConfigPayload>,
) where
    V: TransactionValidation,
{
    info!(LogSchema::event_log(
        LogEntry::CoordinatorRuntime,
        LogEvent::Start
    ));
    let smp_events: Vec<_> = network_events
        .into_iter()
        .map(|(network_id, events)| events.map(move |e| (network_id.clone(), e)))
        .collect();
    let mut events = select_all(smp_events).fuse();
    let mempool = smp.mempool.clone();
    let mut scheduled_broadcasts = FuturesUnordered::new();

    // Use a BoundedExecutor to restrict only `workers_available` concurrent
    // worker tasks that can process incoming transactions.
    let workers_available = smp.config.shared_mempool_max_concurrent_inbound_syncs;
    let bounded_executor = BoundedExecutor::new(workers_available, executor.clone());

    loop {
        let _timer = counters::MAIN_LOOP.start_timer();
        ::futures::select! {
            (msg, callback) = client_events.select_next_some() => {
                trace_event!("mempool::client_event", {"txn", msg.sender(), msg.sequence_number()});
                // This timer measures how long it took for the bounded executor to *schedule* the
                // task.
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::CLIENT_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                // This timer measures how long it took for the task to go from scheduled to started.
                let task_start_timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::CLIENT_EVENT_LABEL, counters::START_LABEL])
                    .start_timer();
                bounded_executor
                .spawn(tasks::process_client_transaction_submission(
                    smp.clone(),
                    msg,
                    callback,
                    task_start_timer,
                ))
                .await;
            },
            msg = consensus_requests.select_next_some() => {
                tasks::process_consensus_request(&mempool, msg).await;
            }
            msg = state_sync_requests.select_next_some() => {
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::STATE_SYNC_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                tokio::spawn(tasks::process_state_sync_request(mempool.clone(), msg));
            }
            config_update = mempool_reconfig_events.select_next_some() => {
                info!(LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::Received));
                let _timer = counters::TASK_SPAWN_LATENCY
                    .with_label_values(&[counters::RECONFIG_EVENT_LABEL, counters::SPAWN_LABEL])
                    .start_timer();
                bounded_executor
                    .spawn(tasks::process_config_update(config_update, smp.validator.clone()))
                    .await;
            },
            (peer, backoff) = scheduled_broadcasts.select_next_some() => {
                tasks::execute_broadcast(peer, backoff, &mut smp, &mut scheduled_broadcasts, executor.clone());
            },
            (network_id, event) = events.select_next_some() => {
                handle_event(&executor, &bounded_executor, &mut scheduled_broadcasts, &mut smp, network_id, event).await;
            },
            complete => break,
        }
    }
    error!(LogSchema::event_log(
        LogEntry::CoordinatorRuntime,
        LogEvent::Terminated
    ));
}

async fn handle_event<V>(
    executor: &Handle,
    bounded_executor: &BoundedExecutor,
    scheduled_broadcasts: &mut FuturesUnordered<ScheduledBroadcast>,
    smp: &mut SharedMempool<V>,
    network_id: NodeNetworkId,
    event: Event<MempoolSyncMsg>,
) where
    V: TransactionValidation,
{
    match event {
        Event::NewPeer(metadata) => {
            counters::SHARED_MEMPOOL_EVENTS
                .with_label_values(&["new_peer".to_string().deref()])
                .inc();
            let origin = metadata.origin;
            let peer = PeerNetworkId(network_id, metadata.remote_peer_id);
            let is_new_peer = smp.peer_manager.add_peer(peer.clone(), origin);
            let is_upstream_peer = smp.peer_manager.is_upstream_peer(&peer, Some(origin));
            debug!(LogSchema::new(LogEntry::NewPeer)
                .peer(&peer)
                .is_upstream_peer(is_upstream_peer));
            notify_subscribers(SharedMempoolNotification::PeerStateChange, &smp.subscribers);
            if is_new_peer && is_upstream_peer {
                tasks::execute_broadcast(peer, false, smp, scheduled_broadcasts, executor.clone());
            }
        }
        Event::LostPeer(metadata) => {
            counters::SHARED_MEMPOOL_EVENTS
                .with_label_values(&["lost_peer".to_string().deref()])
                .inc();
            let peer = PeerNetworkId(network_id, metadata.remote_peer_id);
            debug!(LogSchema::new(LogEntry::LostPeer)
                .peer(&peer)
                .is_upstream_peer(
                    smp.peer_manager
                        .is_upstream_peer(&peer, Some(metadata.origin))
                ));
            smp.peer_manager.disable_peer(peer);
            notify_subscribers(SharedMempoolNotification::PeerStateChange, &smp.subscribers);
        }
        Event::Message(peer_id, msg) => {
            counters::SHARED_MEMPOOL_EVENTS
                .with_label_values(&["message".to_string().deref()])
                .inc();
            match msg {
                MempoolSyncMsg::BroadcastTransactionsRequest {
                    request_id,
                    transactions,
                } => {
                    let smp_clone = smp.clone();
                    let peer = PeerNetworkId(network_id, peer_id);
                    let timeline_state = match smp.peer_manager.is_upstream_peer(&peer, None) {
                        true => TimelineState::NonQualified,
                        false => TimelineState::NotReady,
                    };
                    // This timer measures how long it took for the bounded executor to
                    // *schedule* the task.
                    let _timer = counters::TASK_SPAWN_LATENCY
                        .with_label_values(&[
                            counters::PEER_BROADCAST_EVENT_LABEL,
                            counters::SPAWN_LABEL,
                        ])
                        .start_timer();
                    // This timer measures how long it took for the task to go from scheduled
                    // to started.
                    let task_start_timer = counters::TASK_SPAWN_LATENCY
                        .with_label_values(&[
                            counters::PEER_BROADCAST_EVENT_LABEL,
                            counters::START_LABEL,
                        ])
                        .start_timer();
                    bounded_executor
                        .spawn(tasks::process_transaction_broadcast(
                            smp_clone,
                            transactions,
                            request_id,
                            timeline_state,
                            peer,
                            task_start_timer,
                        ))
                        .await;
                }
                MempoolSyncMsg::BroadcastTransactionsResponse {
                    request_id,
                    retry,
                    backoff,
                } => {
                    let ack_timestamp = SystemTime::now();
                    smp.peer_manager.process_broadcast_ack(
                        PeerNetworkId(network_id, peer_id),
                        request_id,
                        retry,
                        backoff,
                        ack_timestamp,
                    );
                }
            }
        }
        Event::RpcRequest(peer_id, _msg, _res_tx) => {
            counters::UNEXPECTED_NETWORK_MSG_COUNT
                .with_label_values(&[&network_id.network_id().to_string(), &peer_id.to_string()])
                .inc();
            sample!(
                SampleRate::Duration(Duration::from_secs(60)),
                warn!(LogSchema::new(LogEntry::UnexpectedNetworkMsg)
                    .peer(&PeerNetworkId(network_id, peer_id)))
            );
        }
    }
}

/// Garbage collect all expired transactions by SystemTTL.
pub(crate) async fn gc_coordinator(mempool: Arc<Mutex<CoreMempool>>, gc_interval_ms: u64) {
    info!(LogSchema::event_log(LogEntry::GCRuntime, LogEvent::Start));
    let mut interval = IntervalStream::new(interval(Duration::from_millis(gc_interval_ms)));
    while let Some(_interval) = interval.next().await {
        sample!(
            SampleRate::Duration(Duration::from_secs(60)),
            info!(LogSchema::event_log(LogEntry::GCRuntime, LogEvent::Live))
        );
        mempool.lock().gc();
    }

    error!(LogSchema::event_log(
        LogEntry::GCRuntime,
        LogEvent::Terminated
    ));
}

/// Periodically logs a snapshot of transactions in core mempool.
/// In the future we may want an interactive way to directly query mempool's internal state.
/// For now, we will rely on this periodic snapshot to observe the internal state.
pub(crate) async fn snapshot_job(mempool: Arc<Mutex<CoreMempool>>, snapshot_interval_secs: u64) {
    let mut interval = IntervalStream::new(interval(Duration::from_secs(snapshot_interval_secs)));
    while let Some(_interval) = interval.next().await {
        let snapshot = mempool.lock().gen_snapshot();
        debug!(LogSchema::new(LogEntry::MempoolSnapshot).txns(snapshot));
    }
}
<|endoftext|>