<commit_before>use gen::GenError;
use lib::std::io::Write;
use lib::std::io;

pub type GenResult<I> = Result<I, GenError>;

pub trait SerializeFn<I>: Fn(I) -> GenResult<I> {}

impl<I, F: Fn(I) -> GenResult<I>> SerializeFn<I> for F {}

pub trait Skip: Write {
    fn skip(self, s: usize) -> Result<Self, GenError> where Self: Sized;
}

pub struct WriteCounter<W>(W, u64);

impl<W: Write> WriteCounter<W> {
    pub fn new(w: W) -> Self {
        WriteCounter(w, 0)
    }

    pub fn position(&self) -> u64 {
        self.1
    }

    pub fn into_inner(self) -> (W, u64) {
        (self.0, self.1)
    }
}

impl<W: Write> Write for WriteCounter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let written = self.0.write(buf)?;
        self.1 += written as u64;
        Ok(written)
    }

    // Our minimal Write trait has no flush()
    #[cfg(feature = "std")]
    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

macro_rules! try_write(($out:ident, $len:ident, $data:expr) => (
    match $out.write($data) {
        Err(io)           => Err(GenError::IoError(io)),
        Ok(n) if n < $len => Err(GenError::BufferTooSmall($len - n)),
        Ok(_)             => Ok($out)
    }
));

/// writes a byte slice to the output
///
/// ```rust
/// use cookie_factory::slice;
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = slice(&b"abcd"[..])(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn slice<S: AsRef<[u8]>, W: Write>(data: S) -> impl SerializeFn<W> {
    let len = data.as_ref().len();

    move |mut out: W| {
        try_write!(out, len, data.as_ref())
    }
}

/// writes a byte slice to the output
///
/// ```rust
/// use cookie_factory::string;
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = string("abcd")(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn string<S: AsRef<str>, W: Write>(data: S) -> impl SerializeFn<W> {
    let len = data.as_ref().len();

    move |mut out: W| {
        try_write!(out, len, data.as_ref().as_bytes())
    }
}

/// writes an hex string to the output
#[cfg(feature = "std")]
pub fn hex<S: ::lib::std::fmt::UpperHex, W: Write>(data: S) -> impl SerializeFn<W> {
  move |mut out: W| {
    match write!(out, "{:X}", data) {
      Err(io) => Err(GenError::IoError(io)),
      Ok(())  => Ok(out)
    }
  }
}

/// skips over some input bytes
///
/// ```rust
/// use cookie_factory::skip;
///
/// let mut buf = [0u8; 100];
///
/// let out = skip(2)(&mut buf[..]).unwrap();
///
/// assert_eq!(out.len(), 98);
/// ```
pub fn skip<W: Write + Skip>(len: usize) -> impl SerializeFn<W> {
    move |out: W| {
        out.skip(len)
    }
}

/// applies 2 serializers in sequence
///
/// ```rust
/// use cookie_factory::{pair, string};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = pair(string("abcd"), string("efgh"))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 8);
/// }
///
/// assert_eq!(&buf[..8], &b"abcdefgh"[..]);
/// ```
pub fn pair<F, G, I: Write>(first: F, second: G) -> impl SerializeFn<I>
where F: SerializeFn<I>,
      G: SerializeFn<I> {

  move |out: I| {
    first(out).and_then(&second)
  }
}

/// applies a serializer if the condition is true
///
/// ```rust
/// use cookie_factory::{cond, string};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = cond(true, string("abcd"))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn cond<F, I: Write>(condition: bool, f: F) -> impl SerializeFn<I>
where F: SerializeFn<I>, {

  move |out: I| {
    if condition {
      f(out)
    } else {
      Ok(out)
    }
  }
}

/// applies an iterator of serializers
///
/// ```rust
/// use cookie_factory::{all, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = all(data.iter().map(string))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 12);
/// }
///
/// assert_eq!(&buf[..12], &b"abcdefghijkl"[..]);
/// ```
pub fn all<G, I: Write, It>(values: It) -> impl SerializeFn<I>
  where G: SerializeFn<I>,
        It: Clone + Iterator<Item=G> {

  move |mut out: I| {
    let it = values.clone();

    for v in it {
      out = v(out)?;
    }

    Ok(out)
  }
}

/// applies an iterator of serializers
///
/// ```rust
/// use cookie_factory::{separated_list, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = separated_list(string(","), data.iter().map(string))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 14);
/// }
///
/// assert_eq!(&buf[..14], &b"abcd,efgh,ijkl"[..]);
/// ```
pub fn separated_list<F, G, I: Write, It>(sep: F, values: It) -> impl SerializeFn<I>
  where F: SerializeFn<I>,
        G: SerializeFn<I>,
        It: Clone + Iterator<Item=G> {

  move |mut out: I| {
    let mut it = values.clone();

    match it.next() {
      None => return Ok(out),
      Some(first) => {
          out = first(out)?;
      }
    }

    for v in it {
      out = sep(out).and_then(v)?;
    }

    Ok(out)
  }
}

pub fn be_u8<W: Write>(i: u8) -> impl SerializeFn<W> {
   let len = 1;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u16<W: Write>(i: u16) -> impl SerializeFn<W> {
   let len = 2;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u24<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 3;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[1..])
    }
}

pub fn be_u32<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 4;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u64<W: Write>(i: u64) -> impl SerializeFn<W> {
   let len = 8;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_i8<W: Write>(i: i8) -> impl SerializeFn<W> {
    be_u8(i as u8)
}

pub fn be_i16<W: Write>(i: i16) -> impl SerializeFn<W> {
    be_u16(i as u16)
}

pub fn be_i24<W: Write>(i: i32) -> impl SerializeFn<W> {
    be_u24(i as u32)
}

pub fn be_i32<W: Write>(i: i32) -> impl SerializeFn<W> {
    be_u32(i as u32)
}

pub fn be_i64<W: Write>(i: i64) -> impl SerializeFn<W> {
    be_u64(i as u64)
}

pub fn be_f32<W: Write>(i: f32) -> impl SerializeFn<W> {
    be_u32(i.to_bits())
}

pub fn be_f64<W: Write>(i: f64) -> impl SerializeFn<W> {
    be_u64(i.to_bits())
}

pub fn le_u8<W: Write>(i: u8) -> impl SerializeFn<W> {
   let len = 1;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u16<W: Write>(i: u16) -> impl SerializeFn<W> {
   let len = 2;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u24<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 3;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[1..])
    }
}

pub fn le_u32<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 4;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u64<W: Write>(i: u64) -> impl SerializeFn<W> {
   let len = 8;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_i8<W: Write>(i: i8) -> impl SerializeFn<W> {
    le_u8(i as u8)
}

pub fn le_i16<W: Write>(i: i16) -> impl SerializeFn<W> {
    le_u16(i as u16)
}

pub fn le_i24<W: Write>(i: i32) -> impl SerializeFn<W> {
    le_u24(i as u32)
}

pub fn le_i32<W: Write>(i: i32) -> impl SerializeFn<W> {
    le_u32(i as u32)
}

pub fn le_i64<W: Write>(i: i64) -> impl SerializeFn<W> {
    le_u64(i as u64)
}

pub fn le_f32<W: Write>(i: f32) -> impl SerializeFn<W> {
    le_u32(i.to_bits())
}

pub fn le_f64<W: Write>(i: f64) -> impl SerializeFn<W> {
    le_u64(i.to_bits())
}

/// applies a generator over an iterator of values, and applies the serializers generated
///
/// ```rust
/// use cookie_factory::{many_ref, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = many_ref(&data, string)(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 12);
/// }
///
/// assert_eq!(&buf[..12], &b"abcdefghijkl"[..]);
/// ```
pub fn many_ref<E, It, I, F, G, O: Write>(items: I, generator: F) -> impl SerializeFn<O>
where
    It: Iterator<Item = E> + Clone,
    I: IntoIterator<Item = E, IntoIter = It>,
    F: Fn(E) -> G,
    G: SerializeFn<O>
{
    let items = items.into_iter();
    move |mut out: O| {
        for item in items.clone() {
            out = generator(item)(out)?;
        }
        Ok(out)
    }
}

//missing combinators:
//or
//empty
//then
//stream
//length_value
//text print
//text upperhex
//text lowerhex

impl Skip for &mut [u8] {
    fn skip(self, len: usize) -> Result<Self, GenError> {
        if self.len() < len {
            Err(GenError::BufferTooSmall(len - self.len()))
        } else {
            Ok(&mut self[len..])
        }
    }
}

impl<'a> Skip for io::Cursor<&'a mut [u8]> {
    fn skip(mut self, len: usize) -> Result<Self, GenError> {
        let remaining = self.get_ref().len().saturating_sub(self.position() as usize);
        if remaining < len {
            Err(GenError::BufferTooSmall(len - remaining))
        } else {
            self.set_position(self.position() + len as u64);
            Ok(self)
        }
    }
}

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

    #[test]
    fn test_pair_with_cursor() {
        let mut buf = [0u8; 8];

        {
            use self::io::Cursor;

            let cursor = Cursor::new(&mut buf[..]);
            let serializer = pair(
                string("1234"),
                string("5678"),
            );

            let cursor = serializer(cursor).unwrap();
            assert_eq!(cursor.position(), 8);
        }

        assert_eq!(&buf[..], b"12345678");
    }

    #[test]
    fn test_gen_with_length() {
        let mut buf = [0; 8];
        {
            let (len_buf, buf) = buf.split_at_mut(4);
            let w = WriteCounter::new(buf);
            let w = string("test")(w).unwrap();
            be_u32(w.position() as u32)(len_buf).unwrap();
        }
        assert_eq!(&buf, &[0, 0, 0, 4, 't' as u8, 'e' as u8, 's' as u8, 't' as u8]);
    }
}
<commit_msg>Fix implementation of le_u24()<commit_after>use gen::GenError;
use lib::std::io::Write;
use lib::std::io;

pub type GenResult<I> = Result<I, GenError>;

pub trait SerializeFn<I>: Fn(I) -> GenResult<I> {}

impl<I, F: Fn(I) -> GenResult<I>> SerializeFn<I> for F {}

pub trait Skip: Write {
    fn skip(self, s: usize) -> Result<Self, GenError> where Self: Sized;
}

pub struct WriteCounter<W>(W, u64);

impl<W: Write> WriteCounter<W> {
    pub fn new(w: W) -> Self {
        WriteCounter(w, 0)
    }

    pub fn position(&self) -> u64 {
        self.1
    }

    pub fn into_inner(self) -> (W, u64) {
        (self.0, self.1)
    }
}

impl<W: Write> Write for WriteCounter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let written = self.0.write(buf)?;
        self.1 += written as u64;
        Ok(written)
    }

    // Our minimal Write trait has no flush()
    #[cfg(feature = "std")]
    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

