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

use crate::error::{Error, Result};
use serde::{ser, Serialize};
use std::collections::BTreeMap;

/// Serialize the given data structure as a `Vec<u8>` of LCS.
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to
/// fail, if `T` contains sequences which are longer than `MAX_SEQUENCE_LENGTH`,
/// or if `T` attempts to serialize an unsupported datatype such as a f32,
/// f64, or char.
///
/// # Examples
///
/// ```
/// use libra_canonical_serialization::to_bytes;
/// use serde::Serialize;
///
/// #[derive(Serialize)]
/// struct Ip([u8; 4]);
///
/// #[derive(Serialize)]
/// struct Port(u16);
///
/// #[derive(Serialize)]
/// struct Service {
///     ip: Ip,
///     port: Vec<Port>,
///     connection_max: Option<u32>,
///     enabled: bool,
/// }
///
/// let service = Service {
///     ip: Ip([192, 168, 1, 1]),
///     port: vec![Port(8001), Port(8002), Port(8003)],
///     connection_max: Some(5000),
///     enabled: false,
/// };
///
/// let bytes = to_bytes(&service).unwrap();
/// let expected = vec![
///     0xc0, 0xa8, 0x01, 0x01, 0x03, 0x41, 0x1f, 0x42,
///     0x1f, 0x43, 0x1f, 0x01, 0x88, 0x13, 0x00, 0x00,
///     0x00,
/// ];
/// assert_eq!(bytes, expected);
/// ```
pub fn to_bytes<T>(value: &T) -> Result<Vec<u8>>
where
    T: ?Sized + Serialize,
{
    let mut serializer = Serializer::new();
    value.serialize(&mut serializer)?;
    Ok(serializer.end())
}

pub fn is_human_readable() -> bool {
    let mut serializer = Serializer::new();
    ser::Serializer::is_human_readable(&&mut serializer)
}

/// Serialization implementation for LCS
struct Serializer {
    output: Vec<u8>,
}

impl Serializer {
    /// Creates a new `Serializer` which will emit LCS.
    fn new() -> Self {
        Self { output: Vec::new() }
    }

    /// The `Serializer::end` method should be called after a type has been
    /// fully serialized. This will return the buffer containing the LCS
    /// serialized data.
    fn end(self) -> Vec<u8> {
        self.output
    }

    fn serialize_u32_as_uleb128(&mut self, mut value: u32) -> Result<()> {
        use serde::ser::Serializer;
        while value >= 0x80 {
            // Write 7 (lowest) bits of data and set the 8th bit to 1.
            let byte = (value & 0x7f) as u8;
            self.serialize_u8(byte | 0x80)?;
            value >>= 7;
        }
        // Write the remaining bits of data and set the highest bit to 0.
        self.serialize_u8(value as u8)
    }

    fn serialize_variant_index(&mut self, v: u32) -> Result<()> {
        self.serialize_u32_as_uleb128(v)
    }

    /// Serialize a sequence length as a u32.
    fn serialize_seq_len(&mut self, len: usize) -> Result<()> {
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }
        self.serialize_u32_as_uleb128(len as u32)
    }
}

impl<'a> ser::Serializer for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    type SerializeSeq = Self;
    type SerializeTuple = Self;
    type SerializeTupleStruct = Self;
    type SerializeTupleVariant = Self;
    type SerializeMap = MapSerializer<'a>;
    type SerializeStruct = Self;
    type SerializeStructVariant = Self;

    fn serialize_bool(self, v: bool) -> Result<()> {
        self.serialize_u8(v.into())
    }

    fn serialize_i8(self, v: i8) -> Result<()> {
        self.serialize_u8(v as u8)
    }

    fn serialize_i16(self, v: i16) -> Result<()> {
        self.serialize_u16(v as u16)
    }

    fn serialize_i32(self, v: i32) -> Result<()> {
        self.serialize_u32(v as u32)
    }

    fn serialize_i64(self, v: i64) -> Result<()> {
        self.serialize_u64(v as u64)
    }

    fn serialize_i128(self, v: i128) -> Result<()> {
        self.serialize_u128(v as u128)
    }

    fn serialize_u8(self, v: u8) -> Result<()> {
        self.output.push(v);
        Ok(())
    }

    fn serialize_u16(self, v: u16) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u32(self, v: u32) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u64(self, v: u64) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u128(self, v: u128) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_f32(self, _v: f32) -> Result<()> {
        Err(Error::NotSupported("serialize_f32"))
    }

    fn serialize_f64(self, _v: f64) -> Result<()> {
        Err(Error::NotSupported("serialize_f64"))
    }

    fn serialize_char(self, _v: char) -> Result<()> {
        Err(Error::NotSupported("serialize_char"))
    }

    // Just serialize the string as a raw byte array
    fn serialize_str(self, v: &str) -> Result<()> {
        self.serialize_bytes(v.as_bytes())
    }

    // Serialize a byte array as an array of bytes.
    fn serialize_bytes(self, v: &[u8]) -> Result<()> {
        self.serialize_seq_len(v.len())?;
        self.output.extend_from_slice(v);
        Ok(())
    }

    // An absent optional is represented as `00`
    fn serialize_none(self) -> Result<()> {
        self.serialize_u8(0)
    }

    // A present optional is represented as `01` followed by the serialized value
    fn serialize_some<T>(self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        self.serialize_u8(1)?;
        value.serialize(self)
    }

    fn serialize_unit(self) -> Result<()> {
        Ok(())
    }

    fn serialize_unit_struct(self, _name: &'static str) -> Result<()> {
        self.serialize_unit()
    }

    fn serialize_unit_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
    ) -> Result<()> {
        self.serialize_variant_index(variant_index)
    }

    fn serialize_newtype_struct<T>(self, _name: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(self)
    }

    fn serialize_newtype_variant<T>(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        value: &T,
    ) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        self.serialize_variant_index(variant_index)?;
        value.serialize(self)
    }

    // The start of the sequence, each value, and the end are three separate
    // method calls. This one is responsible only for serializing the start,
    // which for LCS is either nothing for fixed structures or for variable
    // length structures, the length encoded as a u32.
    fn serialize_seq(self, len: Option<usize>) -> Result<Self::SerializeSeq> {
        if let Some(len) = len {
            self.serialize_seq_len(len)?;
            Ok(self)
        } else {
            Err(Error::MissingLen)
        }
    }

    // Tuples are fixed sized structs so we don't need to encode the length
    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple> {
        Ok(self)
    }

    fn serialize_tuple_struct(
        self,
        _name: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleStruct> {
        Ok(self)
    }

    fn serialize_tuple_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleVariant> {
        self.serialize_variant_index(variant_index)?;
        Ok(self)
    }

    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap> {
        Ok(MapSerializer::new(self))
    }

    fn serialize_struct(self, _name: &'static str, _len: usize) -> Result<Self::SerializeStruct> {
        Ok(self)
    }

    fn serialize_struct_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeStructVariant> {
        self.serialize_variant_index(variant_index)?;
        Ok(self)
    }

    // LCS is not a human readable format
    fn is_human_readable(&self) -> bool {
        false
    }
}

impl<'a> ser::SerializeSeq for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTuple for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTupleStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTupleVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

#[doc(hidden)]
struct MapSerializer<'a> {
    ser: &'a mut Serializer,
    map: BTreeMap<Vec<u8>, Vec<u8>>,
    next_key: Option<Vec<u8>>,
}

