<commit_before>use bit_vec::BitVec;
use fillings::BitsVec;
use num_traits::{Num, NumCast, cast};
use rand::{self, Rng};
use rustc_serialize::{Encodable, Decodable};

use std::u32;

// Prefer this for marking, instead of Option<u32> (as it requires additional byte of memory)
// We could use usize here, but it will consume a great deal of memory. Keeping that aside, even
// the size of the giant human genome is only 70% of this value (~3 billion bases). So, we're good...
const MARKER: u32 = u32::MAX;

fn induced_sort_large<T>(input: &[T], approx_sa: &mut [u32],
                         mut bucket_heads: BitsVec<usize>, type_map: &BitVec)
    where T: Num + NumCast + PartialOrd + Copy
{
    for i in 0..approx_sa.len() {
        let byte = approx_sa[i];
        if byte == MARKER || byte == 0 {
            continue
        }

        let j = (byte - 1) as usize;
        if !type_map.get(j).unwrap() {
            continue    // only the L-types
        }

        let bucket_idx = cast(input[j]).unwrap();
        let bucket_value = bucket_heads.get(bucket_idx);
        approx_sa[bucket_value] = byte - 1;
        bucket_heads.set(bucket_idx, bucket_value + 1);
    }
}

fn induced_sort_small<T>(input: &[T], approx_sa: &mut [u32],
                         mut bucket_tails: BitsVec<usize>, type_map: &BitVec)
    where T: Num + NumCast + PartialOrd + Copy
{
    for i in (0..approx_sa.len()).rev() {
        let byte = approx_sa[i];
        if byte == MARKER || byte == 0 {
            continue
        }

        let j = (byte - 1) as usize;
        if type_map.get(j).unwrap() {
            continue    // only the S-types
        }

        let bucket_idx = cast(input[j]).unwrap();
        let bucket_value = bucket_tails.get(bucket_idx);
        approx_sa[bucket_value] = byte - 1;
        bucket_tails.set(bucket_idx, bucket_value - 1);
    }
}

// Check whether the string between two LMS bytes have the same lengths and same contents
fn is_equal_lms<T>(input: &[T], lms_map: &BitVec, j: usize, k: usize) -> bool
    where T: Num + NumCast + PartialOrd + Copy
{
    let length = input.len();
    if j == length || k == length {
        return false    // null byte
    }

    for i in 0..(length + 1) {
        let first_lms = lms_map.get(i + j).unwrap();
        let second_lms = lms_map.get(i + k).unwrap();
        if first_lms && second_lms && i > 0 {
            return true
        } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
            return false
        }
    }

    false
}

// Increment the counter at an index
pub fn insert<T>(vec: &mut BitsVec<usize>, value: T) -> usize
    where T: Num + NumCast + PartialOrd + Copy
{
    let idx = cast(value).unwrap();
    if vec.len() <= idx {
        vec.extend_with_element(idx + 1, 0);
    }

    let old = vec.get(idx);
    vec.set(idx, old + 1);
    old + 1
}

pub enum Output<T: Num + NumCast + PartialOrd + Copy> {
    BWT(Vec<T>),
    SA(Vec<u32>),
}

pub fn suffix_array(input: Vec<u8>) -> Vec<u32> {
    match suffix_array_or_bwt(input, false) {
        Output::SA(v) => v,
        _ => unreachable!(),
    }
}

pub fn suffix_array_or_bwt<T>(input: Vec<T>, bwt: bool) -> Output<T>
    where T: Num + NumCast + PartialOrd + Copy + Encodable + Decodable
{
    let name = rand::thread_rng().gen_ascii_chars().take(10).collect::<String>();
    suffix_array_(input, 0, &name, bwt)
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
fn suffix_array_<T>(input: Vec<T>, level: usize, name: &str, bwt: bool) -> Output<T>
    where T: Num + NumCast + PartialOrd + Copy + Encodable + Decodable
{
    let length = input.len();
    let length_32 = length as u32;

    let mut type_map = BitVec::from_elem(length + 1, false);    // `false` for S-type and `true` for L-type
    let mut lms_map = BitVec::from_elem(length + 1, false);     // LMS type

    // We'll be adding the frequencies, so input.len() would be the worst case
    // (i.e., same character throughout the string)
    let input_marker = length.next_power_of_two() - 1;
    let input_bits = input_marker.count_ones() as usize;
    let mut bucket_sizes = BitsVec::new(input_bits);      // byte frequency (HashMap will be a killer in recursions)

    lms_map.set(length, true);          // null byte
    type_map.set(length - 1, true);     // should be L-type
    insert(&mut bucket_sizes, input[length - 1]);

    // 1. Group the bytes into S-type or L-type (also mark LMS types)
    for i in (0..length - 1).rev() {
        let prev_type = type_map.get(i + 1).unwrap();
        insert(&mut bucket_sizes, input[i]);

        if input[i] > input[i + 1] ||
           (input[i] == input[i + 1] && prev_type /* L-type */) {
            if !prev_type /* S-type */ {
                lms_map.set(i + 1, true);
            }

            type_map.set(i, true);
        }
    }

    let mut idx = 1;
    let bytes = bucket_sizes.iter().enumerate().filter_map(|(i, c)| {
        if c == 0 { None } else { Some(i) }
    }).collect::<Vec<_>>();

    // BitsVec always requires the max number of bits it should hold. Using the size of `usize`
    // would probably render it useless (as it'd be no different from a vector). So, we get the
    // maximum value from our collection (say, MAX), get its size (MAX::bits) and pass it to BitsVec.
    // This way, we can reduce the memory consumed by more than half.
    let max_byte = bytes[bytes.len() - 1];
    let mut bucket_tails = BitsVec::with_elements(input_bits, max_byte + 1, 0);
    // (bits + 1) would be worst case, since we'll be incrementing the values again in `induced_sort_large`
    let mut bucket_heads = BitsVec::with_elements(input_bits + 1, max_byte + 1, 0);

    // 2. Fill the bucket heads and tails (heads for L-types and tails for S-types)
    let mut j = 0;
    for i in 0..(max_byte + 1) {
        bucket_heads.set(i, idx);
        if i == bytes[j] {
            idx += bucket_sizes.get(bytes[j]);
            j += 1;
        }

        bucket_tails.set(i, idx - 1);
    }

    drop(bytes);
    drop(bucket_sizes);

    // 3. Build the approximate SA for initial induced sorting
    let mut approx_sa = {
        let mut vec = vec![MARKER; length + 1];
        let mut bucket_tails = bucket_tails.clone();
        for (i, byte) in input.iter().enumerate() {
            if !lms_map.get(i).unwrap() {
                continue        // ignore the L and S types (for now)
            }

            let bucket_idx = cast(*byte).unwrap();
            let bucket_value = bucket_tails.get(bucket_idx);
            vec[bucket_value] = i as u32;
            bucket_tails.set(bucket_idx, bucket_value - 1);
        }

        vec[0] = length_32;     // null byte
        vec
    };

    // 4. Induced sort with respect to L & S types (using the buckets)
    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    // 5. Record the indices that share LMS substrings
    let mut label = 0;
    let lms_bytes = {
        // Approx SA is no longer needed (it'll be dropped when it goes out of scope)
        let mut approx_sa = approx_sa;
        let mut last_idx = approx_sa[0] as usize;
        let mut lms_vec = BitsVec::with_elements(input_bits, length + 1, length_32);
        lms_vec.set(last_idx, 0);

        for count in approx_sa.drain(1..) {
            let idx = if count == MARKER { length } else { count as usize };
            if !lms_map.get(idx).unwrap() {
                continue
            }

            if !is_equal_lms(&input, &lms_map, last_idx, idx) {
                label += 1;
            }

            last_idx = idx;
            lms_vec.set(idx, label);
        }

        lms_vec
    };

    drop(lms_map);

    // ... and filter the unnecessary bytes
    let lms_bits = (lms_bytes.len().next_power_of_two() - 1).count_ones() as usize;
    let mut summary_index = BitsVec::new(lms_bits);
    for (i, b) in lms_bytes.iter().enumerate() {
        if b != length_32 {
            summary_index.push(i);
        }
    }

    // 6. Build the final SA
    let mut final_sa = {
        let summary_sa = if label + 1 < summary_index.len() as u32 {
            // recursion (we don't have enough labels - multiple LMS substrings are same)
            let mapped = summary_index.iter().map(|i| lms_bytes.get(i)).collect::<Vec<_>>();
            drop(lms_bytes);
            match suffix_array_(mapped, level + 1, name, bwt) {
                Output::SA(v) => v,
                _ => unreachable!(),
            }
        } else {
            let mut sum_sa = vec![0; summary_index.len() + 1];
            sum_sa[0] = summary_index.len() as u32;
            for (i, val) in summary_index.iter().enumerate() {
                let idx = lms_bytes.get(val) as usize;
                sum_sa[idx + 1] = i as u32;
            }

            drop(lms_bytes);
            sum_sa      // recursion begins to unwind
        };

        let mut bucket_tails = bucket_tails.clone();
        let mut suffix_idx = vec![MARKER; length + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index.get(summary_sa[i] as usize);
            let bucket_idx = cast(input[idx]).unwrap();
            let bucket_value = bucket_tails.get(bucket_idx);
            suffix_idx[bucket_value] = idx as u32;
            bucket_tails.set(bucket_idx, bucket_value - 1);
        }

        suffix_idx[0] = length_32;
        suffix_idx
    };

    // ... and sort it one last time
    induced_sort_large(&input, &mut final_sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut final_sa, bucket_tails, &type_map);

    if level == 0 && bwt {      // peek of memory consumption
        Output::BWT(
            final_sa.into_iter().map(|i| {
                // BWT[i] = S[SA[i] - 1]
                if i == 0 { cast(0).unwrap() } else { input[(i - 1) as usize] }
            }).collect()
        )
    } else {
        Output::SA(final_sa)
    }
}

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

    #[test]
    fn test_suffix_array() {
        let text = b"ATCGAATCGAGAGATCATCGAATCGAGATCATCGAAATCATCGAATCGTC";
        let sa = suffix_array(text.iter().map(|i| *i).collect::<Vec<_>>());

        let mut rotations = (0..text.len()).map(|i| &text[i as usize..]).collect::<Vec<_>>();
        rotations.sort();

        assert_eq!(sa.into_iter().skip(1).map(|i| &text[i as usize..]).collect::<Vec<_>>(),
                   rotations);
    }
}
<commit_msg>Prefer Vec over BitsVec in some of the implementations<commit_after>use bit_vec::BitVec;
use fillings::BitsVec;
use num_traits::{Num, NumCast, cast};
use rand::{self, Rng};
use rustc_serialize::{Encodable, Decodable};