macro_rules! try_write(($out:ident, $len:ident, $data:expr) => (
    match $out.write($data) {
        Err(io)           => Err(GenError::IoError(io)),
        Ok(n) if n < $len => Err(GenError::BufferTooSmall($len - n)),
        Ok(_)             => Ok($out)
    }
));

/// writes a byte slice to the output
///
/// ```rust
/// use cookie_factory::slice;
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = slice(&b"abcd"[..])(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn slice<S: AsRef<[u8]>, W: Write>(data: S) -> impl SerializeFn<W> {
    let len = data.as_ref().len();

    move |mut out: W| {
        try_write!(out, len, data.as_ref())
    }
}

/// writes a byte slice to the output
///
/// ```rust
/// use cookie_factory::string;
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = string("abcd")(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn string<S: AsRef<str>, W: Write>(data: S) -> impl SerializeFn<W> {
    let len = data.as_ref().len();

    move |mut out: W| {
        try_write!(out, len, data.as_ref().as_bytes())
    }
}

/// writes an hex string to the output
#[cfg(feature = "std")]
pub fn hex<S: ::lib::std::fmt::UpperHex, W: Write>(data: S) -> impl SerializeFn<W> {
  move |mut out: W| {
    match write!(out, "{:X}", data) {
      Err(io) => Err(GenError::IoError(io)),
      Ok(())  => Ok(out)
    }
  }
}

/// skips over some input bytes
///
/// ```rust
/// use cookie_factory::skip;
///
/// let mut buf = [0u8; 100];
///
/// let out = skip(2)(&mut buf[..]).unwrap();
///
/// assert_eq!(out.len(), 98);
/// ```
pub fn skip<W: Write + Skip>(len: usize) -> impl SerializeFn<W> {
    move |out: W| {
        out.skip(len)
    }
}

/// applies 2 serializers in sequence
///
/// ```rust
/// use cookie_factory::{pair, string};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = pair(string("abcd"), string("efgh"))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 8);
/// }
///
/// assert_eq!(&buf[..8], &b"abcdefgh"[..]);
/// ```
pub fn pair<F, G, I: Write>(first: F, second: G) -> impl SerializeFn<I>
where F: SerializeFn<I>,
      G: SerializeFn<I> {

  move |out: I| {
    first(out).and_then(&second)
  }
}

/// applies a serializer if the condition is true
///
/// ```rust
/// use cookie_factory::{cond, string};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = cond(true, string("abcd"))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn cond<F, I: Write>(condition: bool, f: F) -> impl SerializeFn<I>
where F: SerializeFn<I>, {

  move |out: I| {
    if condition {
      f(out)
    } else {
      Ok(out)
    }
  }
}

/// applies an iterator of serializers
///
/// ```rust
/// use cookie_factory::{all, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = all(data.iter().map(string))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 12);
/// }
///
/// assert_eq!(&buf[..12], &b"abcdefghijkl"[..]);
/// ```
pub fn all<G, I: Write, It>(values: It) -> impl SerializeFn<I>
  where G: SerializeFn<I>,
        It: Clone + Iterator<Item=G> {

  move |mut out: I| {
    let it = values.clone();

    for v in it {
      out = v(out)?;
    }

    Ok(out)
  }
}

/// applies an iterator of serializers
///
/// ```rust
/// use cookie_factory::{separated_list, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = separated_list(string(","), data.iter().map(string))(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 14);
/// }
///
/// assert_eq!(&buf[..14], &b"abcd,efgh,ijkl"[..]);
/// ```
pub fn separated_list<F, G, I: Write, It>(sep: F, values: It) -> impl SerializeFn<I>
  where F: SerializeFn<I>,
        G: SerializeFn<I>,
        It: Clone + Iterator<Item=G> {

  move |mut out: I| {
    let mut it = values.clone();

    match it.next() {
      None => return Ok(out),
      Some(first) => {
          out = first(out)?;
      }
    }

    for v in it {
      out = sep(out).and_then(v)?;
    }

    Ok(out)
  }
}

pub fn be_u8<W: Write>(i: u8) -> impl SerializeFn<W> {
   let len = 1;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u16<W: Write>(i: u16) -> impl SerializeFn<W> {
   let len = 2;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u24<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 3;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[1..])
    }
}

pub fn be_u32<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 4;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_u64<W: Write>(i: u64) -> impl SerializeFn<W> {
   let len = 8;

    move |mut out: W| {
        try_write!(out, len, &i.to_be_bytes()[..])
    }
}

pub fn be_i8<W: Write>(i: i8) -> impl SerializeFn<W> {
    be_u8(i as u8)
}

pub fn be_i16<W: Write>(i: i16) -> impl SerializeFn<W> {
    be_u16(i as u16)
}

pub fn be_i24<W: Write>(i: i32) -> impl SerializeFn<W> {
    be_u24(i as u32)
}

pub fn be_i32<W: Write>(i: i32) -> impl SerializeFn<W> {
    be_u32(i as u32)
}

pub fn be_i64<W: Write>(i: i64) -> impl SerializeFn<W> {
    be_u64(i as u64)
}

pub fn be_f32<W: Write>(i: f32) -> impl SerializeFn<W> {
    be_u32(i.to_bits())
}

pub fn be_f64<W: Write>(i: f64) -> impl SerializeFn<W> {
    be_u64(i.to_bits())
}

pub fn le_u8<W: Write>(i: u8) -> impl SerializeFn<W> {
   let len = 1;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u16<W: Write>(i: u16) -> impl SerializeFn<W> {
   let len = 2;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u24<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 3;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[0..3])
    }
}

pub fn le_u32<W: Write>(i: u32) -> impl SerializeFn<W> {
   let len = 4;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_u64<W: Write>(i: u64) -> impl SerializeFn<W> {
   let len = 8;

    move |mut out: W| {
        try_write!(out, len, &i.to_le_bytes()[..])
    }
}

pub fn le_i8<W: Write>(i: i8) -> impl SerializeFn<W> {
    le_u8(i as u8)
}

pub fn le_i16<W: Write>(i: i16) -> impl SerializeFn<W> {
    le_u16(i as u16)
}

pub fn le_i24<W: Write>(i: i32) -> impl SerializeFn<W> {
    le_u24(i as u32)
}

pub fn le_i32<W: Write>(i: i32) -> impl SerializeFn<W> {
    le_u32(i as u32)
}

pub fn le_i64<W: Write>(i: i64) -> impl SerializeFn<W> {
    le_u64(i as u64)
}

pub fn le_f32<W: Write>(i: f32) -> impl SerializeFn<W> {
    le_u32(i.to_bits())
}

pub fn le_f64<W: Write>(i: f64) -> impl SerializeFn<W> {
    le_u64(i.to_bits())
}

/// applies a generator over an iterator of values, and applies the serializers generated
///
/// ```rust
/// use cookie_factory::{many_ref, string};
///
/// let mut buf = [0u8; 100];
///
/// let data = vec!["abcd", "efgh", "ijkl"];
/// {
///   let buf = many_ref(&data, string)(&mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 12);
/// }
///
/// assert_eq!(&buf[..12], &b"abcdefghijkl"[..]);
/// ```
pub fn many_ref<E, It, I, F, G, O: Write>(items: I, generator: F) -> impl SerializeFn<O>
where
    It: Iterator<Item = E> + Clone,
    I: IntoIterator<Item = E, IntoIter = It>,
    F: Fn(E) -> G,
    G: SerializeFn<O>
{
    let items = items.into_iter();
    move |mut out: O| {
        for item in items.clone() {
            out = generator(item)(out)?;
        }
        Ok(out)
    }
}

//missing combinators:
//or
//empty
//then
//stream
//length_value
//text print
//text upperhex
//text lowerhex

impl Skip for &mut [u8] {
    fn skip(self, len: usize) -> Result<Self, GenError> {
        if self.len() < len {
            Err(GenError::BufferTooSmall(len - self.len()))
        } else {
            Ok(&mut self[len..])
        }
    }
}

impl<'a> Skip for io::Cursor<&'a mut [u8]> {
    fn skip(mut self, len: usize) -> Result<Self, GenError> {
        let remaining = self.get_ref().len().saturating_sub(self.position() as usize);
        if remaining < len {
            Err(GenError::BufferTooSmall(len - remaining))
        } else {
            self.set_position(self.position() + len as u64);
            Ok(self)
        }
    }
}

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

    #[test]
    fn test_pair_with_cursor() {
        let mut buf = [0u8; 8];

        {
            use self::io::Cursor;

            let cursor = Cursor::new(&mut buf[..]);
            let serializer = pair(
                string("1234"),
                string("5678"),
            );

            let cursor = serializer(cursor).unwrap();
            assert_eq!(cursor.position(), 8);
        }

        assert_eq!(&buf[..], b"12345678");
    }

    #[test]
    fn test_gen_with_length() {
        let mut buf = [0; 8];
        {
            let (len_buf, buf) = buf.split_at_mut(4);
            let w = WriteCounter::new(buf);
            let w = string("test")(w).unwrap();
            be_u32(w.position() as u32)(len_buf).unwrap();
        }
        assert_eq!(&buf, &[0, 0, 0, 4, 't' as u8, 'e' as u8, 's' as u8, 't' as u8]);
    }
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

//! Wire format traits.
//!
//! This module provides [`FromWire`] and [`ToWire`], a pair of traits similar
//! to the core traits in the [`serde`] library. Rather than representing a
//! generically serializeable type, they represent types that can be converted
//! to and from Cerberus's wire format, which has a unique, ad-hoc data model.
//!
//! [`FromWire`]: trait.FromWire.html
//! [`ToWire`]: trait.ToWire.html
//! [`serde`]: https://serde.rs

use crate::io;
use crate::io::LeInt;
use crate::io::Read;
use crate::io::Write;

/// A type which can be deserialized from the Cerberus wire format.
///
/// The lifetime `'wire` indicates that the type can be deserialized from a
/// buffer of lifetime `'wire`.
pub trait FromWire<'wire>: Sized {
    /// Deserializes a `Self` w of `r`.
    fn from_wire<R: Read<'wire>>(r: R) -> Result<Self, FromWireError>;
}