impl<'a> MapSerializer<'a> {
    fn new(ser: &'a mut Serializer) -> Self {
        MapSerializer {
            ser,
            map: BTreeMap::new(),
            next_key: None,
        }
    }
}

impl<'a> ser::SerializeMap for MapSerializer<'a> {
    type Ok = ();
    type Error = Error;

    fn serialize_key<T>(&mut self, key: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        if self.next_key.is_some() {
            return Err(Error::ExpectedMapValue);
        }

        let mut s = Serializer::new();
        key.serialize(&mut s)?;
        self.next_key = Some(s.output);
        Ok(())
    }

    fn serialize_value<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        match self.next_key.take() {
            Some(key) => {
                let mut s = Serializer::new();
                value.serialize(&mut s)?;
                self.map.insert(key, s.output);
            }
            None => {
                return Err(Error::ExpectedMapKey);
            }
        }

        Ok(())
    }

    fn end(self) -> Result<()> {
        if self.next_key.is_some() {
            return Err(Error::ExpectedMapValue);
        }

        let len = self.map.len();
        self.ser.serialize_seq_len(len)?;

        for (key, value) in self.map {
            self.ser.output.extend(key.into_iter());
            self.ser.output.extend(value.into_iter());
        }

        Ok(())
    }
}

impl<'a> ser::SerializeStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, _key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeStructVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, _key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}
<commit_msg>[LCS] faster serialization of maps<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::error::{Error, Result};
use serde::{ser, Serialize};

/// Serialize the given data structure as a `Vec<u8>` of LCS.
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to
/// fail, if `T` contains sequences which are longer than `MAX_SEQUENCE_LENGTH`,
/// or if `T` attempts to serialize an unsupported datatype such as a f32,
/// f64, or char.
///
/// # Examples
///
/// ```
/// use libra_canonical_serialization::to_bytes;
/// use serde::Serialize;
///
/// #[derive(Serialize)]
/// struct Ip([u8; 4]);
///
/// #[derive(Serialize)]
/// struct Port(u16);
///
/// #[derive(Serialize)]
/// struct Service {
///     ip: Ip,
///     port: Vec<Port>,
///     connection_max: Option<u32>,
///     enabled: bool,
/// }
///
/// let service = Service {
///     ip: Ip([192, 168, 1, 1]),
///     port: vec![Port(8001), Port(8002), Port(8003)],
///     connection_max: Some(5000),
///     enabled: false,
/// };
///
/// let bytes = to_bytes(&service).unwrap();
/// let expected = vec![
///     0xc0, 0xa8, 0x01, 0x01, 0x03, 0x41, 0x1f, 0x42,
///     0x1f, 0x43, 0x1f, 0x01, 0x88, 0x13, 0x00, 0x00,
///     0x00,
/// ];
/// assert_eq!(bytes, expected);
/// ```
pub fn to_bytes<T>(value: &T) -> Result<Vec<u8>>
where
    T: ?Sized + Serialize,
{
    let mut serializer = Serializer::new();
    value.serialize(&mut serializer)?;
    Ok(serializer.end())
}

pub fn is_human_readable() -> bool {
    let mut serializer = Serializer::new();
    ser::Serializer::is_human_readable(&&mut serializer)
}

/// Serialization implementation for LCS
struct Serializer {
    output: Vec<u8>,
}

impl Serializer {
    /// Creates a new `Serializer` which will emit LCS.
    fn new() -> Self {
        Self { output: Vec::new() }
    }

    /// The `Serializer::end` method should be called after a type has been
    /// fully serialized. This will return the buffer containing the LCS
    /// serialized data.
    fn end(self) -> Vec<u8> {
        self.output
    }

    fn serialize_u32_as_uleb128(&mut self, mut value: u32) -> Result<()> {
        use serde::ser::Serializer;
        while value >= 0x80 {
            // Write 7 (lowest) bits of data and set the 8th bit to 1.
            let byte = (value & 0x7f) as u8;
            self.serialize_u8(byte | 0x80)?;
            value >>= 7;
        }
        // Write the remaining bits of data and set the highest bit to 0.
        self.serialize_u8(value as u8)
    }

    fn serialize_variant_index(&mut self, v: u32) -> Result<()> {
        self.serialize_u32_as_uleb128(v)
    }

    /// Serialize a sequence length as a u32.
    fn serialize_seq_len(&mut self, len: usize) -> Result<()> {
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }
        self.serialize_u32_as_uleb128(len as u32)
    }
}

impl<'a> ser::Serializer for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    type SerializeSeq = Self;
    type SerializeTuple = Self;
    type SerializeTupleStruct = Self;
    type SerializeTupleVariant = Self;
    type SerializeMap = MapSerializer<'a>;
    type SerializeStruct = Self;
    type SerializeStructVariant = Self;

    fn serialize_bool(self, v: bool) -> Result<()> {
        self.serialize_u8(v.into())
    }

    fn serialize_i8(self, v: i8) -> Result<()> {
        self.serialize_u8(v as u8)
    }

    fn serialize_i16(self, v: i16) -> Result<()> {
        self.serialize_u16(v as u16)
    }

    fn serialize_i32(self, v: i32) -> Result<()> {
        self.serialize_u32(v as u32)
    }

    fn serialize_i64(self, v: i64) -> Result<()> {
        self.serialize_u64(v as u64)
    }

    fn serialize_i128(self, v: i128) -> Result<()> {
        self.serialize_u128(v as u128)
    }

    fn serialize_u8(self, v: u8) -> Result<()> {
        self.output.push(v);
        Ok(())
    }

    fn serialize_u16(self, v: u16) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u32(self, v: u32) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u64(self, v: u64) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_u128(self, v: u128) -> Result<()> {
        self.output.extend_from_slice(&v.to_le_bytes());
        Ok(())
    }

    fn serialize_f32(self, _v: f32) -> Result<()> {
        Err(Error::NotSupported("serialize_f32"))
    }

    fn serialize_f64(self, _v: f64) -> Result<()> {
        Err(Error::NotSupported("serialize_f64"))
    }

    fn serialize_char(self, _v: char) -> Result<()> {
        Err(Error::NotSupported("serialize_char"))
    }

    // Just serialize the string as a raw byte array
    fn serialize_str(self, v: &str) -> Result<()> {
        self.serialize_bytes(v.as_bytes())
    }

    // Serialize a byte array as an array of bytes.
    fn serialize_bytes(self, v: &[u8]) -> Result<()> {
        self.serialize_seq_len(v.len())?;
        self.output.extend_from_slice(v);
        Ok(())
    }

    // An absent optional is represented as `00`
    fn serialize_none(self) -> Result<()> {
        self.serialize_u8(0)
    }

    // A present optional is represented as `01` followed by the serialized value
    fn serialize_some<T>(self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        self.serialize_u8(1)?;
        value.serialize(self)
    }

    fn serialize_unit(self) -> Result<()> {
        Ok(())
    }

    fn serialize_unit_struct(self, _name: &'static str) -> Result<()> {
        self.serialize_unit()
    }

    fn serialize_unit_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
    ) -> Result<()> {
        self.serialize_variant_index(variant_index)
    }

    fn serialize_newtype_struct<T>(self, _name: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(self)
    }

    fn serialize_newtype_variant<T>(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        value: &T,
    ) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        self.serialize_variant_index(variant_index)?;
        value.serialize(self)
    }

    // The start of the sequence, each value, and the end are three separate
    // method calls. This one is responsible only for serializing the start,
    // which for LCS is either nothing for fixed structures or for variable
    // length structures, the length encoded as a u32.
    fn serialize_seq(self, len: Option<usize>) -> Result<Self::SerializeSeq> {
        if let Some(len) = len {
            self.serialize_seq_len(len)?;
            Ok(self)
        } else {
            Err(Error::MissingLen)
        }
    }

    // Tuples are fixed sized structs so we don't need to encode the length
    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple> {
        Ok(self)
    }

    fn serialize_tuple_struct(
        self,
        _name: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleStruct> {
        Ok(self)
    }

    fn serialize_tuple_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleVariant> {
        self.serialize_variant_index(variant_index)?;
        Ok(self)
    }

    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap> {
        Ok(MapSerializer::new(self))
    }

    fn serialize_struct(self, _name: &'static str, _len: usize) -> Result<Self::SerializeStruct> {
        Ok(self)
    }

    fn serialize_struct_variant(
        self,
        _name: &'static str,
        variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeStructVariant> {
        self.serialize_variant_index(variant_index)?;
        Ok(self)
    }

    // LCS is not a human readable format
    fn is_human_readable(&self) -> bool {
        false
    }
}

impl<'a> ser::SerializeSeq for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTuple for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTupleStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeTupleVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

#[doc(hidden)]
struct MapSerializer<'a> {
    ser: &'a mut Serializer,
    entries: Vec<(Vec<u8>, Vec<u8>)>,
    next_key: Option<Vec<u8>>,
}