use std::u32;

// Prefer this for marking, instead of Option<u32> (as it requires additional byte of memory)
// We could use usize here, but it will consume a great deal of memory. Keeping that aside, even
// the size of the giant human genome is only 70% of this value (~3 billion bases). So, we're good...
const MARKER: u32 = u32::MAX;

fn induced_sort_large<T>(input: &[T], approx_sa: &mut [u32],
                         mut bucket_heads: Vec<u32>, type_map: &BitVec)
    where T: Num + NumCast + PartialOrd + Copy
{
    for i in 0..approx_sa.len() {
        let byte = approx_sa[i];
        if byte == MARKER || byte == 0 {
            continue
        }

        let j = (byte - 1) as usize;
        if !type_map.get(j).unwrap() {
            continue    // only the L-types
        }

        let bucket_idx: usize = cast(input[j]).unwrap();
        let bucket_value = bucket_heads[bucket_idx];
        approx_sa[bucket_value as usize] = byte - 1;
        bucket_heads[bucket_idx] += 1;
    }
}

fn induced_sort_small<T>(input: &[T], approx_sa: &mut [u32],
                         mut bucket_tails: Vec<u32>, type_map: &BitVec)
    where T: Num + NumCast + PartialOrd + Copy
{
    for i in (0..approx_sa.len()).rev() {
        let byte = approx_sa[i];
        if byte == MARKER || byte == 0 {
            continue
        }

        let j = (byte - 1) as usize;
        if type_map.get(j).unwrap() {
            continue    // only the S-types
        }

        let bucket_idx: usize = cast(input[j]).unwrap();
        let bucket_value = bucket_tails[bucket_idx];
        approx_sa[bucket_value as usize] = byte - 1;
        bucket_tails[bucket_idx] -= 1;
    }
}

// Check whether the string between two LMS bytes have the same lengths and same contents
fn is_equal_lms<T>(input: &[T], lms_map: &BitVec, j: usize, k: usize) -> bool
    where T: Num + NumCast + PartialOrd + Copy
{
    let length = input.len();
    if j == length || k == length {
        return false    // null byte
    }

    for i in 0..(length + 1) {
        let first_lms = lms_map.get(i + j).unwrap();
        let second_lms = lms_map.get(i + k).unwrap();
        if first_lms && second_lms && i > 0 {
            return true
        } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
            return false
        }
    }

    false
}

// Increment the counter at an index
pub fn insert<T>(vec: &mut BitsVec<usize>, value: T) -> usize
    where T: Num + NumCast + PartialOrd + Copy
{
    let idx = cast(value).unwrap();
    if vec.len() <= idx {
        vec.extend_with_element(idx + 1, 0);
    }

    let old = vec.get(idx);
    vec.set(idx, old + 1);
    old + 1
}

pub fn insert_vec<T>(vec: &mut Vec<u32>, value: T) -> u32
    where T: Num + NumCast + PartialOrd + Copy
{
    let idx = cast(value).unwrap();
    if vec.len() <= idx {
        vec.resize(idx + 1, 0);
    }

    vec[idx] += 1;
    vec[idx]
}

pub enum Output<T: Num + NumCast + PartialOrd + Copy> {
    BWT(Vec<T>),
    SA(Vec<u32>),
}

pub fn suffix_array(input: Vec<u8>) -> Vec<u32> {
    match suffix_array_or_bwt(input, false) {
        Output::SA(v) => v,
        _ => unreachable!(),
    }
}

pub fn suffix_array_or_bwt<T>(input: Vec<T>, bwt: bool) -> Output<T>
    where T: Num + NumCast + PartialOrd + Copy + Encodable + Decodable
{
    let name = rand::thread_rng().gen_ascii_chars().take(10).collect::<String>();
    suffix_array_(input, 0, &name, bwt)
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
fn suffix_array_<T>(input: Vec<T>, level: usize, name: &str, bwt: bool) -> Output<T>
    where T: Num + NumCast + PartialOrd + Copy + Encodable + Decodable
{
    let length = input.len();
    let length_32 = length as u32;

    let mut type_map = BitVec::from_elem(length + 1, false);    // `false` for S-type and `true` for L-type
    let mut lms_map = BitVec::from_elem(length + 1, false);     // LMS type

    // We'll be adding the frequencies, so input.len() would be the worst case
    // (i.e., same character throughout the string)
    let input_marker = length.next_power_of_two() - 1;
    let input_bits = input_marker.count_ones() as usize;
    let mut bucket_sizes = Vec::new();      // byte frequency (HashMap will be a killer in recursions)

    lms_map.set(length, true);          // null byte
    type_map.set(length - 1, true);     // should be L-type
    insert_vec(&mut bucket_sizes, input[length - 1]);

    // 1. Group the bytes into S-type or L-type (also mark LMS types)
    for i in (0..length - 1).rev() {
        let prev_type = type_map.get(i + 1).unwrap();
        insert_vec(&mut bucket_sizes, input[i]);

        if input[i] > input[i + 1] ||
           (input[i] == input[i + 1] && prev_type /* L-type */) {
            if !prev_type /* S-type */ {
                lms_map.set(i + 1, true);
            }

            type_map.set(i, true);
        }
    }

    let mut idx = 1;
    let bytes = bucket_sizes.iter().enumerate().filter_map(|(i, c)| {
        if *c == 0 { None } else { Some(i) }
    }).collect::<Vec<_>>();

    let max_byte = bytes[bytes.len() - 1];
    let mut bucket_tails = vec![0u32; max_byte + 1];
    let mut bucket_heads = vec![0u32; max_byte + 1];

    // 2. Fill the bucket heads and tails (heads for L-types and tails for S-types)
    let mut j = 0;
    for i in 0..(max_byte + 1) {
        bucket_heads[i] = idx;
        if i == bytes[j] {
            idx += bucket_sizes[bytes[j]];
            j += 1;
        }

        bucket_tails[i] = idx - 1;
    }

    drop(bytes);
    drop(bucket_sizes);

    // 3. Build the approximate SA for initial induced sorting
    let mut approx_sa = {
        let mut vec = vec![MARKER; length + 1];
        let mut bucket_tails = bucket_tails.clone();
        for (i, byte) in input.iter().enumerate() {
            if !lms_map.get(i).unwrap() {
                continue        // ignore the L and S types (for now)
            }

            let bucket_idx: usize = cast(*byte).unwrap();
            let bucket_value = bucket_tails[bucket_idx];
            vec[bucket_value as usize] = i as u32;
            bucket_tails[bucket_idx] -= 1;
        }

        vec[0] = length_32;     // null byte
        vec
    };

    // 4. Induced sort with respect to L & S types (using the buckets)
    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    // 5. Record the indices that share LMS substrings
    let mut label = 0;
    let lms_bytes = {
        // Approx SA is no longer needed (it'll be dropped when it goes out of scope)
        let mut approx_sa = approx_sa;
        let mut last_idx = approx_sa[0] as usize;
        let mut lms_vec = BitsVec::with_elements(input_bits, length + 1, length_32);
        lms_vec.set(last_idx, 0);

        for count in approx_sa.drain(1..) {
            let idx = if count == MARKER { length } else { count as usize };
            if !lms_map.get(idx).unwrap() {
                continue
            }

            if !is_equal_lms(&input, &lms_map, last_idx, idx) {
                label += 1;
            }

            last_idx = idx;
            lms_vec.set(idx, label);
        }

        lms_vec
    };

    drop(lms_map);

    // ... and filter the unnecessary bytes
    let lms_bits = (lms_bytes.len().next_power_of_two() - 1).count_ones() as usize;
    let mut summary_index = BitsVec::new(lms_bits);
    for (i, b) in lms_bytes.iter().enumerate() {
        if b != length_32 {
            summary_index.push(i);
        }
    }

    // 6. Build the final SA
    let mut final_sa = {
        let summary_sa = if label + 1 < summary_index.len() as u32 {
            // recursion (we don't have enough labels - multiple LMS substrings are same)
            let mapped = summary_index.iter().map(|i| lms_bytes.get(i)).collect::<Vec<_>>();
            drop(lms_bytes);
            match suffix_array_(mapped, level + 1, name, bwt) {
                Output::SA(v) => v,
                _ => unreachable!(),
            }
        } else {
            let mut sum_sa = vec![0; summary_index.len() + 1];
            sum_sa[0] = summary_index.len() as u32;
            for (i, val) in summary_index.iter().enumerate() {
                let idx = lms_bytes.get(val) as usize;
                sum_sa[idx + 1] = i as u32;
            }

            drop(lms_bytes);
            sum_sa      // recursion begins to unwind
        };

        let mut bucket_tails = bucket_tails.clone();
        let mut suffix_idx = vec![MARKER; length + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index.get(summary_sa[i] as usize);
            let bucket_idx: usize = cast(input[idx]).unwrap();
            let bucket_value = bucket_tails[bucket_idx];
            suffix_idx[bucket_value as usize] = idx as u32;
            bucket_tails[bucket_idx] -= 1;
        }

        suffix_idx[0] = length_32;
        suffix_idx
    };

    // ... and sort it one last time
    induced_sort_large(&input, &mut final_sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut final_sa, bucket_tails, &type_map);

    if level == 0 && bwt {      // peek of memory consumption
        Output::BWT(
            final_sa.into_iter().map(|i| {
                // BWT[i] = S[SA[i] - 1]
                if i == 0 { cast(0).unwrap() } else { input[(i - 1) as usize] }
            }).collect()
        )
    } else {
        Output::SA(final_sa)
    }
}

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

    #[test]
    fn test_suffix_array() {
        let text = b"ATCGAATCGAGAGATCATCGAATCGAGATCATCGAAATCATCGAATCGTC";
        let sa = suffix_array(text.iter().map(|i| *i).collect::<Vec<_>>());

        let mut rotations = (0..text.len()).map(|i| &text[i as usize..]).collect::<Vec<_>>();
        rotations.sort();

        assert_eq!(sa.into_iter().skip(1).map(|i| &text[i as usize..]).collect::<Vec<_>>(),
                   rotations);
    }
}
<|endoftext|>