/// An deserialization error.
#[derive(Clone, Copy, Debug)]
pub enum FromWireError {
    /// Indicates that something went wrong in an `io` operation.
    ///
    /// [`io`]: ../../io/index.html
    Io(io::Error),

    /// Indicates that some field within the request was outside of its
    /// valid range.
    OutOfRange,
}

impl From<io::Error> for FromWireError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}

/// A type which can be serialized into the Cerberus wire format.
pub trait ToWire: Sized {
    /// Serializes `self` into `w`.
    fn to_wire<W: Write>(&self, w: W) -> Result<(), ToWireError>;
}

/// A serializerion error.
#[derive(Clone, Copy, Debug)]
pub enum ToWireError {
    /// Indicates that something went wrong in an [`io`] operation.
    ///
    /// [`io`]: ../../io/index.html
    Io(io::Error),
}

impl From<io::Error> for ToWireError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}

/// Represents a C-like enum that can be converted to and from a wire
/// representation.
///
/// An implementation of this trait can be thought of as an unsigned
/// integer with a limited range: every enum variant can be converted
/// to the wire format and back, though not every value of the wire
/// representation can be converted into an enum variant.
///
/// In particular the following identity must hold for all types T:
/// ```
/// # use manticore::protocol::wire::WireEnum;
/// # fn test<T: WireEnum + Copy + PartialEq + std::fmt::Debug>(x: T) {
/// assert_eq!(T::from_wire_value(T::to_wire_value(x)), Some(x));
/// # }
/// ```
pub trait WireEnum: Sized + Copy {
    /// The unrelying "wire type". This is almost always some kind of
    /// unsigned integer.
    type Wire: LeInt;

    /// Converts `self` into its underlying wire representation.
    fn to_wire_value(self) -> Self::Wire;

    /// Attemts to parse a value of `Self` from the underlying wire
    /// representation.
    fn from_wire_value(wire: Self::Wire) -> Option<Self>;
}

impl<'wire, E> FromWire<'wire> for E
where
    E: WireEnum,
{
    fn from_wire<R: Read<'wire>>(mut r: R) -> Result<Self, FromWireError> {
        let wire = <Self as WireEnum>::Wire::read_from(&mut r)?;
        Self::from_wire_value(wire).ok_or(FromWireError::OutOfRange)
    }
}

impl<E> ToWire for E
where
    E: WireEnum,
{
    fn to_wire<W: Write>(&self, mut w: W) -> Result<(), ToWireError> {
        self.to_wire_value().write_to(&mut w)?;
        Ok(())
    }
}

/// A conveinence macro for generating `WireEnum`-implementing enums.
///
///
/// Syntax is as follows:
/// ```text
/// wire_enum! {
///     /// This is my enum.
///     pub enum MyEnum : u8 {
///         /// Variant `A`.
///         A = 0x00,
///         /// Variant `B`.
///         B = 0x01,
///     }
/// }
/// ```
/// This macro will generate an implementation of `WireEnum<Wire=u8>` for
/// the above enum.
macro_rules! wire_enum {
    ($(#[$meta:meta])* $vis:vis enum $name:ident : $wire:ident {
        $($(#[$meta_variant:meta])* $variant:ident = $value:literal,)*
    }) => {
        $(#[$meta])*
        #[repr($wire)]
        #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
        $vis enum $name {
           $(
               $(#[$meta_variant])*
               $variant = $value,
           )*
        }

        impl $crate::protocol::wire::WireEnum for $name {
            type Wire = $wire;
            fn to_wire_value(self) -> Self::Wire {
                self as $wire
            }
            fn from_wire_value(wire: Self::Wire) -> Option<Self> {
                match wire {
                    $(
                        $value => Some(Self::$variant),
                    )*
                    _ => None,
                }
            }
        }
    }
}
<commit_msg>Add string serialization/deserialization to WireEnum.<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

//! Wire format traits.
//!
//! This module provides [`FromWire`] and [`ToWire`], a pair of traits similar
//! to the core traits in the [`serde`] library. Rather than representing a
//! generically serializeable type, they represent types that can be converted
//! to and from Cerberus's wire format, which has a unique, ad-hoc data model.
//!
//! [`FromWire`]: trait.FromWire.html
//! [`ToWire`]: trait.ToWire.html
//! [`serde`]: https://serde.rs

use crate::io;
use crate::io::LeInt;
use crate::io::Read;
use crate::io::Write;

/// A type which can be deserialized from the Cerberus wire format.
///
/// The lifetime `'wire` indicates that the type can be deserialized from a
/// buffer of lifetime `'wire`.
pub trait FromWire<'wire>: Sized {
    /// Deserializes a `Self` w of `r`.
    fn from_wire<R: Read<'wire>>(r: R) -> Result<Self, FromWireError>;
}

/// An deserialization error.
#[derive(Clone, Copy, Debug)]
pub enum FromWireError {
    /// Indicates that something went wrong in an `io` operation.
    ///
    /// [`io`]: ../../io/index.html
    Io(io::Error),

    /// Indicates that some field within the request was outside of its
    /// valid range.
    OutOfRange,
}

impl From<io::Error> for FromWireError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}

/// A type which can be serialized into the Cerberus wire format.
pub trait ToWire: Sized {
    /// Serializes `self` into `w`.
    fn to_wire<W: Write>(&self, w: W) -> Result<(), ToWireError>;
}

/// A serializerion error.
#[derive(Clone, Copy, Debug)]
pub enum ToWireError {
    /// Indicates that something went wrong in an [`io`] operation.
    ///
    /// [`io`]: ../../io/index.html
    Io(io::Error),
}

impl From<io::Error> for ToWireError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}

/// Represents a C-like enum that can be converted to and from a wire
/// representation as well as to and from a string representation.
///
/// An implementation of this trait can be thought of as an unsigned
/// integer with a limited range: every enum variant can be converted
/// to the wire format and back, though not every value of the wire
/// representation can be converted into an enum variant.
///
/// In particular the following identity must hold for all types T:
/// ```
/// # use manticore::protocol::wire::WireEnum;
/// # fn test<T: WireEnum + Copy + PartialEq + std::fmt::Debug>(x: T) {
/// assert_eq!(T::from_wire_value(T::to_wire_value(x)), Some(x));
/// # }
/// ```
///
/// Also, the following identity must hold for all types T:
/// ```
/// # use manticore::protocol::wire::WireEnum;
/// # fn test<T: WireEnum + Copy + PartialEq + std::fmt::Debug>(x: T) {
/// assert_eq!(T::from_name(T::name(x)), Some(x));
/// # }
/// ```
pub trait WireEnum: Sized + Copy {
    /// The unrelying "wire type". This is almost always some kind of
    /// unsigned integer.
    type Wire: LeInt;

    /// Converts `self` into its underlying wire representation.
    fn to_wire_value(self) -> Self::Wire;

    /// Attempts to parse a value of `Self` from the underlying wire
    /// representation.
    fn from_wire_value(wire: Self::Wire) -> Option<Self>;

    /// Converts `self` into a string representation.
    fn name(self) -> &'static str;

    /// Attempts to convert a value of `Self` from a string representation.
    fn from_name(str: &str) -> Option<Self>;
}

impl<'wire, E> FromWire<'wire> for E
where
    E: WireEnum,
{
    fn from_wire<R: Read<'wire>>(mut r: R) -> Result<Self, FromWireError> {
        let wire = <Self as WireEnum>::Wire::read_from(&mut r)?;
        Self::from_wire_value(wire).ok_or(FromWireError::OutOfRange)
    }
}