impl<'a> MapSerializer<'a> {
    fn new(ser: &'a mut Serializer) -> Self {
        MapSerializer {
            ser,
            entries: Vec::new(),
            next_key: None,
        }
    }
}

impl<'a> ser::SerializeMap for MapSerializer<'a> {
    type Ok = ();
    type Error = Error;

    fn serialize_key<T>(&mut self, key: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        if self.next_key.is_some() {
            return Err(Error::ExpectedMapValue);
        }

        let mut s = Serializer::new();
        key.serialize(&mut s)?;
        self.next_key = Some(s.output);
        Ok(())
    }

    fn serialize_value<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        match self.next_key.take() {
            Some(key) => {
                let mut s = Serializer::new();
                value.serialize(&mut s)?;
                self.entries.push((key, s.output));
            }
            None => {
                return Err(Error::ExpectedMapKey);
            }
        }

        Ok(())
    }

    fn end(mut self) -> Result<()> {
        if self.next_key.is_some() {
            return Err(Error::ExpectedMapValue);
        }
        self.entries.sort_by(|e1, e2| e1.0.cmp(&e2.0));
        self.entries.dedup_by(|e1, e2| e1.0.eq(&e2.0));

        let len = self.entries.len();
        self.ser.serialize_seq_len(len)?;

        for (key, value) in self.entries {
            self.ser.output.extend(key.into_iter());
            self.ser.output.extend(value.into_iter());
        }

        Ok(())
    }
}

impl<'a> ser::SerializeStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, _key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}

impl<'a> ser::SerializeStructVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;

    fn serialize_field<T>(&mut self, _key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(&mut **self)
    }

    fn end(self) -> Result<()> {
        Ok(())
    }
}
<|endoftext|>


<commit_before><commit_msg>convert `_mm_mask_cmp_epu8_mask` to const generics<commit_after><|endoftext|>


<commit_before>//! Command like parsing for rust-analyzer.
//!
//! If run started args, we run the LSP server loop. With a subcommand, we do a
//! one-time batch processing.

use std::{env, fmt::Write, path::PathBuf};

use anyhow::{bail, format_err, Result};
use pico_args::Arguments;
use rust_analyzer::cli::{AnalysisStatsCmd, BenchCmd, BenchWhat, Position, Verbosity};
use ssr::{SsrPattern, SsrRule};
use vfs::AbsPathBuf;

pub(crate) struct Args {
    pub(crate) verbosity: Verbosity,
    pub(crate) log_file: Option<PathBuf>,
    pub(crate) command: Command,
}

pub(crate) enum Command {
    Parse { no_dump: bool },
    Symbols,
    Highlight { rainbow: bool },
    AnalysisStats(AnalysisStatsCmd),
    Bench(BenchCmd),
    Diagnostics { path: PathBuf, load_output_dirs: bool, with_proc_macro: bool },
    Ssr { rules: Vec<SsrRule> },
    StructuredSearch { debug_snippet: Option<String>, patterns: Vec<SsrPattern> },
    ProcMacro,
    RunServer,
    Version,
    Help,
}

const HELP: &str = "\
rust-analyzer

USAGE:
    rust-analyzer [FLAGS] [COMMAND] [COMMAND_OPTIONS]

FLAGS:
    --version         Print version
    -h, --help        Print this help

    -v,  --verbose
    -vv, --spammy
    -q,  --quiet      Set verbosity

    --log-file <PATH> Log to the specified filed instead of stderr

ENVIRONMENTAL VARIABLES:
    RA_LOG            Set log filter in env_logger format
    RA_PROFILE        Enable hierarchical profiler

COMMANDS:

not specified         Launch LSP server

parse < main.rs       Parse tree
    --no-dump         Suppress printing

symbols < main.rs     Parse input an print the list of symbols

highlight < main.rs   Highlight input as html
    --rainbow         Enable rainbow highlighting of identifiers

analysis-stats <PATH> Batch typecheck project and print summary statistics
    <PATH>            Directory with Cargo.toml
    --randomize       Randomize order in which crates, modules, and items are processed
    --parallel        Run type inference in parallel
    --memory-usage    Collect memory usage statistics
    -o, --only <PATH> Only analyze items matching this path
    --with-deps       Also analyze all dependencies
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

analysis-bench <PATH> Benchmark specific analysis operation
    <PATH>            Directory with Cargo.toml
    --highlight <PATH>
                      Compute syntax highlighting for this file
    --complete <PATH:LINE:COLUMN>
                      Compute completions at this location
    --goto-def <PATH:LINE:COLUMN>
                      Compute goto definition at this location
    --memory-usage    Collect memory usage statistics
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

diagnostics <PATH>
    <PATH>            Directory with Cargo.toml
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

ssr [RULE...]
    <RULE>            A structured search replace rule (`$a.foo($b) ==> bar($a, $b)`)

search [PATTERN..]
    <PATTERN>         A structured search replace pattern (`$a.foo($b)`)
    --debug <snippet> Prints debug information for any nodes with source exactly
                      equal to <snippet>
";