<commit_before>
pub struct Vm {

}
<commit_msg>Remove unnecessary source file<commit_after><|endoftext|>


<commit_before>macro_rules! write_bytes_le {
    ($buf:ident, $n:ident) => {
        let bytes = $n.to_le_bytes();
        &mut $buf[..bytes.len()].copy_from_slice(&bytes);
    };
}

macro_rules! write_bytes_be {
    ($buf:ident, $n:ident) => {
        let bytes = $n.to_be_bytes();
        &mut $buf[..bytes.len()].copy_from_slice(&bytes);
    };
}

/// Writes an unsigned byte at the start of a buffer.
#[inline]
pub fn put_u8(buf: &mut [u8], n: u8) {
    buf[0] = n;
}

/// Converts a `i8` into an unsigned byte and
/// writes it at the start of a buffer.
#[inline]
pub fn put_i8(buf: &mut [u8], n: i8) {
    buf[0] = n as u8;
}

/// Writes 2 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u16l(buf: &mut [u8], n: u16) {
    write_bytes_le!(buf, n);
}

/// Writes 2 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u16b(buf: &mut [u8], n: u16) {
    write_bytes_be!(buf, n);
}

/// Writes 4 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u32l(buf: &mut [u8], n: u32) {
    write_bytes_le!(buf, n);
}

/// Writes 4 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u32b(buf: &mut [u8], n: u32) {
    write_bytes_be!(buf, n);
}

/// Writes 8 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u64l(buf: &mut [u8], n: u64) {
    write_bytes_le!(buf, n);
}

/// Writes 8 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u64b(buf: &mut [u8], n: u64) {
    write_bytes_be!(buf, n);
}

/// Converts an `i16` into 2 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i16l(buf: &mut [u8], n: i16) {
    put_u16l(buf, n as u16);
}

/// Converts an `i16` into 2 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i16b(buf: &mut [u8], n: i16) {
    put_u16b(buf, n as u16);
}

/// Converts an `i32` into 4 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i32l(buf: &mut [u8], n: i32) {
    put_u32l(buf, n as u32);
}

/// Converts an `i32` into 4 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i32b(buf: &mut [u8], n: i32) {
    put_u32b(buf, n as u32);
}

/// Converts an `i64` into 8 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i64l(buf: &mut [u8], n: i64) {
    put_u64l(buf, n as u64);
}

/// Converts an `i64` into 8 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i64b(buf: &mut [u8], n: i64) {
    put_u64b(buf, n as u64);
}

/// Converts a `f32` into 4 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_f32l(buf: &mut [u8], n: f32) {
    write_bytes_le!(buf, n);
}

/// Converts a `f32` into 4 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_f32b(buf: &mut [u8], n: f32) {
    write_bytes_be!(buf, n);
}

/// Converts a `f64` into 8 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_f64l(buf: &mut [u8], n: f64) {
    write_bytes_le!(buf, n);
}

/// Converts a `f64` into 8 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_f64b(buf: &mut [u8], n: f64) {
    write_bytes_be!(buf, n);
}

#[cfg(test)]
mod test {

    use super::*;
    use crate::byteread::*;
    use std::mem::size_of;

    #[test]
    fn put_and_get_u8() {
        let mut buf = [0; 3];
        put_u8(&mut buf, 1);
        assert!(1 == get_u8(&buf));

        put_u8(&mut buf[1..], 2);
        assert!(2 == get_u8(&buf[1..]));

        put_u8(&mut buf[2..], 255);
        assert!(255 == get_u8(&buf[2..]));
    }

    #[test]
    fn put_and_get_i8() {
        let mut buf = [0; 3];
        put_i8(&mut buf, 1);
        assert!(1 == get_i8(&buf));

        put_i8(&mut buf[1..], 2);
        assert!(2 == get_i8(&buf[1..]));

        put_i8(&mut buf[2..], -128);
        assert!(-128 == get_i8(&buf[2..]));
    }

    macro_rules! decl_put_and_get_endian_tests {
        ($($TYPE:ty),*) => {
            $(
                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _be>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE b>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE == [<get_ $TYPE b>](&buf));

                        [<put_ $TYPE b>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE == [<get_ $TYPE b>](&buf[(1 * item_size)..]));

                        [<put_ $TYPE b>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE == [<get_ $TYPE b>](&buf[(2 * item_size)..]));
                    }
                }

                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _l>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE l>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE == [<get_ $TYPE l>](&buf));

                        [<put_ $TYPE l>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE == [<get_ $TYPE l>](&buf[(1 * item_size)..]));

                        [<put_ $TYPE l>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE == [<get_ $TYPE l>](&buf[(2 * item_size)..]));
                    }
                }
            )*
        };
    }

    macro_rules! decl_put_and_get_endian_float_tests {
        ($($TYPE:ty),*) => {
            $(
                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _be>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE b>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE - [<get_ $TYPE b>](&buf) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE b>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE - [<get_ $TYPE b>](&buf[(1 * item_size)..]) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE b>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE - [<get_ $TYPE b>](&buf[(2 * item_size)..]) < ::std::$TYPE::EPSILON);
                    }
                }

                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _l>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE l>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE - [<get_ $TYPE l>](&buf) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE l>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE - [<get_ $TYPE l>](&buf[(1 * item_size)..]) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE l>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE - [<get_ $TYPE l>](&buf[(2 * item_size)..]) < ::std::$TYPE::EPSILON);
                    }
                }
            )*
        };
    }

    decl_put_and_get_endian_tests!(u16, i16, u32, i32, u64, i64);
    decl_put_and_get_endian_float_tests!(f32, f64);
}
<commit_msg>Remove unnecessary mutable borrow<commit_after>macro_rules! write_bytes_le {
    ($buf:ident, $n:ident) => {
        let bytes = $n.to_le_bytes();
        $buf[..bytes.len()].copy_from_slice(&bytes);
    };
}

macro_rules! write_bytes_be {
    ($buf:ident, $n:ident) => {
        let bytes = $n.to_be_bytes();
        $buf[..bytes.len()].copy_from_slice(&bytes);
    };
}

/// Writes an unsigned byte at the start of a buffer.
#[inline]
pub fn put_u8(buf: &mut [u8], n: u8) {
    buf[0] = n;
}