impl<E> ToWire for E
where
    E: WireEnum,
{
    fn to_wire<W: Write>(&self, mut w: W) -> Result<(), ToWireError> {
        self.to_wire_value().write_to(&mut w)?;
        Ok(())
    }
}

/// A conveinence macro for generating `WireEnum`-implementing enums.
///
///
/// Syntax is as follows:
/// ```text
/// wire_enum! {
///     /// This is my enum.
///     pub enum MyEnum : u8 {
///         /// Variant `A`.
///         A = 0x00,
///         /// Variant `B`.
///         B = 0x01,
///     }
/// }
/// ```
/// This macro will generate an implementation of `WireEnum<Wire=u8>` for
/// the above enum.
macro_rules! wire_enum {
    ($(#[$meta:meta])* $vis:vis enum $name:ident : $wire:ident {
        $($(#[$meta_variant:meta])* $variant:ident = $value:literal,)*
    }) => {
        $(#[$meta])*
        #[repr($wire)]
        #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
        $vis enum $name {
           $(
               $(#[$meta_variant])*
               $variant = $value,
           )*
        }

        impl $crate::protocol::wire::WireEnum for $name {
            type Wire = $wire;
            fn to_wire_value(self) -> Self::Wire {
                self as $wire
            }
            fn from_wire_value(wire: Self::Wire) -> Option<Self> {
                match wire {
                    $(
                        $value => Some(Self::$variant),
                    )*
                    _ => None,
                }
            }

            fn name(self) -> &'static str {
                match self {
                    $(
                        Self::$variant => stringify!($variant),
                    )*
                }
            }

            fn from_name(name: &str) -> Option<Self> {
                match name {
                    $(
                        stringify!($variant) => Some(Self::$variant),
                    )*
                    _ => None,
                }
            }
        }

        impl core::fmt::Display for $name {
            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
                use $crate::protocol::wire::WireEnum;

                write!(f, "{}", self.name())
            }
        }

        impl core::str::FromStr for $name {
            type Err = ();

            fn from_str(s: &str) -> Result<Self, ()> {
                use $crate::protocol::wire::WireEnum;

                match $name::from_name(s) {
                    Some(val) => Ok(val),
                    None => Err(()),
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    wire_enum! {
        /// An enum for testing.
        #[cfg_attr(feature = "arbitrary-derive", derive(Arbitrary))]
        #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
        pub enum DemoEnum: u8 {
            /// Unknown value
            Unknown = 0x00,

            /// First enum value
            First = 0x01,

            /// Second enum value
            Second = 0x02,
        }
    }

    #[test]
    fn from_name() {
        use crate::protocol::wire::*;

        let value = DemoEnum::from_name("Second").expect("from_name failed");
        assert_eq!(value, DemoEnum::Second);

        let value = DemoEnum::from_name("First").expect("from_name failed");
        assert_eq!(value, DemoEnum::First);

        assert_eq!(None, DemoEnum::from_name("does not exist"));
    }

    #[test]
    fn name() {
        use crate::protocol::wire::*;

        assert_eq!(DemoEnum::First.name(), "First");
        assert_eq!(DemoEnum::Second.name(), "Second");
    }
}
<|endoftext|>


<commit_before>import std::os;
import std::fs;
import std::os_fs;
import std::vec;
import std::map;
import std::str;
import std::uint;
import metadata::cstore;
import driver::session;
import util::filesearch;

// FIXME #721: Despite the compiler warning, test does exist and needs
// to be exported
export get_rpath_flags, test;

fn get_rpath_flags(sess: session::session, out_filename: str) -> [str] {
    // No rpath on windows
    if sess.get_targ_cfg().os == session::os_win32 {
        ret [];
    }

    log "preparing the RPATH!";

    let cwd = os::getcwd();
    let sysroot = sess.filesearch().sysroot();
    let output = out_filename;
    let libs = cstore::get_used_crate_files(sess.get_cstore());
    // We don't currently rpath native libraries, but we know
    // where rustrt is and we know every rust program needs it
    let libs = libs + [get_sysroot_absolute_rt_lib(sess)];

    let target_triple = sess.get_opts().target_triple;
    let rpaths = get_rpaths(cwd, sysroot, output, libs, target_triple);
    rpaths_to_flags(rpaths)
}

fn get_sysroot_absolute_rt_lib(sess: session::session) -> fs::path {
    let path = [sess.filesearch().sysroot()]
        + filesearch::relative_target_lib_path(
            sess.get_opts().target_triple)
        + [os::dylib_filename("rustrt")];
    check vec::is_not_empty(path);
    fs::connect_many(path)
}

fn rpaths_to_flags(rpaths: [str]) -> [str] {
    vec::map({ |rpath| #fmt("-Wl,-rpath,%s",rpath)}, rpaths)
}

fn get_rpaths(cwd: fs::path, sysroot: fs::path,
              output: fs::path, libs: [fs::path],
              target_triple: str) -> [str] {
    log #fmt("cwd: %s", cwd);
    log #fmt("sysroot: %s", sysroot);
    log #fmt("output: %s", output);
    log #fmt("libs:");
    for libpath in libs {
        log #fmt("    %s", libpath);
    }
    log #fmt("target_triple: %s", target_triple);

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(cwd, output, libs);

    // Make backup absolute paths to the libraries. Binaries can
    // be moved as long as the crates they link against don't move.
    let abs_rpaths = get_absolute_rpaths(cwd, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = [get_install_prefix_rpath(target_triple)];

    fn log_rpaths(desc: str, rpaths: [str]) {
        log #fmt("%s rpaths:", desc);
        for rpath in rpaths {
            log #fmt("    %s", rpath);
        }
    }

    log_rpaths("relative", rel_rpaths);
    log_rpaths("absolute", abs_rpaths);
    log_rpaths("fallback", fallback_rpaths);

    let rpaths = rel_rpaths + abs_rpaths + fallback_rpaths;

    // Remove duplicates
    let rpaths = minimize_rpaths(rpaths);
    ret rpaths;
}

fn get_rpaths_relative_to_output(cwd: fs::path,
                                 output: fs::path,
                                 libs: [fs::path]) -> [str] {
    vec::map(bind get_rpath_relative_to_output(cwd, output, _), libs)
}

fn get_rpath_relative_to_output(cwd: fs::path,
                                output: fs::path,
                                lib: fs::path) -> str {
    "$ORIGIN" + fs::path_sep() + get_relative_to(
        get_absolute(cwd, output),
        get_absolute(cwd, lib))
}

// Find the relative path from one file to another
fn get_relative_to(abs1: fs::path, abs2: fs::path) -> fs::path {
    assert fs::path_is_absolute(abs1);
    assert fs::path_is_absolute(abs2);
    log #fmt("finding relative path from %s to %s",
             abs1, abs2);
    let normal1 = fs::normalize(abs1);
    let normal2 = fs::normalize(abs2);
    let split1 = str::split(normal1, os_fs::path_sep as u8);
    let split2 = str::split(normal2, os_fs::path_sep as u8);
    let len1 = vec::len(split1);
    let len2 = vec::len(split2);
    assert len1 > 0u;
    assert len2 > 0u;

    let max_common_path = uint::min(len1, len2) - 1u;
    let start_idx = 0u;
    while start_idx < max_common_path
        && split1[start_idx] == split2[start_idx] {
        start_idx += 1u;
    }

    let path = [];

    for each _ in uint::range(start_idx, len1 - 1u) {
        path += [".."];
    }

    path += vec::slice(split2, start_idx, len2 - 1u);

    if check vec::is_not_empty(path) {
        ret fs::connect_many(path);
    } else {
        ret ".";
    }
}

fn get_absolute_rpaths(cwd: fs::path, libs: [fs::path]) -> [str] {
    vec::map(bind get_absolute_rpath(cwd, _), libs)
}

fn get_absolute_rpath(cwd: fs::path, lib: fs::path) -> str {
    fs::dirname(get_absolute(cwd, lib))
}

fn get_absolute(cwd: fs::path, lib: fs::path) -> fs::path {
    if fs::path_is_absolute(lib) {
        lib
    } else {
        fs::connect(cwd, lib)
    }
}

fn get_install_prefix_rpath(target_triple: str) -> str {
    let install_prefix = #env("CFG_PREFIX");

    if install_prefix == "" {
        fail "rustc compiled without CFG_PREFIX environment variable";
    }

    let path = [install_prefix]
        + filesearch::relative_target_lib_path(target_triple);
    check vec::is_not_empty(path);
    fs::connect_many(path)
}

fn minimize_rpaths(rpaths: [str]) -> [str] {
    let set = map::new_str_hash::<()>();
    let minimized = [];
    for rpath in rpaths {
        if !set.contains_key(rpath) {
            minimized += [rpath];
            set.insert(rpath, ());
        }
    }
    ret minimized;
}

#[cfg(target_os = "linux")]
#[cfg(target_os = "macos")]
#[cfg(test)]
mod test {
    #[test]
    fn test_rpaths_to_flags() {
        let flags = rpaths_to_flags(["path1", "path2"]);
        assert flags == ["-Wl,-rpath,path1", "-Wl,-rpath,path2"];
    }

    #[test]
    fn test_get_absolute1() {
        let cwd = "/dir";
        let lib = "some/path/lib";
        let res = get_absolute(cwd, lib);
        assert res == "/dir/some/path/lib";
    }

    #[test]
    fn test_get_absolute2() {
        let cwd = "/dir";
        let lib = "/some/path/lib";
        let res = get_absolute(cwd, lib);
        assert res == "/some/path/lib";
    }

    #[test]
    fn test_prefix_rpath() {
        let res = get_install_prefix_rpath("triple");
        assert res == #env("CFG_PREFIX") + "/lib/rustc/triple/lib";
    }

    #[test]
    fn test_minimize1() {
        let res = minimize_rpaths(["rpath1", "rpath2", "rpath1"]);
        assert res == ["rpath1", "rpath2"];
    }

    #[test]
    fn test_minimize2() {
        let res = minimize_rpaths(["1a", "2", "2", "1a", "4a",
                                   "1a", "2", "3", "4a", "3"]);
        assert res == ["1a", "2", "4a", "3"];
    }

    #[test]
    fn test_relative_to1() {
        let p1 = "/usr/bin/rustc";
        let p2 = "/usr/lib/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to2() {
        let p1 = "/usr/bin/rustc";
        let p2 = "/usr/bin/../lib/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to3() {
        let p1 = "/usr/bin/whatever/rustc";
        let p2 = "/usr/lib/whatever/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../../lib/whatever";
    }

    #[test]
    fn test_relative_to4() {
        let p1 = "/usr/bin/whatever/../rustc";
        let p2 = "/usr/lib/whatever/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib/whatever";
    }

    #[test]
    fn test_relative_to5() {
        let p1 = "/usr/bin/whatever/../rustc";
        let p2 = "/usr/lib/whatever/../mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to6() {
        let p1 = "/1";
        let p2 = "/2/3";
        let res = get_relative_to(p1, p2);
        assert res == "2";
    }

    #[test]
    fn test_relative_to7() {
        let p1 = "/1/2";
        let p2 = "/3";
        let res = get_relative_to(p1, p2);
        assert res == "..";
    }

    #[test]
    fn test_relative_to8() {
        let p1 = "/home/brian/Dev/rust/build/"
            + "stage2/lib/rustc/i686-unknown-linux-gnu/lib/librustc.so";
        let p2 = "/home/brian/Dev/rust/build/stage2/bin/.."
            + "/lib/rustc/i686-unknown-linux-gnu/lib/libstd.so";
        let res = get_relative_to(p1, p2);
        assert res == ".";
    }

    #[test]
    fn test_rpath_relative() {
        let res = get_rpath_relative_to_output(
            "/usr", "bin/rustc", "lib/libstd.so");
        assert res == "$ORIGIN/../lib";
    }

    #[test]
    fn test_get_absolute_rpath() {
        let res = get_absolute_rpath("/usr", "lib/libstd.so");
        assert res == "/usr/lib";
    }
}
<commit_msg>Mac doesn't seem to support rpath $ORIGIN<commit_after>import std::os;
import std::fs;
import std::os_fs;
import std::vec;
import std::map;
import std::str;
import std::uint;
import metadata::cstore;
import driver::session;
import util::filesearch;

// FIXME #721: Despite the compiler warning, test does exist and needs
// to be exported
export get_rpath_flags, test;

fn get_rpath_flags(sess: session::session, out_filename: str) -> [str] {
    let os = sess.get_targ_cfg().os;

    // No rpath on windows
    if os == session::os_win32 {
        ret [];
    }

    log "preparing the RPATH!";

    let cwd = os::getcwd();
    let sysroot = sess.filesearch().sysroot();
    let output = out_filename;
    let libs = cstore::get_used_crate_files(sess.get_cstore());
    // We don't currently rpath native libraries, but we know
    // where rustrt is and we know every rust program needs it
    let libs = libs + [get_sysroot_absolute_rt_lib(sess)];

    let target_triple = sess.get_opts().target_triple;
    let rpaths = get_rpaths(os, cwd, sysroot, output, libs, target_triple);
    rpaths_to_flags(rpaths)
}

fn get_sysroot_absolute_rt_lib(sess: session::session) -> fs::path {
    let path = [sess.filesearch().sysroot()]
        + filesearch::relative_target_lib_path(
            sess.get_opts().target_triple)
        + [os::dylib_filename("rustrt")];
    check vec::is_not_empty(path);
    fs::connect_many(path)
}

fn rpaths_to_flags(rpaths: [str]) -> [str] {
    vec::map({ |rpath| #fmt("-Wl,-rpath,%s",rpath)}, rpaths)
}

fn get_rpaths(os: session::os, cwd: fs::path, sysroot: fs::path,
              output: fs::path, libs: [fs::path],
              target_triple: str) -> [str] {
    log #fmt("cwd: %s", cwd);
    log #fmt("sysroot: %s", sysroot);
    log #fmt("output: %s", output);
    log #fmt("libs:");
    for libpath in libs {
        log #fmt("    %s", libpath);
    }
    log #fmt("target_triple: %s", target_triple);

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(os, cwd, output, libs);

    // Make backup absolute paths to the libraries. Binaries can
    // be moved as long as the crates they link against don't move.
    let abs_rpaths = get_absolute_rpaths(cwd, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = [get_install_prefix_rpath(target_triple)];

    fn log_rpaths(desc: str, rpaths: [str]) {
        log #fmt("%s rpaths:", desc);
        for rpath in rpaths {
            log #fmt("    %s", rpath);
        }
    }

    log_rpaths("relative", rel_rpaths);
    log_rpaths("absolute", abs_rpaths);
    log_rpaths("fallback", fallback_rpaths);

    let rpaths = rel_rpaths + abs_rpaths + fallback_rpaths;

    // Remove duplicates
    let rpaths = minimize_rpaths(rpaths);
    ret rpaths;
}

fn get_rpaths_relative_to_output(os: session::os,
                                 cwd: fs::path,
                                 output: fs::path,
                                 libs: [fs::path]) -> [str] {
    vec::map(bind get_rpath_relative_to_output(os, cwd, output, _), libs)
}

fn get_rpath_relative_to_output(os: session::os,
                                cwd: fs::path,
                                output: fs::path,
                                lib: fs::path) -> str {
    // Mac doesn't appear to support $ORIGIN
    let prefix = alt os {
        session::os_linux. { "$ORIGIN" + fs::path_sep() }
        session::os_macos. { "" }
    };

    prefix + get_relative_to(
        get_absolute(cwd, output),
        get_absolute(cwd, lib))
}

// Find the relative path from one file to another
fn get_relative_to(abs1: fs::path, abs2: fs::path) -> fs::path {
    assert fs::path_is_absolute(abs1);
    assert fs::path_is_absolute(abs2);
    log #fmt("finding relative path from %s to %s",
             abs1, abs2);
    let normal1 = fs::normalize(abs1);
    let normal2 = fs::normalize(abs2);
    let split1 = str::split(normal1, os_fs::path_sep as u8);
    let split2 = str::split(normal2, os_fs::path_sep as u8);
    let len1 = vec::len(split1);
    let len2 = vec::len(split2);
    assert len1 > 0u;
    assert len2 > 0u;

    let max_common_path = uint::min(len1, len2) - 1u;
    let start_idx = 0u;
    while start_idx < max_common_path
        && split1[start_idx] == split2[start_idx] {
        start_idx += 1u;
    }

    let path = [];

    for each _ in uint::range(start_idx, len1 - 1u) {
        path += [".."];
    }

    path += vec::slice(split2, start_idx, len2 - 1u);

    if check vec::is_not_empty(path) {
        ret fs::connect_many(path);
    } else {
        ret ".";
    }
}

fn get_absolute_rpaths(cwd: fs::path, libs: [fs::path]) -> [str] {
    vec::map(bind get_absolute_rpath(cwd, _), libs)
}

fn get_absolute_rpath(cwd: fs::path, lib: fs::path) -> str {
    fs::dirname(get_absolute(cwd, lib))
}

fn get_absolute(cwd: fs::path, lib: fs::path) -> fs::path {
    if fs::path_is_absolute(lib) {
        lib
    } else {
        fs::connect(cwd, lib)
    }
}

fn get_install_prefix_rpath(target_triple: str) -> str {
    let install_prefix = #env("CFG_PREFIX");

    if install_prefix == "" {
        fail "rustc compiled without CFG_PREFIX environment variable";
    }

    let path = [install_prefix]
        + filesearch::relative_target_lib_path(target_triple);
    check vec::is_not_empty(path);
    fs::connect_many(path)
}

fn minimize_rpaths(rpaths: [str]) -> [str] {
    let set = map::new_str_hash::<()>();
    let minimized = [];
    for rpath in rpaths {
        if !set.contains_key(rpath) {
            minimized += [rpath];
            set.insert(rpath, ());
        }
    }
    ret minimized;
}

#[cfg(target_os = "linux")]
#[cfg(target_os = "macos")]
#[cfg(test)]
mod test {
    #[test]
    fn test_rpaths_to_flags() {
        let flags = rpaths_to_flags(["path1", "path2"]);
        assert flags == ["-Wl,-rpath,path1", "-Wl,-rpath,path2"];
    }

    #[test]
    fn test_get_absolute1() {
        let cwd = "/dir";
        let lib = "some/path/lib";
        let res = get_absolute(cwd, lib);
        assert res == "/dir/some/path/lib";
    }

    #[test]
    fn test_get_absolute2() {
        let cwd = "/dir";
        let lib = "/some/path/lib";
        let res = get_absolute(cwd, lib);
        assert res == "/some/path/lib";
    }

    #[test]
    fn test_prefix_rpath() {
        let res = get_install_prefix_rpath("triple");
        assert res == #env("CFG_PREFIX") + "/lib/rustc/triple/lib";
    }

    #[test]
    fn test_minimize1() {
        let res = minimize_rpaths(["rpath1", "rpath2", "rpath1"]);
        assert res == ["rpath1", "rpath2"];
    }

    #[test]
    fn test_minimize2() {
        let res = minimize_rpaths(["1a", "2", "2", "1a", "4a",
                                   "1a", "2", "3", "4a", "3"]);
        assert res == ["1a", "2", "4a", "3"];
    }

    #[test]
    fn test_relative_to1() {
        let p1 = "/usr/bin/rustc";
        let p2 = "/usr/lib/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to2() {
        let p1 = "/usr/bin/rustc";
        let p2 = "/usr/bin/../lib/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to3() {
        let p1 = "/usr/bin/whatever/rustc";
        let p2 = "/usr/lib/whatever/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../../lib/whatever";
    }

    #[test]
    fn test_relative_to4() {
        let p1 = "/usr/bin/whatever/../rustc";
        let p2 = "/usr/lib/whatever/mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib/whatever";
    }

    #[test]
    fn test_relative_to5() {
        let p1 = "/usr/bin/whatever/../rustc";
        let p2 = "/usr/lib/whatever/../mylib";
        let res = get_relative_to(p1, p2);
        assert res == "../lib";
    }

    #[test]
    fn test_relative_to6() {
        let p1 = "/1";
        let p2 = "/2/3";
        let res = get_relative_to(p1, p2);
        assert res == "2";
    }

    #[test]
    fn test_relative_to7() {
        let p1 = "/1/2";
        let p2 = "/3";
        let res = get_relative_to(p1, p2);
        assert res == "..";
    }

    #[test]
    fn test_relative_to8() {
        let p1 = "/home/brian/Dev/rust/build/"
            + "stage2/lib/rustc/i686-unknown-linux-gnu/lib/librustc.so";
        let p2 = "/home/brian/Dev/rust/build/stage2/bin/.."
            + "/lib/rustc/i686-unknown-linux-gnu/lib/libstd.so";
        let res = get_relative_to(p1, p2);
        assert res == ".";
    }

    #[test]
    #[cfg(target_os = "linux")]
    fn test_rpath_relative() {
        let res = get_rpath_relative_to_output(session::os_linux,
            "/usr", "bin/rustc", "lib/libstd.so");
        assert res == "$ORIGIN/../lib";
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_rpath_relative() {
        let res = get_rpath_relative_to_output(session::os_macos,
            "/usr", "bin/rustc", "lib/libstd.so");
        assert res == "../lib";
    }

    #[test]
    fn test_get_absolute_rpath() {
        let res = get_absolute_rpath("/usr", "lib/libstd.so");
        assert res == "/usr/lib";
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix benchmarks.<commit_after><|endoftext|>


<commit_before>//! `AstIdMap` allows to create stable IDs for "large" syntax nodes like items
//! and macro calls.
//!
//! Specifically, it enumerates all items in a file and uses position of a an
//! item as an ID. That way, id's don't change unless the set of items itself
//! changes.

use std::{
    any::type_name,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
};

use la_arena::{Arena, Idx};
use syntax::{ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};

/// `AstId` points to an AST node in a specific file.
pub struct FileAstId<N: AstNode> {
    raw: ErasedFileAstId,
    _ty: PhantomData<fn() -> N>,
}

impl<N: AstNode> Clone for FileAstId<N> {
    fn clone(&self) -> FileAstId<N> {
        *self
    }
}
impl<N: AstNode> Copy for FileAstId<N> {}

impl<N: AstNode> PartialEq for FileAstId<N> {
    fn eq(&self, other: &Self) -> bool {
        self.raw == other.raw
    }
}
impl<N: AstNode> Eq for FileAstId<N> {}
impl<N: AstNode> Hash for FileAstId<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.raw.hash(hasher);
    }
}

impl<N: AstNode> fmt::Debug for FileAstId<N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "FileAstId::<{}>({})", type_name::<N>(), self.raw.into_raw())
    }
}

impl<N: AstNode> FileAstId<N> {
    // Can't make this a From implementation because of coherence
    pub fn upcast<M: AstNode>(self) -> FileAstId<M>
    where
        N: Into<M>,
    {
        FileAstId { raw: self.raw, _ty: PhantomData }
    }
}

type ErasedFileAstId = Idx<SyntaxNodePtr>;

/// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.
#[derive(Debug, PartialEq, Eq, Default)]
pub struct AstIdMap {
    arena: Arena<SyntaxNodePtr>,
}

impl AstIdMap {
    pub(crate) fn from_source(node: &SyntaxNode) -> AstIdMap {
        assert!(node.parent().is_none());
        let mut res = AstIdMap { arena: Arena::default() };
        // By walking the tree in breadth-first order we make sure that parents
        // get lower ids then children. That is, adding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not change ids of top-level items, which helps caching.
        bdfs(node, |it| match ast::Item::cast(it) {
            Some(module_item) => {
                res.alloc(module_item.syntax());
                true
            }
            None => false,
        });
        res
    }

    pub fn ast_id<N: AstNode>(&self, item: &N) -> FileAstId<N> {
        let raw = self.erased_ast_id(item.syntax());
        FileAstId { raw, _ty: PhantomData }
    }
    fn erased_ast_id(&self, item: &SyntaxNode) -> ErasedFileAstId {
        let ptr = SyntaxNodePtr::new(item);
        match self.arena.iter().find(|(_id, i)| **i == ptr) {
            Some((it, _)) => it,
            None => panic!(
                "Can't find {:?} in AstIdMap:\n{:?}",
                item,
                self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
            ),
        }
    }

    pub fn get<N: AstNode>(&self, id: FileAstId<N>) -> AstPtr<N> {
        self.arena[id.raw].clone().cast::<N>().unwrap()
    }

    fn alloc(&mut self, item: &SyntaxNode) -> ErasedFileAstId {
        self.arena.alloc(SyntaxNodePtr::new(item))
    }
}

/// Walks the subtree in bdfs order, calling `f` for each node. What is bdfs
/// order? It is a mix of breadth-first and depth first orders. Nodes for which
/// `f` returns true are visited breadth-first, all the other nodes are explored
/// depth-first.
///
/// In other words, the size of the bfs queue is bound by the number of "true"
/// nodes.
fn bdfs(node: &SyntaxNode, mut f: impl FnMut(SyntaxNode) -> bool) {
    let mut curr_layer = vec![node.clone()];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            let mut preorder = node.preorder();
            while let Some(event) = preorder.next() {
                match event {
                    syntax::WalkEvent::Enter(node) => {
                        if f(node.clone()) {
                            next_layer.extend(node.children());
                            preorder.skip_subtree();
                        }
                    }
                    syntax::WalkEvent::Leave(_) => {}
                }
            }
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<commit_msg>Record `FileAstId`s for block expressiosn<commit_after>//! `AstIdMap` allows to create stable IDs for "large" syntax nodes like items
//! and macro calls.
//!
//! Specifically, it enumerates all items in a file and uses position of a an
//! item as an ID. That way, id's don't change unless the set of items itself
//! changes.

use std::{
    any::type_name,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
};

use la_arena::{Arena, Idx};
use syntax::{ast, match_ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};

/// `AstId` points to an AST node in a specific file.
pub struct FileAstId<N: AstNode> {
    raw: ErasedFileAstId,
    _ty: PhantomData<fn() -> N>,
}

impl<N: AstNode> Clone for FileAstId<N> {
    fn clone(&self) -> FileAstId<N> {
        *self
    }
}
impl<N: AstNode> Copy for FileAstId<N> {}

impl<N: AstNode> PartialEq for FileAstId<N> {
    fn eq(&self, other: &Self) -> bool {
        self.raw == other.raw
    }
}
impl<N: AstNode> Eq for FileAstId<N> {}
impl<N: AstNode> Hash for FileAstId<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.raw.hash(hasher);
    }
}

impl<N: AstNode> fmt::Debug for FileAstId<N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "FileAstId::<{}>({})", type_name::<N>(), self.raw.into_raw())
    }
}

impl<N: AstNode> FileAstId<N> {
    // Can't make this a From implementation because of coherence
    pub fn upcast<M: AstNode>(self) -> FileAstId<M>
    where
        N: Into<M>,
    {
        FileAstId { raw: self.raw, _ty: PhantomData }
    }
}

type ErasedFileAstId = Idx<SyntaxNodePtr>;

/// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.
#[derive(Debug, PartialEq, Eq, Default)]
pub struct AstIdMap {
    arena: Arena<SyntaxNodePtr>,
}

impl AstIdMap {
    pub(crate) fn from_source(node: &SyntaxNode) -> AstIdMap {
        assert!(node.parent().is_none());
        let mut res = AstIdMap { arena: Arena::default() };
        // By walking the tree in breadth-first order we make sure that parents
        // get lower ids then children. That is, adding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not change ids of top-level items, which helps caching.
        bdfs(node, |it| {
            match_ast! {
                match it {
                    ast::Item(module_item) => {
                        res.alloc(module_item.syntax());
                        true
                    },
                    ast::BlockExpr(block) => {
                        res.alloc(block.syntax());
                        true
                    },
                    _ => false,
                }
            }
        });
        res
    }

    pub fn ast_id<N: AstNode>(&self, item: &N) -> FileAstId<N> {
        let raw = self.erased_ast_id(item.syntax());
        FileAstId { raw, _ty: PhantomData }
    }
    fn erased_ast_id(&self, item: &SyntaxNode) -> ErasedFileAstId {
        let ptr = SyntaxNodePtr::new(item);
        match self.arena.iter().find(|(_id, i)| **i == ptr) {
            Some((it, _)) => it,
            None => panic!(
                "Can't find {:?} in AstIdMap:\n{:?}",
                item,
                self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
            ),
        }
    }

    pub fn get<N: AstNode>(&self, id: FileAstId<N>) -> AstPtr<N> {
        self.arena[id.raw].clone().cast::<N>().unwrap()
    }

    fn alloc(&mut self, item: &SyntaxNode) -> ErasedFileAstId {
        self.arena.alloc(SyntaxNodePtr::new(item))
    }
}

/// Walks the subtree in bdfs order, calling `f` for each node. What is bdfs
/// order? It is a mix of breadth-first and depth first orders. Nodes for which
/// `f` returns true are visited breadth-first, all the other nodes are explored
/// depth-first.
///
/// In other words, the size of the bfs queue is bound by the number of "true"
/// nodes.
fn bdfs(node: &SyntaxNode, mut f: impl FnMut(SyntaxNode) -> bool) {
    let mut curr_layer = vec![node.clone()];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            let mut preorder = node.preorder();
            while let Some(event) = preorder.next() {
                match event {
                    syntax::WalkEvent::Enter(node) => {
                        if f(node.clone()) {
                            next_layer.extend(node.children());
                            preorder.skip_subtree();
                        }
                    }
                    syntax::WalkEvent::Leave(_) => {}
                }
            }
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<|endoftext|>


<commit_before>use hir::Semantics;
use ide_db::{
    base_db::FilePosition,
    defs::Definition,
    helpers::pick_best_token,
    search::{FileReference, ReferenceAccess, SearchScope},
    RootDatabase,
};
use syntax::{ast, match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, WalkEvent, T};

use crate::{display::TryToNav, references, NavigationTarget};

pub struct HighlightedRange {
    pub range: TextRange,
    pub access: Option<ReferenceAccess>,
}

// Feature: Highlight Related
//
// Highlights constructs related to the thing under the cursor:
// - if on an identifier, highlights all references to that identifier in the current file
// - if on an `async` or `await token, highlights all yield points for that async context
// - if on a `return` token, `?` character or `->` return type arrow, highlights all exit points for that context
pub(crate) fn highlight_related(
    sema: &Semantics<RootDatabase>,
    position: FilePosition,
) -> Option<Vec<HighlightedRange>> {
    let _p = profile::span("highlight_related");
    let syntax = sema.parse(position.file_id).syntax().clone();

    let token = pick_best_token(syntax.token_at_offset(position.offset), |kind| match kind {
        T![?] => 2, // prefer `?` when the cursor is sandwiched like `await$0?`
        T![await] | T![async] | T![->] | T![return] => 1,
        _ => 0,
    })?;

    match token.kind() {
        T![?] | T![return] | T![->] => highlight_exit_points(sema, token),
        T![await] | T![async] => highlight_yield_points(token),
        _ => highlight_references(sema, &syntax, position),
    }
}

fn highlight_references(
    sema: &Semantics<RootDatabase>,
    syntax: &SyntaxNode,
    FilePosition { offset, file_id }: FilePosition,
) -> Option<Vec<HighlightedRange>> {
    let def = references::find_def(sema, syntax, offset)?;
    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();

    let declaration = match def {
        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
            Some(NavigationTarget::from_module_to_decl(sema.db, module))
        }
        def => def.try_to_nav(sema.db),
    }
    .filter(|decl| decl.file_id == file_id)
    .and_then(|decl| {
        let range = decl.focus_range?;
        let access = references::decl_access(&def, syntax, range);
        Some(HighlightedRange { range, access })
    });

    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
    let mut res = Vec::with_capacity(file_refs.len() + 1);
    res.extend(declaration);
    res.extend(
        file_refs
            .iter()
            .map(|&FileReference { access, range, .. }| HighlightedRange { range, access }),
    );
    Some(res)
}

fn highlight_exit_points(
    sema: &Semantics<RootDatabase>,
    token: SyntaxToken,
) -> Option<Vec<HighlightedRange>> {
    fn hl(
        sema: &Semantics<RootDatabase>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<HighlightedRange>> {
        let mut highlights = Vec::new();
        let body = body?;
        walk(&body, |node| {
            match_ast! {
                match node {
                    ast::ReturnExpr(expr) => if let Some(token) = expr.return_token() {
                        highlights.push(HighlightedRange {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::TryExpr(try_) => if let Some(token) = try_.question_mark_token() {
                        highlights.push(HighlightedRange {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::Expr(expr) => match expr {
                        ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroCall(_) => {
                            if sema.type_of_expr(&expr).map_or(false, |ty| ty.is_never()) {
                                highlights.push(HighlightedRange {
                                    access: None,
                                    range: expr.syntax().text_range(),
                                });
                            }
                        },
                        ast::Expr::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::Expr::ClosureExpr(_) => return true,
                        _ => (),
                    },
                    ast::Item(__) => return true,
                    // Don't look into const args
                    ast::Path(__) => return true,
                    _ => (),
                }
            }
            false
        });
        let tail = match body {
            ast::Expr::BlockExpr(b) => b.tail_expr(),
            e => Some(e),
        };
        if let Some(tail) = tail {
            highlights.push(HighlightedRange { access: None, range: tail.syntax().text_range() });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(sema, fn_.body().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(sema, closure.body()),
                ast::EffectExpr(effect) => if effect.async_token().is_some() || effect.try_token().is_some() {
                    hl(sema, effect.block_expr().map(ast::Expr::BlockExpr))
                } else {
                    continue;
                },
                _ => continue,
            }
        };
    }
    None
}

fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
    fn hl(
        async_token: Option<SyntaxToken>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<HighlightedRange>> {
        let mut highlights = Vec::new();
        highlights.push(HighlightedRange { access: None, range: async_token?.text_range() });
        if let Some(body) = body {
            walk(&body, |node| {
                match_ast! {
                    match node {
                        ast::AwaitExpr(expr) => if let Some(token) = expr.await_token() {
                            highlights.push(HighlightedRange {
                                access: None,
                                range: token.text_range(),
                            });
                        },
                        // All the following are different contexts so skip them
                        ast::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::ClosureExpr(__) => return true,
                        ast::Item(__) => return true,
                        // Don't look into const args
                        ast::Path(__) => return true,
                        _ => (),
                    }
                }
                false
            });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(fn_.async_token(), fn_.body().map(ast::Expr::BlockExpr)),
                ast::EffectExpr(effect) => hl(effect.async_token(), effect.block_expr().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(closure.async_token(), closure.body()),
                _ => continue,
            }
        };
    }
    None
}

/// Preorder walk the expression node skipping a node's subtrees if the callback returns `true` for the node.
fn walk(expr: &ast::Expr, mut cb: impl FnMut(SyntaxNode) -> bool) {
    let mut preorder = expr.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        if cb(node) {
            preorder.skip_subtree();
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::fixture;

    use super::*;

    fn check(ra_fixture: &str) {
        let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
        let hls = analysis.highlight_related(pos).unwrap().unwrap();

        let mut expected = annotations
            .into_iter()
            .map(|(r, access)| (r.range, (!access.is_empty()).then(|| access)))
            .collect::<Vec<_>>();

        let mut actual = hls
            .into_iter()
            .map(|hl| {
                (
                    hl.range,
                    hl.access.map(|it| {
                        match it {
                            ReferenceAccess::Read => "read",
                            ReferenceAccess::Write => "write",
                        }
                        .to_string()
                    }),
                )
            })
            .collect::<Vec<_>>();
        actual.sort_by_key(|(range, _)| range.start());
        expected.sort_by_key(|(range, _)| range.start());

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_hl_module() {
        check(
            r#"
//- /lib.rs
mod foo$0;
 // ^^^
//- /foo.rs
struct Foo;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_crate_root() {
        check(
            r#"
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_module() {
        check(
            r#"
//- /lib.rs
mod foo;
//- /foo.rs
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_local() {
        check(
            r#"
fn foo() {
    let mut bar = 3;
         // ^^^ write
    bar$0;
 // ^^^ read
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points() {
        check(
            r#"
pub async fn foo() {
 // ^^^^^
    let x = foo()
        .await$0
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points2() {
        check(
            r#"
pub async$0 fn foo() {
 // ^^^^^
    let x = foo()
        .await
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_fn() {
        check(
            r#"
async fn foo() {
    async fn foo2() {
 // ^^^^^
        async fn foo3() {
            0.await
        }
        0.await$0
       // ^^^^^
    }
    0.await
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_async_blocks() {
        check(
            r#"
async fn foo() {
    (async {
  // ^^^^^
        (async {
           0.await
        }).await$0 }
        // ^^^^^
    ).await;
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points() {
        check(
            r#"
fn foo() -> u32 {
    if true {
        return$0 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points2() {
        check(
            r#"
fn foo() ->$0 u32 {
    if true {
        return 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_prefer_ref_over_tail_exit() {
        check(
            r#"
fn foo() -> u32 {
// ^^^
    if true {
        return 0;
    }

    0?;

    foo$0()
 // ^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_never_call_is_exit_point() {
        check(
            r#"
struct Never;
impl Never {
    fn never(self) -> ! { loop {} }
}
macro_rules! never {
    () => { never() }
}
fn never() -> ! { loop {} }
fn foo() ->$0 u32 {
    never();
 // ^^^^^^^
    never!();
 // FIXME sema doesnt give us types for macrocalls

    Never.never();
 // ^^^^^^^^^^^^^

    0
 // ^
}
"#,
        );
    }
}
<commit_msg>Simplify<commit_after>use hir::Semantics;
use ide_db::{
    base_db::FilePosition,
    defs::Definition,
    helpers::pick_best_token,
    search::{FileReference, ReferenceAccess, SearchScope},
    RootDatabase,
};
use syntax::{ast, match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, WalkEvent, T};

use crate::{display::TryToNav, references, NavigationTarget};

pub struct HighlightedRange {
    pub range: TextRange,
    pub access: Option<ReferenceAccess>,
}

// Feature: Highlight Related
//
// Highlights constructs related to the thing under the cursor:
// - if on an identifier, highlights all references to that identifier in the current file
// - if on an `async` or `await token, highlights all yield points for that async context
// - if on a `return` token, `?` character or `->` return type arrow, highlights all exit points for that context
pub(crate) fn highlight_related(
    sema: &Semantics<RootDatabase>,
    position: FilePosition,
) -> Option<Vec<HighlightedRange>> {
    let _p = profile::span("highlight_related");
    let syntax = sema.parse(position.file_id).syntax().clone();

    let token = pick_best_token(syntax.token_at_offset(position.offset), |kind| match kind {
        T![?] => 2, // prefer `?` when the cursor is sandwiched like `await$0?`
        T![await] | T![async] | T![return] | T![->] => 1,
        _ => 0,
    })?;

    match token.kind() {
        T![return] | T![?] | T![->] => highlight_exit_points(sema, token),
        T![await] | T![async] => highlight_yield_points(token),
        _ => highlight_references(sema, &syntax, position),
    }
}

fn highlight_references(
    sema: &Semantics<RootDatabase>,
    syntax: &SyntaxNode,
    FilePosition { offset, file_id }: FilePosition,
) -> Option<Vec<HighlightedRange>> {
    let def = references::find_def(sema, syntax, offset)?;
    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();

    let declaration = match def {
        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
            Some(NavigationTarget::from_module_to_decl(sema.db, module))
        }
        def => def.try_to_nav(sema.db),
    }
    .filter(|decl| decl.file_id == file_id)
    .and_then(|decl| {
        let range = decl.focus_range?;
        let access = references::decl_access(&def, syntax, range);
        Some(HighlightedRange { range, access })
    });

    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
    let mut res = Vec::with_capacity(file_refs.len() + 1);
    res.extend(declaration);
    res.extend(
        file_refs
            .iter()
            .map(|&FileReference { access, range, .. }| HighlightedRange { range, access }),
    );
    Some(res)
}

fn highlight_exit_points(
    sema: &Semantics<RootDatabase>,
    token: SyntaxToken,
) -> Option<Vec<HighlightedRange>> {
    fn hl(
        sema: &Semantics<RootDatabase>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<HighlightedRange>> {
        let mut highlights = Vec::new();
        let body = body?;
        walk(&body, &mut |expr| {
            match expr {
                ast::Expr::ReturnExpr(expr) => {
                    if let Some(token) = expr.return_token() {
                        highlights
                            .push(HighlightedRange { access: None, range: token.text_range() });
                    }
                }
                ast::Expr::TryExpr(try_) => {
                    if let Some(token) = try_.question_mark_token() {
                        highlights
                            .push(HighlightedRange { access: None, range: token.text_range() });
                    }
                }
                ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroCall(_) => {
                    if sema.type_of_expr(&expr).map_or(false, |ty| ty.is_never()) {
                        highlights.push(HighlightedRange {
                            access: None,
                            range: expr.syntax().text_range(),
                        });
                    }
                }
                ast::Expr::EffectExpr(effect) => {
                    return effect.async_token().is_some() || effect.try_token().is_some()
                }
                ast::Expr::ClosureExpr(_) => return true,
                _ => (),
            }
            false
        });
        let tail = match body {
            ast::Expr::BlockExpr(b) => b.tail_expr(),
            e => Some(e),
        };

        if let Some(tail) = tail {
            highlights.push(HighlightedRange { access: None, range: tail.syntax().text_range() })
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(sema, fn_.body().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(sema, closure.body()),
                ast::EffectExpr(effect) => if matches!(effect.effect(), ast::Effect::Async(_) | ast::Effect::Try(_)| ast::Effect::Const(_)) {
                    hl(sema, effect.block_expr().map(ast::Expr::BlockExpr))
                } else {
                    continue;
                },
                _ => continue,
            }
        };
    }
    None
}

fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
    fn hl(
        async_token: Option<SyntaxToken>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<HighlightedRange>> {
        let mut highlights = Vec::new();
        highlights.push(HighlightedRange { access: None, range: async_token?.text_range() });
        if let Some(body) = body {
            walk(&body, &mut |expr| {
                match expr {
                    ast::Expr::AwaitExpr(expr) => {
                        if let Some(token) = expr.await_token() {
                            highlights
                                .push(HighlightedRange { access: None, range: token.text_range() });
                        }
                    }
                    // All the following are different contexts so skip them
                    ast::Expr::EffectExpr(effect) => {
                        return matches!(
                            effect.effect(),
                            ast::Effect::Async(_) | ast::Effect::Try(_) | ast::Effect::Const(_)
                        )
                    }
                    ast::Expr::ClosureExpr(__) => return true,
                    _ => (),
                }
                false
            });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(fn_.async_token(), fn_.body().map(ast::Expr::BlockExpr)),
                ast::EffectExpr(effect) => hl(effect.async_token(), effect.block_expr().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(closure.async_token(), closure.body()),
                _ => continue,
            }
        };
    }
    None
}

/// Preorder walk the expression node skipping a node's subtrees if the callback returns `true` for the node.
fn walk(expr: &ast::Expr, cb: &mut dyn FnMut(ast::Expr) -> bool) {
    let mut preorder = expr.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        match ast::Stmt::cast(node.clone()) {
            Some(ast::Stmt::LetStmt(l)) => {
                if let Some(expr) = l.initializer() {
                    walk(&expr, cb);
                }
            }
            // Don't skip subtree since we want to process the expression behind this next
            Some(ast::Stmt::ExprStmt(_)) => continue,
            // skip inner items which might have their own expressions
            Some(ast::Stmt::Item(_)) => (),
            None => {
                if let Some(expr) = ast::Expr::cast(node) {
                    if !cb(expr) {
                        continue;
                    }
                }
            }
        }
        preorder.skip_subtree();
    }
}

#[cfg(test)]
mod tests {
    use crate::fixture;

    use super::*;

    fn check(ra_fixture: &str) {
        let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
        let hls = analysis.highlight_related(pos).unwrap().unwrap();

        let mut expected = annotations
            .into_iter()
            .map(|(r, access)| (r.range, (!access.is_empty()).then(|| access)))
            .collect::<Vec<_>>();

        let mut actual = hls
            .into_iter()
            .map(|hl| {
                (
                    hl.range,
                    hl.access.map(|it| {
                        match it {
                            ReferenceAccess::Read => "read",
                            ReferenceAccess::Write => "write",
                        }
                        .to_string()
                    }),
                )
            })
            .collect::<Vec<_>>();
        actual.sort_by_key(|(range, _)| range.start());
        expected.sort_by_key(|(range, _)| range.start());

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_hl_module() {
        check(
            r#"
//- /lib.rs
mod foo$0;
 // ^^^
//- /foo.rs
struct Foo;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_crate_root() {
        check(
            r#"
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_module() {
        check(
            r#"
//- /lib.rs
mod foo;
//- /foo.rs
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_local() {
        check(
            r#"
fn foo() {
    let mut bar = 3;
         // ^^^ write
    bar$0;
 // ^^^ read
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points() {
        check(
            r#"
pub async fn foo() {
 // ^^^^^
    let x = foo()
        .await$0
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points2() {
        check(
            r#"
pub async$0 fn foo() {
 // ^^^^^
    let x = foo()
        .await
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_fn() {
        check(
            r#"
async fn foo() {
    async fn foo2() {
 // ^^^^^
        async fn foo3() {
            0.await
        }
        0.await$0
       // ^^^^^
    }
    0.await
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_async_blocks() {
        check(
            r#"
async fn foo() {
    (async {
  // ^^^^^
        (async {
           0.await
        }).await$0 }
        // ^^^^^
    ).await;
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points() {
        check(
            r#"
fn foo() -> u32 {
    if true {
        return$0 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points2() {
        check(
            r#"
fn foo() ->$0 u32 {
    if true {
        return 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_prefer_ref_over_tail_exit() {
        check(
            r#"
fn foo() -> u32 {
// ^^^
    if true {
        return 0;
    }

    0?;

    foo$0()
 // ^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_never_call_is_exit_point() {
        check(
            r#"
struct Never;
impl Never {
    fn never(self) -> ! { loop {} }
}
macro_rules! never {
    () => { never() }
}
fn never() -> ! { loop {} }
fn foo() ->$0 u32 {
    never();
 // ^^^^^^^
    never!();
 // FIXME sema doesn't give us types for macrocalls

    Never.never();
 // ^^^^^^^^^^^^^

    0
 // ^
}
"#,
        );
    }
}
<|endoftext|>