impl Args {
    pub(crate) fn parse() -> Result<Args> {
        let mut matches = Arguments::from_env();

        if matches.contains("--version") {
            matches.finish().or_else(handle_extra_flags)?;
            return Ok(Args {
                verbosity: Verbosity::Normal,
                log_file: None,
                command: Command::Version,
            });
        }

        let verbosity = match (
            matches.contains(["-vv", "--spammy"]),
            matches.contains(["-v", "--verbose"]),
            matches.contains(["-q", "--quiet"]),
        ) {
            (true, _, true) => bail!("Invalid flags: -q conflicts with -vv"),
            (true, _, false) => Verbosity::Spammy,
            (false, false, false) => Verbosity::Normal,
            (false, false, true) => Verbosity::Quiet,
            (false, true, false) => Verbosity::Verbose,
            (false, true, true) => bail!("Invalid flags: -q conflicts with -v"),
        };
        let log_file = matches.opt_value_from_str("--log-file")?;

        if matches.contains(["-h", "--help"]) {
            eprintln!("{}", HELP);
            return Ok(Args { verbosity, log_file: None, command: Command::Help });
        }

        let subcommand = match matches.subcommand()? {
            Some(it) => it,
            None => {
                matches.finish().or_else(handle_extra_flags)?;
                return Ok(Args { verbosity, log_file, command: Command::RunServer });
            }
        };
        let command = match subcommand.as_str() {
            "parse" => Command::Parse { no_dump: matches.contains("--no-dump") },
            "symbols" => Command::Symbols,
            "highlight" => Command::Highlight { rainbow: matches.contains("--rainbow") },
            "analysis-stats" => Command::AnalysisStats(AnalysisStatsCmd {
                randomize: matches.contains("--randomize"),
                parallel: matches.contains("--parallel"),
                memory_usage: matches.contains("--memory-usage"),
                only: matches.opt_value_from_str(["-o", "--only"])?,
                with_deps: matches.contains("--with-deps"),
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            }),
            "analysis-bench" => Command::Bench(BenchCmd {
                what: {
                    let highlight_path: Option<String> =
                        matches.opt_value_from_str("--highlight")?;
                    let complete_path: Option<Position> =
                        matches.opt_value_from_str("--complete")?;
                    let goto_def_path: Option<Position> =
                        matches.opt_value_from_str("--goto-def")?;
                    match (highlight_path, complete_path, goto_def_path) {
                            (Some(path), None, None) => {
                                let path = env::current_dir().unwrap().join(path);
                                BenchWhat::Highlight { path: AbsPathBuf::assert(path) }
                            }
                            (None, Some(position), None) => BenchWhat::Complete(position),
                            (None, None, Some(position)) => BenchWhat::GotoDef(position),
                            _ => panic!(
                                "exactly one of  `--highlight`, `--complete` or `--goto-def` must be set"
                            ),
                        }
                },
                memory_usage: matches.contains("--memory-usage"),
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            }),
            "diagnostics" => Command::Diagnostics {
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            },
            "proc-macro" => Command::ProcMacro,
            "ssr" => Command::Ssr {
                rules: {
                    let mut acc = Vec::new();
                    while let Some(rule) = matches.free_from_str()? {
                        acc.push(rule);
                    }
                    acc
                },
            },
            "search" => Command::StructuredSearch {
                debug_snippet: matches.opt_value_from_str("--debug")?,
                patterns: {
                    let mut acc = Vec::new();
                    while let Some(rule) = matches.free_from_str()? {
                        acc.push(rule);
                    }
                    acc
                },
            },
            _ => {
                eprintln!("{}", HELP);
                return Ok(Args { verbosity, log_file: None, command: Command::Help });
            }
        };
        matches.finish().or_else(handle_extra_flags)?;
        Ok(Args { verbosity, log_file, command })
    }
}

fn handle_extra_flags(e: pico_args::Error) -> Result<()> {
    if let pico_args::Error::UnusedArgsLeft(flags) = e {
        let mut invalid_flags = String::new();
        for flag in flags {
            write!(&mut invalid_flags, "{}, ", flag)?;
        }
        let (invalid_flags, _) = invalid_flags.split_at(invalid_flags.len() - 2);
        bail!("Invalid flags: {}", invalid_flags);
    } else {
        bail!(e);
    }
}
<commit_msg>Simplify<commit_after>//! Command like parsing for rust-analyzer.
//!
//! If run started args, we run the LSP server loop. With a subcommand, we do a
//! one-time batch processing.

use std::{env, path::PathBuf};

use anyhow::{bail, format_err, Result};
use pico_args::Arguments;
use rust_analyzer::cli::{AnalysisStatsCmd, BenchCmd, BenchWhat, Position, Verbosity};
use ssr::{SsrPattern, SsrRule};
use vfs::AbsPathBuf;

pub(crate) struct Args {
    pub(crate) verbosity: Verbosity,
    pub(crate) log_file: Option<PathBuf>,
    pub(crate) command: Command,
}

pub(crate) enum Command {
    Parse { no_dump: bool },
    Symbols,
    Highlight { rainbow: bool },
    AnalysisStats(AnalysisStatsCmd),
    Bench(BenchCmd),
    Diagnostics { path: PathBuf, load_output_dirs: bool, with_proc_macro: bool },
    Ssr { rules: Vec<SsrRule> },
    StructuredSearch { debug_snippet: Option<String>, patterns: Vec<SsrPattern> },
    ProcMacro,
    RunServer,
    Version,
    Help,
}

const HELP: &str = "\
rust-analyzer

USAGE:
    rust-analyzer [FLAGS] [COMMAND] [COMMAND_OPTIONS]

FLAGS:
    --version         Print version
    -h, --help        Print this help

    -v,  --verbose
    -vv, --spammy
    -q,  --quiet      Set verbosity

    --log-file <PATH> Log to the specified filed instead of stderr

ENVIRONMENTAL VARIABLES:
    RA_LOG            Set log filter in env_logger format
    RA_PROFILE        Enable hierarchical profiler

COMMANDS:

not specified         Launch LSP server

parse < main.rs       Parse tree
    --no-dump         Suppress printing

symbols < main.rs     Parse input an print the list of symbols

highlight < main.rs   Highlight input as html
    --rainbow         Enable rainbow highlighting of identifiers

analysis-stats <PATH> Batch typecheck project and print summary statistics
    <PATH>            Directory with Cargo.toml
    --randomize       Randomize order in which crates, modules, and items are processed
    --parallel        Run type inference in parallel
    --memory-usage    Collect memory usage statistics
    -o, --only <PATH> Only analyze items matching this path
    --with-deps       Also analyze all dependencies
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

analysis-bench <PATH> Benchmark specific analysis operation
    <PATH>            Directory with Cargo.toml
    --highlight <PATH>
                      Compute syntax highlighting for this file
    --complete <PATH:LINE:COLUMN>
                      Compute completions at this location
    --goto-def <PATH:LINE:COLUMN>
                      Compute goto definition at this location
    --memory-usage    Collect memory usage statistics
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

diagnostics <PATH>
    <PATH>            Directory with Cargo.toml
    --load-output-dirs
                      Load OUT_DIR values by running `cargo check` before analysis
    --with-proc-macro Use proc-macro-srv for proc-macro expanding

ssr [RULE...]
    <RULE>            A structured search replace rule (`$a.foo($b) ==> bar($a, $b)`)

search [PATTERN..]
    <PATTERN>         A structured search replace pattern (`$a.foo($b)`)
    --debug <snippet> Prints debug information for any nodes with source exactly
                      equal to <snippet>
";