/// Converts a `i8` into an unsigned byte and
/// writes it at the start of a buffer.
#[inline]
pub fn put_i8(buf: &mut [u8], n: i8) {
    buf[0] = n as u8;
}

/// Writes 2 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u16l(buf: &mut [u8], n: u16) {
    write_bytes_le!(buf, n);
}

/// Writes 2 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u16b(buf: &mut [u8], n: u16) {
    write_bytes_be!(buf, n);
}

/// Writes 4 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u32l(buf: &mut [u8], n: u32) {
    write_bytes_le!(buf, n);
}

/// Writes 4 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u32b(buf: &mut [u8], n: u32) {
    write_bytes_be!(buf, n);
}

/// Writes 8 unsigned bytes in a little-endian order at the start of a buffer.
#[inline]
pub fn put_u64l(buf: &mut [u8], n: u64) {
    write_bytes_le!(buf, n);
}

/// Writes 8 unsigned bytes in a big-endian order at the start of a buffer.
#[inline]
pub fn put_u64b(buf: &mut [u8], n: u64) {
    write_bytes_be!(buf, n);
}

/// Converts an `i16` into 2 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i16l(buf: &mut [u8], n: i16) {
    put_u16l(buf, n as u16);
}

/// Converts an `i16` into 2 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i16b(buf: &mut [u8], n: i16) {
    put_u16b(buf, n as u16);
}

/// Converts an `i32` into 4 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i32l(buf: &mut [u8], n: i32) {
    put_u32l(buf, n as u32);
}

/// Converts an `i32` into 4 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i32b(buf: &mut [u8], n: i32) {
    put_u32b(buf, n as u32);
}

/// Converts an `i64` into 8 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_i64l(buf: &mut [u8], n: i64) {
    put_u64l(buf, n as u64);
}

/// Converts an `i64` into 8 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_i64b(buf: &mut [u8], n: i64) {
    put_u64b(buf, n as u64);
}

/// Converts a `f32` into 4 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_f32l(buf: &mut [u8], n: f32) {
    write_bytes_le!(buf, n);
}

/// Converts a `f32` into 4 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_f32b(buf: &mut [u8], n: f32) {
    write_bytes_be!(buf, n);
}

/// Converts a `f64` into 8 unsigned bytes and
/// writes them in a little-endian order at the start of a buffer.
#[inline]
pub fn put_f64l(buf: &mut [u8], n: f64) {
    write_bytes_le!(buf, n);
}

/// Converts a `f64` into 8 unsigned bytes and
/// writes them in a big-endian order at the start of a buffer.
#[inline]
pub fn put_f64b(buf: &mut [u8], n: f64) {
    write_bytes_be!(buf, n);
}

#[cfg(test)]
mod test {

    use super::*;
    use crate::byteread::*;
    use std::mem::size_of;

    #[test]
    fn put_and_get_u8() {
        let mut buf = [0; 3];
        put_u8(&mut buf, 1);
        assert!(1 == get_u8(&buf));

        put_u8(&mut buf[1..], 2);
        assert!(2 == get_u8(&buf[1..]));

        put_u8(&mut buf[2..], 255);
        assert!(255 == get_u8(&buf[2..]));
    }

    #[test]
    fn put_and_get_i8() {
        let mut buf = [0; 3];
        put_i8(&mut buf, 1);
        assert!(1 == get_i8(&buf));

        put_i8(&mut buf[1..], 2);
        assert!(2 == get_i8(&buf[1..]));

        put_i8(&mut buf[2..], -128);
        assert!(-128 == get_i8(&buf[2..]));
    }

    macro_rules! decl_put_and_get_endian_tests {
        ($($TYPE:ty),*) => {
            $(
                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _be>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE b>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE == [<get_ $TYPE b>](&buf));

                        [<put_ $TYPE b>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE == [<get_ $TYPE b>](&buf[(1 * item_size)..]));

                        [<put_ $TYPE b>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE == [<get_ $TYPE b>](&buf[(2 * item_size)..]));
                    }
                }

                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _l>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE l>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE == [<get_ $TYPE l>](&buf));

                        [<put_ $TYPE l>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE == [<get_ $TYPE l>](&buf[(1 * item_size)..]));

                        [<put_ $TYPE l>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE == [<get_ $TYPE l>](&buf[(2 * item_size)..]));
                    }
                }
            )*
        };
    }

    macro_rules! decl_put_and_get_endian_float_tests {
        ($($TYPE:ty),*) => {
            $(
                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _be>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE b>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE - [<get_ $TYPE b>](&buf) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE b>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE - [<get_ $TYPE b>](&buf[(1 * item_size)..]) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE b>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE - [<get_ $TYPE b>](&buf[(2 * item_size)..]) < ::std::$TYPE::EPSILON);
                    }
                }

                paste::item! {
                    #[test]
                    fn [<put_and_get_ $TYPE _l>]() {
                        let item_size = size_of::<$TYPE>();
                        let mut buf = vec![0; 3 * item_size];
                        [<put_ $TYPE l>](&mut buf, 1 as $TYPE);
                        assert!(1 as $TYPE - [<get_ $TYPE l>](&buf) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE l>](&mut buf[(1 * item_size)..], 2 as $TYPE);
                        assert!(2 as $TYPE - [<get_ $TYPE l>](&buf[(1 * item_size)..]) < ::std::$TYPE::EPSILON);

                        [<put_ $TYPE l>](&mut buf[(2 * item_size)..], 255 as $TYPE);
                        assert!(255 as $TYPE - [<get_ $TYPE l>](&buf[(2 * item_size)..]) < ::std::$TYPE::EPSILON);
                    }
                }
            )*
        };
    }

    decl_put_and_get_endian_tests!(u16, i16, u32, i32, u64, i64);
    decl_put_and_get_endian_float_tests!(f32, f64);
}
<|endoftext|>


<commit_before>use base::{AbstractCFType, Boolean, CFAllocatorRef, CFIndex, CFRelease, CFTypeRef};
use base::{kCFAllocatorDefault, kCFAllocatorNull};
use cast::reinterpret_cast;
use libc::c_char;

type CFStringEncoding = u32;

const kCFStringEncodingMacRoman: u32 = 0;
const kCFStringEncodingWindowsLatin1: u32 = 0x0500;
const kCFStringEncodingUTF8: u32 = 0x08000100;

struct __CFString { private: () }
pub type CFStringRef = *__CFString;

pub struct CFString {
    obj: CFStringRef,

    drop {
        unsafe {
            CFRelease(reinterpret_cast(&self.obj));
        }
    }
}

pub impl CFString {
    static fn wrap(obj: CFStringRef) -> CFString {
        CFString { obj: obj }
    }

    static fn new_static(string: &static/str) -> CFString {
        let string_ref = do str::as_buf(string) |bytes, len| {
            CFStringCreateWithBytesNoCopy(kCFAllocatorDefault,
                                          bytes,
                                          len as CFIndex,
                                          kCFStringEncodingUTF8,
                                          false as Boolean,
                                          kCFAllocatorNull)
        };
        CFString::wrap(string_ref)
    }
}

impl CFString : AbstractCFType {
    pure fn as_type_ref(&self) -> CFTypeRef {
        unsafe {
            reinterpret_cast(&self.obj)
        }
    }
}

#[link_args="-framework CoreFoundation"]
#[nolink]
extern {
    fn CFStringCreateWithBytesNoCopy(alloc: CFAllocatorRef,
                                     bytes: *u8,
                                     numBytes: CFIndex,
                                     encoding: CFStringEncoding,
                                     isExternalRepresentation: Boolean,
                                     contentsDeallocator: CFAllocatorRef)
                                  -> CFStringRef;
}

<commit_msg>Add UniChar typedef.<commit_after>use base::{AbstractCFType, Boolean, CFAllocatorRef, CFIndex, CFRelease, CFTypeRef};
use base::{kCFAllocatorDefault, kCFAllocatorNull};
use cast::reinterpret_cast;
use libc::c_char;

pub type UniChar = libc::c_ushort;
pub type CFStringEncoding = u32;

const kCFStringEncodingMacRoman: u32 = 0;
const kCFStringEncodingWindowsLatin1: u32 = 0x0500;
const kCFStringEncodingUTF8: u32 = 0x08000100;

struct __CFString { private: () }
pub type CFStringRef = *__CFString;

pub struct CFString {
    obj: CFStringRef,

    drop {
        unsafe {
            CFRelease(reinterpret_cast(&self.obj));
        }
    }
}

pub impl CFString {
    static fn wrap(obj: CFStringRef) -> CFString {
        CFString { obj: obj }
    }

    static fn new_static(string: &static/str) -> CFString {
        let string_ref = do str::as_buf(string) |bytes, len| {
            CFStringCreateWithBytesNoCopy(kCFAllocatorDefault,
                                          bytes,
                                          len as CFIndex,
                                          kCFStringEncodingUTF8,
                                          false as Boolean,
                                          kCFAllocatorNull)
        };
        CFString::wrap(string_ref)
    }
}

impl CFString : AbstractCFType {
    pure fn as_type_ref(&self) -> CFTypeRef {
        unsafe {
            reinterpret_cast(&self.obj)
        }
    }
}

#[link_args="-framework CoreFoundation"]
#[nolink]
extern {
    fn CFStringCreateWithBytesNoCopy(alloc: CFAllocatorRef,
                                     bytes: *u8,
                                     numBytes: CFIndex,
                                     encoding: CFStringEncoding,
                                     isExternalRepresentation: Boolean,
                                     contentsDeallocator: CFAllocatorRef)
                                  -> CFStringRef;
}

<|endoftext|>


<commit_before>use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::sha1::Sha1;
use futures::{Async, Future, Poll};
use protobuf::{self, Message};
use rand::thread_rng;
use std::io::{self, Read};
use std::marker::PhantomData;
use tokio_io::codec::Framed;
use tokio_io::io::{read_exact, write_all, ReadExact, Window, WriteAll};
use tokio_io::{AsyncRead, AsyncWrite};

use super::codec::APCodec;
use diffie_hellman::DHLocalKeys;
use protocol;
use protocol::keyexchange::{APResponseMessage, ClientHello, ClientResponsePlaintext};
use util;

pub struct Handshake<T> {
    keys: DHLocalKeys,
    state: HandshakeState<T>,
}

enum HandshakeState<T> {
    ClientHello(WriteAll<T, Vec<u8>>),
    APResponse(RecvPacket<T, APResponseMessage>),
    ClientResponse(Option<APCodec>, WriteAll<T, Vec<u8>>),
}

pub fn handshake<T: AsyncRead + AsyncWrite>(connection: T) -> Handshake<T> {
    let local_keys = DHLocalKeys::random(&mut thread_rng());
    let client_hello = client_hello(connection, local_keys.public_key());

    Handshake {
        keys: local_keys,
        state: HandshakeState::ClientHello(client_hello),
    }
}

impl<T: AsyncRead + AsyncWrite> Future for Handshake<T> {
    type Item = Framed<T, APCodec>;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, io::Error> {
        use self::HandshakeState::*;
        loop {
            self.state = match self.state {
                ClientHello(ref mut write) => {
                    let (connection, accumulator) = try_ready!(write.poll());

                    let read = recv_packet(connection, accumulator);
                    APResponse(read)
                }

                APResponse(ref mut read) => {
                    let (connection, message, accumulator) = try_ready!(read.poll());
                    let remote_key = message
                        .get_challenge()
                        .get_login_crypto_challenge()
                        .get_diffie_hellman()
                        .get_gs()
                        .to_owned();

                    let shared_secret = self.keys.shared_secret(&remote_key);
                    let (challenge, send_key, recv_key) = compute_keys(&shared_secret, &accumulator);
                    let codec = APCodec::new(&send_key, &recv_key);

                    let write = client_response(connection, challenge);
                    ClientResponse(Some(codec), write)
                }

                ClientResponse(ref mut codec, ref mut write) => {
                    let (connection, _) = try_ready!(write.poll());
                    let codec = codec.take().unwrap();
                    let framed = connection.framed(codec);
                    return Ok(Async::Ready(framed));
                }
            }
        }
    }
}

fn client_hello<T: AsyncWrite>(connection: T, gc: Vec<u8>) -> WriteAll<T, Vec<u8>> {
    let mut packet = ClientHello::new();
    packet
        .mut_build_info()
        .set_product(protocol::keyexchange::Product::PRODUCT_PARTNER);
    packet
        .mut_build_info()
        .set_platform(protocol::keyexchange::Platform::PLATFORM_LINUX_X86);
    packet.mut_build_info().set_version(0x10800000000);
    packet
        .mut_cryptosuites_supported()
        .push(protocol::keyexchange::Cryptosuite::CRYPTO_SUITE_SHANNON);
    packet.mut_login_crypto_hello().mut_diffie_hellman().set_gc(gc);
    packet
        .mut_login_crypto_hello()
        .mut_diffie_hellman()
        .set_server_keys_known(1);
    packet.set_client_nonce(util::rand_vec(&mut thread_rng(), 0x10));
    packet.set_padding(vec![0x1e]);

    let mut buffer = vec![0, 4];
    let size = 2 + 4 + packet.compute_size();
    buffer.write_u32::<BigEndian>(size).unwrap();
    packet.write_to_vec(&mut buffer).unwrap();

    write_all(connection, buffer)
}

fn client_response<T: AsyncWrite>(connection: T, challenge: Vec<u8>) -> WriteAll<T, Vec<u8>> {
    let mut packet = ClientResponsePlaintext::new();
    packet
        .mut_login_crypto_response()
        .mut_diffie_hellman()
        .set_hmac(challenge);
    packet.mut_pow_response();
    packet.mut_crypto_response();

    let mut buffer = vec![];
    let size = 4 + packet.compute_size();
    buffer.write_u32::<BigEndian>(size).unwrap();
    packet.write_to_vec(&mut buffer).unwrap();

    write_all(connection, buffer)
}

enum RecvPacket<T, M: Message> {
    Header(ReadExact<T, Window<Vec<u8>>>, PhantomData<M>),
    Body(ReadExact<T, Window<Vec<u8>>>, PhantomData<M>),
}

fn recv_packet<T: AsyncRead, M>(connection: T, acc: Vec<u8>) -> RecvPacket<T, M>
where
    T: Read,
    M: Message,
{
    RecvPacket::Header(read_into_accumulator(connection, 4, acc), PhantomData)
}

impl<T: AsyncRead, M> Future for RecvPacket<T, M>
where
    T: Read,
    M: Message,
{
    type Item = (T, M, Vec<u8>);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, io::Error> {
        use self::RecvPacket::*;
        loop {
            *self = match *self {
                Header(ref mut read, _) => {
                    let (connection, header) = try_ready!(read.poll());
                    let size = BigEndian::read_u32(header.as_ref()) as usize;

                    let acc = header.into_inner();
                    let read = read_into_accumulator(connection, size - 4, acc);
                    RecvPacket::Body(read, PhantomData)
                }

                Body(ref mut read, _) => {
                    let (connection, data) = try_ready!(read.poll());
                    let message = protobuf::parse_from_bytes(data.as_ref()).unwrap();

                    let acc = data.into_inner();
                    return Ok(Async::Ready((connection, message, acc)));
                }
            }
        }
    }
}

fn read_into_accumulator<T: AsyncRead>(
    connection: T,
    size: usize,
    mut acc: Vec<u8>,
) -> ReadExact<T, Window<Vec<u8>>> {
    let offset = acc.len();
    acc.resize(offset + size, 0);

    let mut window = Window::new(acc);
    window.set_start(offset);

    read_exact(connection, window)
}

fn compute_keys(shared_secret: &[u8], packets: &[u8]) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
    let mut data = Vec::with_capacity(0x64);
    let mut mac = Hmac::new(Sha1::new(), &shared_secret);

    for i in 1..6 {
        mac.input(packets);
        mac.input(&[i]);
        data.extend_from_slice(&mac.result().code());
        mac.reset();
    }

    mac = Hmac::new(Sha1::new(), &data[..0x14]);
    mac.input(packets);

    (
        mac.result().code().to_vec(),
        data[0x14..0x34].to_vec(),
        data[0x34..0x54].to_vec(),
    )
}
<commit_msg>Updated version number<commit_after>use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::sha1::Sha1;
use futures::{Async, Future, Poll};
use protobuf::{self, Message};
use rand::thread_rng;
use std::io::{self, Read};
use std::marker::PhantomData;
use tokio_io::codec::Framed;
use tokio_io::io::{read_exact, write_all, ReadExact, Window, WriteAll};
use tokio_io::{AsyncRead, AsyncWrite};

use super::codec::APCodec;
use diffie_hellman::DHLocalKeys;
use protocol;
use protocol::keyexchange::{APResponseMessage, ClientHello, ClientResponsePlaintext};
use util;

pub struct Handshake<T> {
    keys: DHLocalKeys,
    state: HandshakeState<T>,
}

enum HandshakeState<T> {
    ClientHello(WriteAll<T, Vec<u8>>),
    APResponse(RecvPacket<T, APResponseMessage>),
    ClientResponse(Option<APCodec>, WriteAll<T, Vec<u8>>),
}

pub fn handshake<T: AsyncRead + AsyncWrite>(connection: T) -> Handshake<T> {
    let local_keys = DHLocalKeys::random(&mut thread_rng());
    let client_hello = client_hello(connection, local_keys.public_key());

    Handshake {
        keys: local_keys,
        state: HandshakeState::ClientHello(client_hello),
    }
}