impl Args {
    pub(crate) fn parse() -> Result<Args> {
        let mut matches = Arguments::from_env();

        if matches.contains("--version") {
            matches.finish()?;
            return Ok(Args {
                verbosity: Verbosity::Normal,
                log_file: None,
                command: Command::Version,
            });
        }

        let verbosity = match (
            matches.contains(["-vv", "--spammy"]),
            matches.contains(["-v", "--verbose"]),
            matches.contains(["-q", "--quiet"]),
        ) {
            (true, _, true) => bail!("Invalid flags: -q conflicts with -vv"),
            (true, _, false) => Verbosity::Spammy,
            (false, false, false) => Verbosity::Normal,
            (false, false, true) => Verbosity::Quiet,
            (false, true, false) => Verbosity::Verbose,
            (false, true, true) => bail!("Invalid flags: -q conflicts with -v"),
        };
        let log_file = matches.opt_value_from_str("--log-file")?;

        if matches.contains(["-h", "--help"]) {
            eprintln!("{}", HELP);
            return Ok(Args { verbosity, log_file: None, command: Command::Help });
        }

        let subcommand = match matches.subcommand()? {
            Some(it) => it,
            None => {
                matches.finish()?;
                return Ok(Args { verbosity, log_file, command: Command::RunServer });
            }
        };
        let command = match subcommand.as_str() {
            "parse" => Command::Parse { no_dump: matches.contains("--no-dump") },
            "symbols" => Command::Symbols,
            "highlight" => Command::Highlight { rainbow: matches.contains("--rainbow") },
            "analysis-stats" => Command::AnalysisStats(AnalysisStatsCmd {
                randomize: matches.contains("--randomize"),
                parallel: matches.contains("--parallel"),
                memory_usage: matches.contains("--memory-usage"),
                only: matches.opt_value_from_str(["-o", "--only"])?,
                with_deps: matches.contains("--with-deps"),
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            }),
            "analysis-bench" => Command::Bench(BenchCmd {
                what: {
                    let highlight_path: Option<String> =
                        matches.opt_value_from_str("--highlight")?;
                    let complete_path: Option<Position> =
                        matches.opt_value_from_str("--complete")?;
                    let goto_def_path: Option<Position> =
                        matches.opt_value_from_str("--goto-def")?;
                    match (highlight_path, complete_path, goto_def_path) {
                            (Some(path), None, None) => {
                                let path = env::current_dir().unwrap().join(path);
                                BenchWhat::Highlight { path: AbsPathBuf::assert(path) }
                            }
                            (None, Some(position), None) => BenchWhat::Complete(position),
                            (None, None, Some(position)) => BenchWhat::GotoDef(position),
                            _ => panic!(
                                "exactly one of  `--highlight`, `--complete` or `--goto-def` must be set"
                            ),
                        }
                },
                memory_usage: matches.contains("--memory-usage"),
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            }),
            "diagnostics" => Command::Diagnostics {
                load_output_dirs: matches.contains("--load-output-dirs"),
                with_proc_macro: matches.contains("--with-proc-macro"),
                path: matches
                    .free_from_str()?
                    .ok_or_else(|| format_err!("expected positional argument"))?,
            },
            "proc-macro" => Command::ProcMacro,
            "ssr" => Command::Ssr {
                rules: {
                    let mut acc = Vec::new();
                    while let Some(rule) = matches.free_from_str()? {
                        acc.push(rule);
                    }
                    acc
                },
            },
            "search" => Command::StructuredSearch {
                debug_snippet: matches.opt_value_from_str("--debug")?,
                patterns: {
                    let mut acc = Vec::new();
                    while let Some(rule) = matches.free_from_str()? {
                        acc.push(rule);
                    }
                    acc
                },
            },
            _ => {
                eprintln!("{}", HELP);
                return Ok(Args { verbosity, log_file: None, command: Command::Help });
            }
        };
        matches.finish()?;
        Ok(Args { verbosity, log_file, command })
    }
}
<|endoftext|>


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

use std::{env, fs, path::Path, process};

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

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

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

fn main() {
    if std::env::var("RA_RUSTC_WRAPPER").is_ok() {
        let mut args = std::env::args_os();
        let _me = args.next().unwrap();
        let rustc = args.next().unwrap();
        let code = match rustc_wrapper::run_rustc_skipping_cargo_checking(rustc, args.collect()) {
            Ok(rustc_wrapper::ExitCode(code)) => code.unwrap_or(102),
            Err(err) => {
                eprintln!("{}", err);
                101
            }
        };
        process::exit(code);
    }

    if let Err(err) = try_main() {
        tracing::error!("Unexpected error: {}", err);
        eprintln!("{}", err);
        process::exit(101);
    }
}