impl<T: AsyncRead + AsyncWrite> Future for Handshake<T> {
    type Item = Framed<T, APCodec>;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, io::Error> {
        use self::HandshakeState::*;
        loop {
            self.state = match self.state {
                ClientHello(ref mut write) => {
                    let (connection, accumulator) = try_ready!(write.poll());

                    let read = recv_packet(connection, accumulator);
                    APResponse(read)
                }

                APResponse(ref mut read) => {
                    let (connection, message, accumulator) = try_ready!(read.poll());
                    let remote_key = message
                        .get_challenge()
                        .get_login_crypto_challenge()
                        .get_diffie_hellman()
                        .get_gs()
                        .to_owned();

                    let shared_secret = self.keys.shared_secret(&remote_key);
                    let (challenge, send_key, recv_key) = compute_keys(&shared_secret, &accumulator);
                    let codec = APCodec::new(&send_key, &recv_key);

                    let write = client_response(connection, challenge);
                    ClientResponse(Some(codec), write)
                }

                ClientResponse(ref mut codec, ref mut write) => {
                    let (connection, _) = try_ready!(write.poll());
                    let codec = codec.take().unwrap();
                    let framed = connection.framed(codec);
                    return Ok(Async::Ready(framed));
                }
            }
        }
    }
}

fn client_hello<T: AsyncWrite>(connection: T, gc: Vec<u8>) -> WriteAll<T, Vec<u8>> {
    let mut packet = ClientHello::new();
    packet
        .mut_build_info()
        .set_product(protocol::keyexchange::Product::PRODUCT_PARTNER);
    packet
        .mut_build_info()
        .set_platform(protocol::keyexchange::Platform::PLATFORM_LINUX_X86);
    packet.mut_build_info().set_version(109800078);
    packet
        .mut_cryptosuites_supported()
        .push(protocol::keyexchange::Cryptosuite::CRYPTO_SUITE_SHANNON);
    packet.mut_login_crypto_hello().mut_diffie_hellman().set_gc(gc);
    packet
        .mut_login_crypto_hello()
        .mut_diffie_hellman()
        .set_server_keys_known(1);
    packet.set_client_nonce(util::rand_vec(&mut thread_rng(), 0x10));
    packet.set_padding(vec![0x1e]);

    let mut buffer = vec![0, 4];
    let size = 2 + 4 + packet.compute_size();
    buffer.write_u32::<BigEndian>(size).unwrap();
    packet.write_to_vec(&mut buffer).unwrap();

    write_all(connection, buffer)
}

fn client_response<T: AsyncWrite>(connection: T, challenge: Vec<u8>) -> WriteAll<T, Vec<u8>> {
    let mut packet = ClientResponsePlaintext::new();
    packet
        .mut_login_crypto_response()
        .mut_diffie_hellman()
        .set_hmac(challenge);
    packet.mut_pow_response();
    packet.mut_crypto_response();

    let mut buffer = vec![];
    let size = 4 + packet.compute_size();
    buffer.write_u32::<BigEndian>(size).unwrap();
    packet.write_to_vec(&mut buffer).unwrap();

    write_all(connection, buffer)
}

enum RecvPacket<T, M: Message> {
    Header(ReadExact<T, Window<Vec<u8>>>, PhantomData<M>),
    Body(ReadExact<T, Window<Vec<u8>>>, PhantomData<M>),
}

fn recv_packet<T: AsyncRead, M>(connection: T, acc: Vec<u8>) -> RecvPacket<T, M>
where
    T: Read,
    M: Message,
{
    RecvPacket::Header(read_into_accumulator(connection, 4, acc), PhantomData)
}

impl<T: AsyncRead, M> Future for RecvPacket<T, M>
where
    T: Read,
    M: Message,
{
    type Item = (T, M, Vec<u8>);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, io::Error> {
        use self::RecvPacket::*;
        loop {
            *self = match *self {
                Header(ref mut read, _) => {
                    let (connection, header) = try_ready!(read.poll());
                    let size = BigEndian::read_u32(header.as_ref()) as usize;

                    let acc = header.into_inner();
                    let read = read_into_accumulator(connection, size - 4, acc);
                    RecvPacket::Body(read, PhantomData)
                }

                Body(ref mut read, _) => {
                    let (connection, data) = try_ready!(read.poll());
                    let message = protobuf::parse_from_bytes(data.as_ref()).unwrap();

                    let acc = data.into_inner();
                    return Ok(Async::Ready((connection, message, acc)));
                }
            }
        }
    }
}

fn read_into_accumulator<T: AsyncRead>(
    connection: T,
    size: usize,
    mut acc: Vec<u8>,
) -> ReadExact<T, Window<Vec<u8>>> {
    let offset = acc.len();
    acc.resize(offset + size, 0);

    let mut window = Window::new(acc);
    window.set_start(offset);

    read_exact(connection, window)
}

fn compute_keys(shared_secret: &[u8], packets: &[u8]) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
    let mut data = Vec::with_capacity(0x64);
    let mut mac = Hmac::new(Sha1::new(), &shared_secret);

    for i in 1..6 {
        mac.input(packets);
        mac.input(&[i]);
        data.extend_from_slice(&mac.result().code());
        mac.reset();
    }

    mac = Hmac::new(Sha1::new(), &data[..0x14]);
    mac.input(packets);

    (
        mac.result().code().to_vec(),
        data[0x14..0x34].to_vec(),
        data[0x34..0x54].to_vec(),
    )
}
<|endoftext|>


<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath
}

fn main() {
    let req_vec: ~[sched_msg] = ~[];
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    let visitor_count: uint = 0;
    let shared_visitor_count = arc::RWArc::new(visitor_count);

    //Wahoo Scheduling
    let wahoo_index:uint = 0;
    let shared_wahoo_index = arc::RWArc::new(wahoo_index);
    let local_wahoo_index = shared_wahoo_index.clone();

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                    Ok(file_data) => {
                        println(fmt!("begin serving file [%?]", tf.filepath));
                        // A web server should always reply a HTTP header for any legal HTTP request.
                        tf.stream.write("HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n".as_bytes());
                        tf.stream.write(file_data);
                        println(fmt!("finish file [%?]", tf.filepath));
                    }
                    Err(err) => {
                        println(err);
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {
                    // LIFO didn't make sense in service scheduling, so we modify it as FIFO by using shift_opt() rather than pop().
                    let tf_opt: Option<sched_msg> = (*vec).shift_opt();
                    let tf = tf_opt.unwrap();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.

                    //Decrement wahoo index
                    do local_wahoo_index.write |count| {
                        if ((*count) > 0 ) {
                            (*count) -= 1;
                            println(fmt!("wahoo index decremented to %u", (*count)));
                        }
                    }
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    for stream in acceptor.incoming() {
        //Getting the ip address of the socket
        let mut stream = stream;
        let mut incoming_ip: ~str = ~"";
        match stream {
            Some(ref mut theTcpStream) => {
                match theTcpStream.peer_name() {
                    Some(peername) => {
                        incoming_ip = peername.to_str();
                    }
                    None => ()
                }
            }
            None => ()
        }

        //Cells to pass in the strem and ip address to the closure function
        let stream = Cell::new(stream);
        let incoming_ip = Cell::new(incoming_ip);

        let local_visitor_count = shared_visitor_count.clone();
        let local_wahoo_index = shared_wahoo_index.clone();
        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();

        do spawn {
            let mut visitor_count = 0;
            do local_visitor_count.write |count| {
                (*count) += 1;
                visitor_count = (*count);
            }

            let mut stream = stream.take();
            let mut incoming_ip = incoming_ip.take();

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));

                let file_path = ~os::getcwd().push(path.replace("/../", ""));
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", visitor_count);

                    stream.write(response.as_bytes());
                }
                else {
                    // Requests scheduling
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone()};
                    let (sm_port, sm_chan) = std::comm::stream();
                    sm_chan.send(msg);

                    incoming_ip = ~"128.143.";
                    let mut wahoo_index = 0;
                    if checkWahoo(incoming_ip)
                    {
                        println("Found a Charlottesville address");
                        do local_wahoo_index.write |count| {
                            (*count) += 1;
                            wahoo_index = (*count);
                        }
                        println(fmt!("Wahoo requests end at index: %u", wahoo_index));
                    }

                    do child_add_vec.write |vec| {
                        let msg = sm_port.recv();
                        if wahoo_index > 0 {
                            (*vec).insert(wahoo_index-1, msg); //place at back of Charlottesville queue
                            println("inserted at back of Charlottesville queue");
                        }
                        else {
                            (*vec).push(msg); // enqueue new request.
                            println("add to queue");
                        }
                    }
                    child_chan.send(""); //notify the new arriving request.
                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}

fn checkWahoo(ip: ~str) -> bool {
    let WAHOO_IP = ~["128.143.", "137.54."];
    println(fmt!("ip: %s", ip));

    for matcher in WAHOO_IP.iter() {
        if ip.contains(*matcher) {
            return true;
        }
    }
    return false;
}<commit_msg>reducing median latency; finding shortest processing time<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath
}

fn main() {
    let req_vec: ~[sched_msg] = ~[];
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    let visitor_count: uint = 0;
    let shared_visitor_count = arc::RWArc::new(visitor_count);

    //Wahoo Scheduling
    let wahoo_index:uint = 0;
    let shared_wahoo_index = arc::RWArc::new(wahoo_index);
    let local_wahoo_index = shared_wahoo_index.clone();

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                    Ok(file_data) => {
                        println(fmt!("begin serving file [%?]", tf.filepath));
                        // A web server should always reply a HTTP header for any legal HTTP request.
                        tf.stream.write("HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n".as_bytes());
                        tf.stream.write(file_data);
                        println(fmt!("finish file [%?]", tf.filepath));
                    }
                    Err(err) => {
                        println(err);
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {
                    // LIFO didn't make sense in service scheduling, so we modify it as FIFO by using shift_opt() rather than pop().
                    /*
                    let tf_opt: Option<sched_msg> = (*vec).shift_opt();
                    let tf = tf_opt.unwrap();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                    */
                    //Find shortest process
                    let mut wahoo_index = 0;
                    let mut min_index: uint = 0;      let mut min_size: int = 0;
                    do local_wahoo_index.write |count| { wahoo_index = (*count) }

                    //Find shortest Charlottesville request, else do other requests
                    if wahoo_index > 0 {
                        for i in range(0, wahoo_index) {
                            let size =  (*vec)[i].filepath.get_size().unwrap_or_zero();
                            if min_size <= 0 || min_size > size as int{
                                min_size = size as int;
                                min_index = i;
                            }
                        }
                    }
                    else {
                        for i in range(0, (*vec).len()) {
                            let size =  (*vec)[i].filepath.get_size().unwrap_or_zero();
                            if min_size <= 0 || min_size > size as int {
                                min_size = size as int;
                                min_index = i;
                            }
                        }
                    }

                    let tf = (*vec).remove(min_index);
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.

                    //Decrement wahoo index
                    do local_wahoo_index.write |count| {
                        if (wahoo_index > 0) {
                            (*count) -= 1;
                            println(fmt!("wahoo index decremented to %u", (*count)));
                        }
                    }
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    for stream in acceptor.incoming() {
        //Getting the ip address of the socket
        let mut stream = stream;
        let mut incoming_ip: ~str = ~"";
        match stream {
            Some(ref mut theTcpStream) => {
                match theTcpStream.peer_name() {
                    Some(peername) => {
                        incoming_ip = peername.to_str();
                    }
                    None => ()
                }
            }
            None => ()
        }

        //Cells to pass in the strem and ip address to the closure function
        let stream = Cell::new(stream);
        let incoming_ip = Cell::new(incoming_ip);

        let local_visitor_count = shared_visitor_count.clone();
        let local_wahoo_index = shared_wahoo_index.clone();
        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();

        do spawn {
            let mut visitor_count = 0;
            do local_visitor_count.write |count| {
                (*count) += 1;
                visitor_count = (*count);
            }

            let mut stream = stream.take();
            let mut incoming_ip = incoming_ip.take();

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));

                let file_path = ~os::getcwd().push(path.replace("/../", ""));
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", visitor_count);

                    stream.write(response.as_bytes());
                }
                else {
                    // Requests scheduling
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone()};
                    let (sm_port, sm_chan) = std::comm::stream();
                    sm_chan.send(msg);

                    incoming_ip = ~"128.143.";
                    let mut wahoo_index = 0;
                    if checkWahoo(incoming_ip)
                    {
                        println("Found a Charlottesville address");
                        do local_wahoo_index.write |count| {
                            (*count) += 1;
                            wahoo_index = (*count);
                        }
                        println(fmt!("Wahoo requests end at index: %u", wahoo_index));
                    }

                    do child_add_vec.write |vec| {
                        let msg = sm_port.recv();
                        if wahoo_index > 0 {
                            (*vec).insert(wahoo_index-1, msg); //place at back of Charlottesville queue
                            println("inserted at back of Charlottesville queue");
                        }
                        else {
                            (*vec).push(msg); // enqueue new request.
                            println("add to queue");
                        }
                    }
                    child_chan.send(""); //notify the new arriving request.
                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}

fn checkWahoo(ip: ~str) -> bool {
    let WAHOO_IP = ~["128.143.", "137.54."];
    println(fmt!("ip: %s", ip));

    for matcher in WAHOO_IP.iter() {
        if ip.contains(*matcher) {
            return true;
        }
    }
    return false;
}<|endoftext|>


<commit_before>//! Advertizes the capabilities of the LSP Server.
use std::env;

use ide::CompletionResolveCapability;
use lsp_types::{
    CallHierarchyServerCapability, ClientCapabilities, CodeActionKind, CodeActionOptions,
    CodeActionProviderCapability, CodeLensOptions, CompletionOptions,
    DocumentOnTypeFormattingOptions, FileOperationFilter, FileOperationPattern,
    FileOperationPatternKind, FileOperationRegistrationOptions, FoldingRangeProviderCapability,
    HoverProviderCapability, ImplementationProviderCapability, OneOf, RenameOptions, SaveOptions,
    SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
    SemanticTokensOptions, ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability,
    TextDocumentSyncKind, TextDocumentSyncOptions, TypeDefinitionProviderCapability,
    WorkDoneProgressOptions, WorkspaceFileOperationsServerCapabilities,
    WorkspaceServerCapabilities,
};
use rustc_hash::FxHashSet;
use serde_json::json;

use crate::semantic_tokens;

pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabilities {
    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_NO_INCREMENTAL_SYNC").is_ok() {
                TextDocumentSyncKind::Full
            } else {
                TextDocumentSyncKind::Incremental
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default().into()),
        })),
        hover_provider: Some(HoverProviderCapability::Simple(true)),
        completion_provider: Some(CompletionOptions {
            resolve_provider: completions_resolve_provider(client_caps),
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(OneOf::Left(true)),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(OneOf::Left(true)),
        document_highlight_provider: Some(OneOf::Left(true)),
        document_symbol_provider: Some(OneOf::Left(true)),
        workspace_symbol_provider: Some(OneOf::Left(true)),
        code_action_provider: Some(code_action_capabilities(client_caps)),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(OneOf::Left(true)),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(OneOf::Right(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        linked_editing_range_provider: None,
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: Some(WorkspaceServerCapabilities {
            workspace_folders: None,
            file_operations: Some(WorkspaceFileOperationsServerCapabilities {
                did_create: None,
                will_create: None,
                did_rename: None,
                will_rename: Some(FileOperationRegistrationOptions {
                    filters: vec![FileOperationFilter {
                        scheme: Some(String::from("file")),
                        pattern: FileOperationPattern {
                            glob: String::from("**/*.rs"),
                            matches: Some(FileOperationPatternKind::File),
                            options: None,
                        },
                    }],
                }),
                did_delete: None,
                will_delete: None,
            }),
        }),
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider: Some(
            SemanticTokensOptions {
                legend: SemanticTokensLegend {
                    token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                    token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
                },

                full: Some(SemanticTokensFullOptions::Delta { delta: Some(true) }),
                range: Some(true),
                work_done_progress_options: Default::default(),
            }
            .into(),
        ),
        moniker_provider: None,
        experimental: Some(json!({
            "joinLines": true,
            "ssr": true,
            "onEnter": true,
            "parentModule": true,
            "runnables": {
                "kinds": [ "cargo" ],
            },
        })),
    }
}

fn completions_resolve_provider(client_caps: &ClientCapabilities) -> Option<bool> {
    if enabled_completions_resolve_capabilities(client_caps)?.is_empty() {
        log::info!("No `additionalTextEdits` completion resolve capability was found in the client capabilities, autoimport completion is disabled");
        None
    } else {
        Some(true)
    }
}

/// Parses client capabilities and returns all completion resolve capabilities rust-analyzer supports.
pub(crate) fn enabled_completions_resolve_capabilities(
    caps: &ClientCapabilities,
) -> Option<FxHashSet<CompletionResolveCapability>> {
    Some(
        caps.text_document
            .as_ref()?
            .completion
            .as_ref()?
            .completion_item
            .as_ref()?
            .resolve_support
            .as_ref()?
            .properties
            .iter()
            .filter_map(|cap_string| match cap_string.as_str() {
                "additionalTextEdits" => Some(CompletionResolveCapability::AdditionalTextEdits),
                "detail" => Some(CompletionResolveCapability::Detail),
                "documentation" => Some(CompletionResolveCapability::Documentation),
                _unsupported => None,
            })
            .collect(),
    )
}

fn code_action_capabilities(client_caps: &ClientCapabilities) -> CodeActionProviderCapability {
    client_caps
        .text_document
        .as_ref()
        .and_then(|it| it.code_action.as_ref())
        .and_then(|it| it.code_action_literal_support.as_ref())
        .map_or(CodeActionProviderCapability::Simple(true), |_| {
            CodeActionProviderCapability::Options(CodeActionOptions {
                // Advertise support for all built-in CodeActionKinds.
                // Ideally we would base this off of the client capabilities
                // but the client is supposed to fall back gracefully for unknown values.
                code_action_kinds: Some(vec![
                    CodeActionKind::EMPTY,
                    CodeActionKind::QUICKFIX,
                    CodeActionKind::REFACTOR,
                    CodeActionKind::REFACTOR_EXTRACT,
                    CodeActionKind::REFACTOR_INLINE,
                    CodeActionKind::REFACTOR_REWRITE,
                ]),
                resolve_provider: Some(true),
                work_done_progress_options: Default::default(),
            })
        })
}
<commit_msg>support folder rename for initial request.<commit_after>//! Advertizes the capabilities of the LSP Server.
use std::env;

use ide::CompletionResolveCapability;
use lsp_types::{
    CallHierarchyServerCapability, ClientCapabilities, CodeActionKind, CodeActionOptions,
    CodeActionProviderCapability, CodeLensOptions, CompletionOptions,
    DocumentOnTypeFormattingOptions, FileOperationFilter, FileOperationPattern,
    FileOperationPatternKind, FileOperationRegistrationOptions, FoldingRangeProviderCapability,
    HoverProviderCapability, ImplementationProviderCapability, OneOf, RenameOptions, SaveOptions,
    SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
    SemanticTokensOptions, ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability,
    TextDocumentSyncKind, TextDocumentSyncOptions, TypeDefinitionProviderCapability,
    WorkDoneProgressOptions, WorkspaceFileOperationsServerCapabilities,
    WorkspaceServerCapabilities,
};
use rustc_hash::FxHashSet;
use serde_json::json;

use crate::semantic_tokens;

pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabilities {
    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_NO_INCREMENTAL_SYNC").is_ok() {
                TextDocumentSyncKind::Full
            } else {
                TextDocumentSyncKind::Incremental
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default().into()),
        })),
        hover_provider: Some(HoverProviderCapability::Simple(true)),
        completion_provider: Some(CompletionOptions {
            resolve_provider: completions_resolve_provider(client_caps),
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(OneOf::Left(true)),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(OneOf::Left(true)),
        document_highlight_provider: Some(OneOf::Left(true)),
        document_symbol_provider: Some(OneOf::Left(true)),
        workspace_symbol_provider: Some(OneOf::Left(true)),
        code_action_provider: Some(code_action_capabilities(client_caps)),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(OneOf::Left(true)),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(OneOf::Right(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        linked_editing_range_provider: None,
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: Some(WorkspaceServerCapabilities {
            workspace_folders: None,
            file_operations: Some(WorkspaceFileOperationsServerCapabilities {
                did_create: None,
                will_create: None,
                did_rename: None,
                will_rename: Some(FileOperationRegistrationOptions {
                    filters: vec![
                        FileOperationFilter {
                            scheme: Some(String::from("file")),
                            pattern: FileOperationPattern {
                                glob: String::from("**/*.rs"),
                                matches: Some(FileOperationPatternKind::File),
                                options: None,
                            },
                        FileOperationFilter {
                            scheme: Some(String::from("untitled")),
                            pattern: FileOperationPattern {
                                glob: String::from("**"),
                                matches: Some(FileOperationPatternKind::Folder),
                                options: None,
                            },
                        }],
                }),
                did_delete: None,
                will_delete: None,
            }),
        }),
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider: Some(
            SemanticTokensOptions {
                legend: SemanticTokensLegend {
                    token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                    token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
                },

                full: Some(SemanticTokensFullOptions::Delta { delta: Some(true) }),
                range: Some(true),
                work_done_progress_options: Default::default(),
            }
            .into(),
        ),
        moniker_provider: None,
        experimental: Some(json!({
            "joinLines": true,
            "ssr": true,
            "onEnter": true,
            "parentModule": true,
            "runnables": {
                "kinds": [ "cargo" ],
            },
        })),
    }
}

fn completions_resolve_provider(client_caps: &ClientCapabilities) -> Option<bool> {
    if enabled_completions_resolve_capabilities(client_caps)?.is_empty() {
        log::info!("No `additionalTextEdits` completion resolve capability was found in the client capabilities, autoimport completion is disabled");
        None
    } else {
        Some(true)
    }
}

/// Parses client capabilities and returns all completion resolve capabilities rust-analyzer supports.
pub(crate) fn enabled_completions_resolve_capabilities(
    caps: &ClientCapabilities,
) -> Option<FxHashSet<CompletionResolveCapability>> {
    Some(
        caps.text_document
            .as_ref()?
            .completion
            .as_ref()?
            .completion_item
            .as_ref()?
            .resolve_support
            .as_ref()?
            .properties
            .iter()
            .filter_map(|cap_string| match cap_string.as_str() {
                "additionalTextEdits" => Some(CompletionResolveCapability::AdditionalTextEdits),
                "detail" => Some(CompletionResolveCapability::Detail),
                "documentation" => Some(CompletionResolveCapability::Documentation),
                _unsupported => None,
            })
            .collect(),
    )
}

fn code_action_capabilities(client_caps: &ClientCapabilities) -> CodeActionProviderCapability {
    client_caps
        .text_document
        .as_ref()
        .and_then(|it| it.code_action.as_ref())
        .and_then(|it| it.code_action_literal_support.as_ref())
        .map_or(CodeActionProviderCapability::Simple(true), |_| {
            CodeActionProviderCapability::Options(CodeActionOptions {
                // Advertise support for all built-in CodeActionKinds.
                // Ideally we would base this off of the client capabilities
                // but the client is supposed to fall back gracefully for unknown values.
                code_action_kinds: Some(vec![
                    CodeActionKind::EMPTY,
                    CodeActionKind::QUICKFIX,
                    CodeActionKind::REFACTOR,
                    CodeActionKind::REFACTOR_EXTRACT,
                    CodeActionKind::REFACTOR_INLINE,
                    CodeActionKind::REFACTOR_REWRITE,
                ]),
                resolve_provider: Some(true),
                work_done_progress_options: Default::default(),
            })
        })
}
<|endoftext|>


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

use std::rc::*;

use super::component::*;
use super::super::tree::*;

struct FunctionComponent;

impl Component for FunctionComponent {
}

impl Drop for FunctionComponent {
    fn drop(&mut self) {
    }
}

///
/// Simplest form of 'component function': a function that receives a `TreeChange` indicating how the
/// input tree has changed, and returns a new change indicating how the output has changed.
///
impl BoxedComponentFactory for Box<Fn(&TreeChange) -> TreeChange> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let action              = self;

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            let change_result = action(change);
            our_publisher.publish(change_result);
        }));

        return Rc::new(FunctionComponent);
    }
}