fn try_main() -> Result<()> {
    let flags = flags::RustAnalyzer::from_env()?;

    #[cfg(debug_assertions)]
    if flags.wait_dbg || env::var("RA_WAIT_DBG").is_ok() {
        #[allow(unused_mut)]
        let mut d = 4;
        while d == 4 {
            d = 4;
        }
    }

    let mut log_file = flags.log_file.as_deref();

    let env_log_file = env::var("RA_LOG_FILE").ok();
    if let Some(env_log_file) = env_log_file.as_deref() {
        log_file = Some(Path::new(env_log_file));
    }

    setup_logging(log_file)?;
    let verbosity = flags.verbosity();

    match flags.subcommand {
        flags::RustAnalyzerCmd::LspServer(cmd) => {
            if cmd.print_config_schema {
                println!("{:#}", Config::json_schema());
                return Ok(());
            }
            if cmd.version {
                println!("rust-analyzer {}", env!("REV"));
                return Ok(());
            }
            if cmd.help {
                println!("{}", flags::RustAnalyzer::HELP);
                return Ok(());
            }
            run_server()?
        }
        flags::RustAnalyzerCmd::ProcMacro(flags::ProcMacro) => proc_macro_srv::cli::run()?,
        flags::RustAnalyzerCmd::Parse(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Symbols(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Highlight(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::AnalysisStats(cmd) => cmd.run(verbosity)?,
        flags::RustAnalyzerCmd::Diagnostics(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Ssr(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Search(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Lsif(cmd) => cmd.run()?,
    }
    Ok(())
}

fn setup_logging(log_file: Option<&Path>) -> Result<()> {
    if cfg!(windows) {
        // This is required so that windows finds our pdb that is placed right beside the exe.
        // By default it doesn't look at the folder the exe resides in, only in the current working
        // directory which we set to the project workspace.
        // https://docs.microsoft.com/en-us/windows-hardware/drivers/debugger/general-environment-variables
        // https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-syminitialize
        if let Ok(path) = env::current_exe() {
            if let Some(path) = path.parent() {
                env::set_var("_NT_SYMBOL_PATH", path);
            }
        }
    }
    if env::var("RUST_BACKTRACE").is_err() {
        env::set_var("RUST_BACKTRACE", "short");
    }

    let log_file = match log_file {
        Some(path) => {
            if let Some(parent) = path.parent() {
                let _ = fs::create_dir_all(parent);
            }
            Some(fs::File::create(path)?)
        }
        None => None,
    };
    let filter = env::var("RA_LOG").ok();
    // deliberately enable all `error` logs if the user has not set RA_LOG, as there is usually useful
    // information in there for debugging
    logger::Logger::new(log_file, filter.as_deref().or(Some("error"))).install()?;

    profile::init();

    Ok(())
}

fn run_server() -> Result<()> {
    tracing::info!("server version {} will start", env!("REV"));

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

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

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

    let mut config = Config::new(root_path, initialize_params.capabilities);
    if let Some(json) = initialize_params.initialization_options {
        if let Err(e) = config.update(json) {
            use lsp_types::{
                notification::{Notification, ShowMessage},
                MessageType, ShowMessageParams,
            };
            let not = lsp_server::Notification::new(
                ShowMessage::METHOD.to_string(),
                ShowMessageParams { typ: MessageType::WARNING, message: e.to_string() },
            );
            connection.sender.send(lsp_server::Message::Notification(not)).unwrap();
        }
    }

    let server_capabilities = rust_analyzer::server_capabilities(&config);

    let initialize_result = lsp_types::InitializeResult {
        capabilities: server_capabilities,
        server_info: Some(lsp_types::ServerInfo {
            name: String::from("rust-analyzer"),
            version: Some(String::from(env!("REV"))),
        }),
        offset_encoding: if supports_utf8(config.caps()) {
            Some("utf-8".to_string())
        } else {
            None
        },
    };

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

    connection.initialize_finish(initialize_id, initialize_result)?;

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

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

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

    rust_analyzer::main_loop(config, connection)?;

    io_threads.join()?;
    tracing::info!("server did shut down");
    Ok(())
}
<commit_msg>Spawn a new thread with a larger stack for the server<commit_after>//! Driver for rust-analyzer.
//!
//! Based on cli flags, either spawns an LSP server, or runs a batch analysis
mod logger;
mod rustc_wrapper;

use std::{env, fs, path::Path, process};

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

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

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

fn main() {
    if std::env::var("RA_RUSTC_WRAPPER").is_ok() {
        let mut args = std::env::args_os();
        let _me = args.next().unwrap();
        let rustc = args.next().unwrap();
        let code = match rustc_wrapper::run_rustc_skipping_cargo_checking(rustc, args.collect()) {
            Ok(rustc_wrapper::ExitCode(code)) => code.unwrap_or(102),
            Err(err) => {
                eprintln!("{}", err);
                101
            }
        };
        process::exit(code);
    }

    if let Err(err) = try_main() {
        tracing::error!("Unexpected error: {}", err);
        eprintln!("{}", err);
        process::exit(101);
    }
}

fn try_main() -> Result<()> {
    let flags = flags::RustAnalyzer::from_env()?;

    #[cfg(debug_assertions)]
    if flags.wait_dbg || env::var("RA_WAIT_DBG").is_ok() {
        #[allow(unused_mut)]
        let mut d = 4;
        while d == 4 {
            d = 4;
        }
    }

    let mut log_file = flags.log_file.as_deref();

    let env_log_file = env::var("RA_LOG_FILE").ok();
    if let Some(env_log_file) = env_log_file.as_deref() {
        log_file = Some(Path::new(env_log_file));
    }

    setup_logging(log_file)?;
    let verbosity = flags.verbosity();

    match flags.subcommand {
        flags::RustAnalyzerCmd::LspServer(cmd) => {
            if cmd.print_config_schema {
                println!("{:#}", Config::json_schema());
                return Ok(());
            }
            if cmd.version {
                println!("rust-analyzer {}", env!("REV"));
                return Ok(());
            }
            if cmd.help {
                println!("{}", flags::RustAnalyzer::HELP);
                return Ok(());
            }
            with_extra_thread("rust-analyzer server thread", run_server)?;
        }
        flags::RustAnalyzerCmd::ProcMacro(flags::ProcMacro) => {
            with_extra_thread("rust-analyzer proc-macro expander", || {
                proc_macro_srv::cli::run().map_err(Into::into)
            })?;
        }
        flags::RustAnalyzerCmd::Parse(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Symbols(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Highlight(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::AnalysisStats(cmd) => cmd.run(verbosity)?,
        flags::RustAnalyzerCmd::Diagnostics(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Ssr(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Search(cmd) => cmd.run()?,
        flags::RustAnalyzerCmd::Lsif(cmd) => cmd.run()?,
    }
    Ok(())
}

fn setup_logging(log_file: Option<&Path>) -> Result<()> {
    if cfg!(windows) {
        // This is required so that windows finds our pdb that is placed right beside the exe.
        // By default it doesn't look at the folder the exe resides in, only in the current working
        // directory which we set to the project workspace.
        // https://docs.microsoft.com/en-us/windows-hardware/drivers/debugger/general-environment-variables
        // https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-syminitialize
        if let Ok(path) = env::current_exe() {
            if let Some(path) = path.parent() {
                env::set_var("_NT_SYMBOL_PATH", path);
            }
        }
    }
    if env::var("RUST_BACKTRACE").is_err() {
        env::set_var("RUST_BACKTRACE", "short");
    }

    let log_file = match log_file {
        Some(path) => {
            if let Some(parent) = path.parent() {
                let _ = fs::create_dir_all(parent);
            }
            Some(fs::File::create(path)?)
        }
        None => None,
    };
    let filter = env::var("RA_LOG").ok();
    // deliberately enable all `error` logs if the user has not set RA_LOG, as there is usually useful
    // information in there for debugging
    logger::Logger::new(log_file, filter.as_deref().or(Some("error"))).install()?;

    profile::init();

    Ok(())
}

const STACK_SIZE: usize = 1024 * 1024 * 8;

/// Parts of rust-analyzer can use a lot of stack space, and some operating systems only give us
/// 1 MB by default (eg. Windows), so this spawns a new thread with hopefully sufficient stack
/// space.
fn with_extra_thread(
    thread_name: impl Into<String>,
    f: impl FnOnce() -> Result<()> + Send + 'static,
) -> Result<()> {
    let handle =
        std::thread::Builder::new().name(thread_name.into()).stack_size(STACK_SIZE).spawn(f)?;
    match handle.join() {
        Ok(res) => res,
        Err(panic) => std::panic::resume_unwind(panic),
    }
}

fn run_server() -> Result<()> {
    tracing::info!("server version {} will start", env!("REV"));

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

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

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

    let mut config = Config::new(root_path, initialize_params.capabilities);
    if let Some(json) = initialize_params.initialization_options {
        if let Err(e) = config.update(json) {
            use lsp_types::{
                notification::{Notification, ShowMessage},
                MessageType, ShowMessageParams,
            };
            let not = lsp_server::Notification::new(
                ShowMessage::METHOD.to_string(),
                ShowMessageParams { typ: MessageType::WARNING, message: e.to_string() },
            );
            connection.sender.send(lsp_server::Message::Notification(not)).unwrap();
        }
    }

    let server_capabilities = rust_analyzer::server_capabilities(&config);

    let initialize_result = lsp_types::InitializeResult {
        capabilities: server_capabilities,
        server_info: Some(lsp_types::ServerInfo {
            name: String::from("rust-analyzer"),
            version: Some(String::from(env!("REV"))),
        }),
        offset_encoding: if supports_utf8(config.caps()) {
            Some("utf-8".to_string())
        } else {
            None
        },
    };

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

    connection.initialize_finish(initialize_id, initialize_result)?;

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

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

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

    rust_analyzer::main_loop(config, connection)?;

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


<commit_before><commit_msg>Removed file with duplicated benchmarks<commit_after><|endoftext|>


<commit_before>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let mut unescaped = String::with_capacity(inside_str.len());
        let mut error = Ok(());
        rustc_lexer::unescape::unescape_str(inside_str, &mut |_, unescaped_char| {
            match unescaped_char {
                Ok(c) => unescaped.push(c),
                Err(_) => error = Err(()),
            }
        });
        if error.is_err() {
            eprintln!("Error unescaping string");
        } else {
            let max_hash_streak = count_hashes(&unescaped);
            let mut hashes = String::with_capacity(max_hash_streak + 1);
            for _ in 0..hashes.capacity() {
                hashes.push('#');
            }
            edit.replace(
                literal.syntax().text_range(),
                format!("r{}\"{}\"{}", hashes, unescaped, hashes),
            );
        }
    });
    ctx.build()
}

fn count_hashes(s: &str) -> usize {
    let indexes: Vec<_> = s.match_indices("\"#").map(|(i, _)| i).collect();
    let mut max_hash_streak = 0usize;
    for idx in indexes {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(literal.syntax().text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(literal.syntax().text_range());
        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(literal.syntax().text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(literal.syntax().text_range(), format!("r{}", result));
    });
    ctx.build()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<commit_msg>feat(assists): Address some PR comments<commit_after>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    let start_of_inside = usual_string_range.start().to_usize() + 1;
    let end_of_inside = usual_string_range.end().to_usize();
    let inside_str = &text[start_of_inside..end_of_inside];
    let mut unescaped = String::with_capacity(inside_str.len());
    let mut error = Ok(());
    rustc_lexer::unescape::unescape_str(
        inside_str,
        &mut |_, unescaped_char| match unescaped_char {
            Ok(c) => unescaped.push(c),
            Err(_) => error = Err(()),
        },
    );
    if error.is_err() {
        return None;
    }
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let max_hash_streak = count_hashes(&unescaped);
        let mut hashes = String::with_capacity(max_hash_streak + 1);
        for _ in 0..hashes.capacity() {
            hashes.push('#');
        }
        edit.replace(
            literal.syntax().text_range(),
            format!("r{}\"{}\"{}", hashes, unescaped, hashes),
        );
    });
    ctx.build()
}

fn count_hashes(s: &str) -> usize {
    let mut max_hash_streak = 0usize;
    for idx in s.match_indices("\"#").map(|(i, _)| i) {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(literal.syntax().text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(literal.syntax().text_range());
        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(literal.syntax().text_range().end(), "#");
    });
    ctx.build()
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
        return None;
    }
    let token = literal.token();
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(literal.syntax().text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(literal.syntax().text_range(), format!("r{}", result));
    });
    ctx.build()
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<|endoftext|>


<commit_before>use std::fmt::Write;

use hir::{AdtDef, Ty, source_binder};
use hir::db::HirDatabase;

use ra_syntax::ast::{self, AstNode};

use crate::{AssistCtx, Assist, AssistId};

pub(crate) fn fill_struct_fields(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let struct_lit = ctx.node_at_offset::<ast::StructLit>()?;
    let named_field_list = struct_lit.named_field_list()?;

    // If we already have existing struct fields, don't provide the assist.
    if named_field_list.fields().count() > 0 {
        return None;
    }

    let function =
        source_binder::function_from_child_node(ctx.db, ctx.frange.file_id, struct_lit.syntax())?;

    let infer_result = function.infer(ctx.db);
    let source_map = function.body_source_map(ctx.db);
    let node_expr = source_map.node_expr(struct_lit.into())?;
    let struct_lit_ty = infer_result[node_expr].clone();
    let struct_def = match struct_lit_ty {
        Ty::Adt { def_id: AdtDef::Struct(s), .. } => s,
        _ => return None,
    };

    let db = ctx.db;
    ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
        let mut buf = String::from("{\n");
        let struct_fields = struct_def.fields(db);
        for field in struct_fields {
            let field_name = field.name(db).to_string();
            write!(&mut buf, "    {}: (),\n", field_name).unwrap();
        }
        buf.push_str("}");

        edit.target(struct_lit.syntax().range());
        edit.set_cursor(struct_lit.syntax().range().start());
        edit.replace_node_and_indent(named_field_list.syntax(), buf);
    });

    ctx.build()
}

#[cfg(test)]
mod tests {
    use crate::helpers::{check_assist, check_assist_target};

    use super::fill_struct_fields;

    #[test]
    fn fill_struct_fields_empty_body() {
        check_assist(
            fill_struct_fields,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                r: &'a str,
            }

            fn main() {
                let s = S<|> {}
            }
            "#,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                r: &'a str,
            }

            fn main() {
                let s = <|>S {
                    a: (),
                    b: (),
                    c: (),
                    d: (),
                    r: (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_struct_fields_target() {
        check_assist_target(
            fill_struct_fields,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                r: &'a str,
            }

            fn main() {
                let s = S<|> {}
            }
            "#,
            "S {}",
        );
    }

    #[test]
    fn fill_struct_fields_preserve_self() {
        check_assist(
            fill_struct_fields,
            r#"
            struct Foo {
                foo: u8,
                bar: String,
                baz: i128,
            }

            impl Foo {
                pub fn new() -> Self {
                    Self <|>{}
                }
            }
            "#,
            r#"
            struct Foo {
                foo: u8,
                bar: String,
                baz: i128,
            }

            impl Foo {
                pub fn new() -> Self {
                    <|>Self {
                        foo: (),
                        bar: (),
                        baz: (),
                    }
                }
            }
            "#,
        );
    }
}
<commit_msg>Fill partial fields<commit_after>use std::fmt::Write;

use hir::{AdtDef, Ty, db::HirDatabase, source_binder::function_from_child_node, StructField};

use ra_syntax::ast::{self, AstNode};

use crate::{AssistCtx, Assist, AssistId};

pub(crate) fn fill_struct_fields(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let struct_lit = ctx.node_at_offset::<ast::StructLit>()?;
    let mut fsf = FillStructFields {
        ctx: &mut ctx,
        named_field_list: struct_lit.named_field_list()?,
        struct_fields: vec![],
        struct_lit,
    };
    fsf.evaluate_struct_def_fields()?;
    if fsf.struct_lit_and_def_have_the_same_number_of_fields() {
        return None;
    }
    fsf.remove_already_included_fields()?;
    fsf.add_action()?;
    ctx.build()
}

struct FillStructFields<'a, 'b: 'a, DB> {
    ctx: &'a mut AssistCtx<'b, DB>,
    named_field_list: &'a ast::NamedFieldList,
    struct_fields: Vec<StructField>,
    struct_lit: &'a ast::StructLit,
}

impl<DB> FillStructFields<'_, '_, DB>
where
    DB: HirDatabase,
{
    fn add_action(&mut self) -> Option<()> {
        let named_field_list = self.named_field_list;
        let struct_fields_string = self.struct_fields_string()?;
        let struct_lit = self.struct_lit;
        self.ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
            edit.target(struct_lit.syntax().range());
            edit.set_cursor(struct_lit.syntax().range().start());
            edit.replace_node_and_indent(named_field_list.syntax(), struct_fields_string);
        });
        Some(())
    }

    fn struct_lit_and_def_have_the_same_number_of_fields(&self) -> bool {
        self.named_field_list.fields().count() == self.struct_fields.len()
    }

    fn evaluate_struct_def_fields(&mut self) -> Option<()> {
        let function = function_from_child_node(
            self.ctx.db,
            self.ctx.frange.file_id,
            self.struct_lit.syntax(),
        )?;
        let infer_result = function.infer(self.ctx.db);
        let source_map = function.body_source_map(self.ctx.db);
        let node_expr = source_map.node_expr(self.struct_lit.into())?;
        let struct_lit_ty = infer_result[node_expr].clone();
        let struct_def = match struct_lit_ty {
            Ty::Adt { def_id: AdtDef::Struct(s), .. } => s,
            _ => return None,
        };
        self.struct_fields = struct_def.fields(self.ctx.db);
        Some(())
    }

    fn remove_already_included_fields(&mut self) -> Option<()> {
        for ast_field in self.named_field_list.fields() {
            let name_from_ast = ast_field.name_ref()?.text().to_string();
            if let Some(idx) = self
                .struct_fields
                .iter()
                .map(|f| f.name(self.ctx.db).to_string())
                .position(|n| n == name_from_ast)
            {
                self.struct_fields.remove(idx);
            }
        }
        Some(())
    }

    fn struct_fields_string(&self) -> Option<String> {
        let mut buf = String::from("{\n");
        for field in self.named_field_list.fields() {
            let expr = field.expr()?.syntax().text().to_string();
            let field_name = field.name_ref()?.syntax().text().to_string();
            write!(&mut buf, "    {}: {},\n", field_name, expr).unwrap();
        }
        for field in &self.struct_fields {
            let field_name = field.name(self.ctx.db).to_string();
            write!(&mut buf, "    {}: (),\n", field_name).unwrap();
        }
        buf.push_str("}");
        Some(buf)
    }
}

#[cfg(test)]
mod tests {
    use crate::helpers::{check_assist, check_assist_target};

    use super::fill_struct_fields;

    #[test]
    fn fill_struct_fields_empty_body() {
        check_assist(
            fill_struct_fields,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                e: &'a str,
            }

            fn main() {
                let s = S<|> {}
            }
            "#,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                e: &'a str,
            }

            fn main() {
                let s = <|>S {
                    a: (),
                    b: (),
                    c: (),
                    d: (),
                    e: (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_struct_fields_target() {
        check_assist_target(
            fill_struct_fields,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                e: &'a str,
            }

            fn main() {
                let s = S<|> {}
            }
            "#,
            "S {}",
        );
    }

    #[test]
    fn fill_struct_fields_preserve_self() {
        check_assist(
            fill_struct_fields,
            r#"
            struct Foo {
                foo: u8,
                bar: String,
                baz: i128,
            }

            impl Foo {
                pub fn new() -> Self {
                    Self <|>{}
                }
            }
            "#,
            r#"
            struct Foo {
                foo: u8,
                bar: String,
                baz: i128,
            }

            impl Foo {
                pub fn new() -> Self {
                    <|>Self {
                        foo: (),
                        bar: (),
                        baz: (),
                    }
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_struct_fields_partial() {
        check_assist(
            fill_struct_fields,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                e: &'a str,
            }

            fn main() {
                let s = S {
                    c: (1, 2),
                    e: "foo",<|>
                }
            }
            "#,
            r#"
            struct S<'a, D> {
                a: u32,
                b: String,
                c: (i32, i32),
                d: D,
                e: &'a str,
            }

            fn main() {
                let s = <|>S {
                    c: (1, 2),
                    e: "foo",
                    a: (),
                    b: (),
                    d: (),
                }
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>use parser::Node;
use std::str;

pub fn generate(ctx_name: &str, tokens: &Vec<Node>) -> String {
    let mut code = String::new();
    code.push_str("impl askama::Template for ");
    code.push_str(ctx_name);
    code.push_str(" {\n");
    code.push_str("    fn render(&self) -> String {\n");
    code.push_str("        let mut buf = String::new();\n");
    for n in tokens {
        code.push_str("        buf.push_str(");
        code.push_str(&match n {
            &Node::Lit(val) => format!("{:#?}", str::from_utf8(val).unwrap()),
            &Node::Expr(val) => format!("&self.{}", str::from_utf8(val).unwrap()),
        });
        code.push_str(");\n");
    }
    code.push_str("        buf");
    code.push_str("    }\n");
    code.push_str("}\n\n");
    code
}
<commit_msg>Introduce code generator impl<commit_after>use parser::Node;
use std::str;

struct Generator {
    buf: String,
}

impl Generator {
    fn new() -> Generator {
        Generator { buf: String::new() }
    }
    fn write(&mut self, s: &str) {
        self.buf.push_str(s);
    }
    fn result(self) -> String {
        self.buf
    }
}

pub fn generate(ctx_name: &str, tokens: &Vec<Node>) -> String {
    let mut gen = Generator::new();
    gen.write("impl askama::Template for ");
    gen.write(ctx_name);
    gen.write(" {\n");
    gen.write("    fn render(&self) -> String {\n");
    gen.write("        let mut buf = String::new();\n");
    for n in tokens {
        gen.write("        buf.push_str(");
        gen.write(&match n {
            &Node::Lit(val) => format!("{:#?}", str::from_utf8(val).unwrap()),
            &Node::Expr(val) => format!("&self.{}", str::from_utf8(val).unwrap()),
        });
        gen.write(");\n");
    }
    gen.write("        buf");
    gen.write("    }\n");
    gen.write("}\n\n");
    gen.result()
}
<|endoftext|>


<commit_before>use std::io::Read;

use common;

#[derive(Default)]
pub struct Day05 {}

impl Day05 {
    pub fn new() -> Day05 {
        Default::default()
    }

    fn reduce(mut data: Vec<u8>) -> usize {
        let mut dptr = 0;
        for iptr in 0..data.len() {
            if dptr > 0 && (data[iptr].eq_ignore_ascii_case(&data[dptr - 1])) && data[iptr] != data[dptr - 1] {
                dptr -= 1;
            } else {
                data[dptr] = data[iptr];
                dptr += 1;
            }
        }

        dptr
    }
}

impl common::Solution for Day05 {
    fn part1(&mut self, input: &mut Read) -> String {
        let mut data = Vec::new();
        input.read_to_end(&mut data).expect("Can't read input!");
        common::trim_back(&mut data);

        format!("{}", Day05::reduce(data))
    }

    fn part2(&mut self, input: &mut Read) -> String {
        let mut data = Vec::new();
        input.read_to_end(&mut data).expect("Can't read input!");
        common::trim_back(&mut data);

        let min_len = (b'a'..=b'z').map(|option| data.iter().filter(|&x| !x.eq_ignore_ascii_case(&option)).map(|&x| x).collect())
            .map(|x| Day05::reduce(x))
            .min().unwrap();

        format!("{}", min_len)
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day05::Day05;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/05");

    #[test]
    fn sample_part1() {
        let mut instance = Day05::new();
        assert_eq!("10", instance.part1(&mut SAMPLE_INPUT));
    }

    #[test]
    fn sample_part2() {
        let mut instance = Day05::new();
        assert_eq!("4", instance.part2(&mut SAMPLE_INPUT));
    }
}
<commit_msg>Bithack optimizations, why not.<commit_after>use std::io::Read;

use common;

#[derive(Default)]
pub struct Day05 {}

impl Day05 {
    pub fn new() -> Day05 {
        Default::default()
    }

    fn reduce(mut data: Vec<u8>) -> usize {
        let mut dptr = 0;
        for iptr in 0..data.len() {
            if dptr > 0 && (data[iptr] ^ data[dptr - 1]) == 32 {
                dptr -= 1;
            } else {
                data[dptr] = data[iptr];
                dptr += 1;
            }
        }

        dptr
    }
}

impl common::Solution for Day05 {
    fn part1(&mut self, input: &mut Read) -> String {
        let mut data = Vec::new();
        input.read_to_end(&mut data).expect("Can't read input!");
        common::trim_back(&mut data);

        format!("{}", Day05::reduce(data))
    }

    fn part2(&mut self, input: &mut Read) -> String {
        let mut data = Vec::new();
        input.read_to_end(&mut data).expect("Can't read input!");
        common::trim_back(&mut data);

        let min_len = (b'a'..=b'z').map(|option| data.iter().filter(|&x| !x.eq_ignore_ascii_case(&option)).map(|&x| x).collect())
            .map(|x| Day05::reduce(x))
            .min().unwrap();

        format!("{}", min_len)
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day05::Day05;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/05");

    #[test]
    fn sample_part1() {
        let mut instance = Day05::new();
        assert_eq!("10", instance.part1(&mut SAMPLE_INPUT));
    }

    #[test]
    fn sample_part2() {
        let mut instance = Day05::new();
        assert_eq!("4", instance.part2(&mut SAMPLE_INPUT));
    }
}
<|endoftext|>