///
/// Provides a component function that
///
impl BoxedComponentFactory for Box<Fn(&TreeRef) -> TreeRef> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let action              = self;

        let mut tree = "empty".to_tree_node();

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            tree = change.apply(&tree);

            let new_tree = action(&tree);

            our_publisher.publish(TreeChange::new(&TreeAddress::Here, TreeChangeType::Child, Some(&new_tree)));
        }));

        return Rc::new(FunctionComponent);
    }
}
<commit_msg>Add variants of the component factories that work on functions with mutable state<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

use std::rc::*;

use super::component::*;
use super::super::tree::*;

struct FunctionComponent;

impl Component for FunctionComponent {
}

impl Drop for FunctionComponent {
    fn drop(&mut self) {
    }
}

///
/// Simplest form of 'component function': a function that receives a `TreeChange` indicating how the
/// input tree has changed, and returns a new change indicating how the output has changed.
///
impl BoxedComponentFactory for Box<Fn(&TreeChange) -> TreeChange> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let action              = self;

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            let change_result = action(change);
            our_publisher.publish(change_result);
        }));

        return Rc::new(FunctionComponent);
    }
}

///
/// Simplest form of 'component function': a function that receives a `TreeChange` indicating how the
/// input tree has changed, and returns a new change indicating how the output has changed.
///
/// This variant allows for mutable state.
///
impl BoxedComponentFactory for Box<FnMut(&TreeChange) -> TreeChange> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let mut action          = self;

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            let change_result = action(change);
            our_publisher.publish(change_result);
        }));

        return Rc::new(FunctionComponent);
    }
}

///
/// Provides a component function that converts an input tree to an output tree
///
impl BoxedComponentFactory for Box<Fn(&TreeRef) -> TreeRef> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let action              = self;

        let mut tree = "empty".to_tree_node();

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            tree = change.apply(&tree);

            let new_tree = action(&tree);

            our_publisher.publish(TreeChange::new(&TreeAddress::Here, TreeChangeType::Child, Some(&new_tree)));
        }));

        return Rc::new(FunctionComponent);
    }
}

///
/// Provides a component function that converts an input tree to an output tree
///
/// This variant allows for mutable state
///
impl BoxedComponentFactory for Box<FnMut(&TreeRef) -> TreeRef> {
    ///
    /// Creates a component that consumes from a particular tree and publishes to a different tree
    ///
    fn create_boxed(self, consumer: ConsumerRef, publisher: PublisherRef) -> ComponentRef {
        let mut our_consumer    = consumer;
        let mut our_publisher   = publisher;
        let mut action          = self;

        let mut tree = "empty".to_tree_node();

        our_consumer.subscribe(TreeAddress::Here, TreeExtent::SubTree, Box::new(move |change| {
            tree = change.apply(&tree);

            let new_tree = action(&tree);

            our_publisher.publish(TreeChange::new(&TreeAddress::Here, TreeChangeType::Child, Some(&new_tree)));
        }));

        return Rc::new(FunctionComponent);
    }
}
<|endoftext|>