<commit_before>//! A future to figure out where we're going to dispatch a TCP socket.
//!
//! When the websocket server receives a TCP connection it may be a websocket
//! request or a general HTTP request. Right now the websocket library we're
//! using, Tungstenite, doesn't have built-in support for handling this
//! situation, so we roll our own.
//!
//! The general idea here is that we're going to read just enough data off the
//! socket to parse an initial HTTP request. This request will be parsed by the
//! `httparse` crate. Once we've got a request we take a look at the headers and
//! if we find a websocket upgrade we classify it as a websocket request. If
//! it's otherwise a `/status` request, we return that we're supposed to get the
//! status, and finally after all that if it doesn't match we return an error.
//!
//! This is basically a "poor man's" HTTP router and while it should be good
//! enough for now it should probably be extended/refactored in the future!
//!
//! Note that also to implement this we buffer the request that we read in
//! memory and then attach that to a socket once we've classified what kind of
//! socket this is. That's done to replay the bytes we read again for the
//! tungstenite library, which'll duplicate header parsing but we don't have
//! many other options for now!

use bytes::BytesMut;
use futures::{Future, Poll};
use httparse;
use tokio_core::net::TcpStream;
use tokio_io::AsyncRead;

use errors::*;
use server::webpush_io::WebpushIo;

pub struct Dispatch {
    socket: Option<TcpStream>,
    data: BytesMut,
}

pub enum RequestType {
    Websocket,
    Status,
}

impl Dispatch {
    pub fn new(socket: TcpStream) -> Dispatch {
        Dispatch {
            socket: Some(socket),
            data: BytesMut::new(),
        }
    }
}

impl Future for Dispatch {
    type Item = (WebpushIo, RequestType);
    type Error = Error;

    fn poll(&mut self) -> Poll<(WebpushIo, RequestType), Error> {
        loop {
            if self.data.len() == self.data.capacity() {
                self.data.reserve(16); // get some extra space
            }
            if try_ready!(self.socket.as_mut().unwrap().read_buf(&mut self.data)) == 0 {
                return Err("early eof".into())
            }
            let ty = {
                let mut headers = [httparse::EMPTY_HEADER; 16];
                let mut req = httparse::Request::new(&mut headers);
                match req.parse(&self.data)? {
                    httparse::Status::Complete(_) => {}
                    httparse::Status::Partial => continue,
                }

                if req.headers.iter().any(|h| h.name == "Upgrade") {
                    RequestType::Websocket
                } else if req.path == Some("/status") {
                    RequestType::Status
                } else {
                    debug!("unknown http request {:?}", req);
                    return Err("unknown http request".into())
                }
            };

            let tcp = self.socket.take().unwrap();
            return Ok((WebpushIo::new(tcp, self.data.take()), ty).into())
        }
    }
}
<commit_msg>fix: more lenient /status matching<commit_after>//! A future to figure out where we're going to dispatch a TCP socket.
//!
//! When the websocket server receives a TCP connection it may be a websocket
//! request or a general HTTP request. Right now the websocket library we're
//! using, Tungstenite, doesn't have built-in support for handling this
//! situation, so we roll our own.
//!
//! The general idea here is that we're going to read just enough data off the
//! socket to parse an initial HTTP request. This request will be parsed by the
//! `httparse` crate. Once we've got a request we take a look at the headers and
//! if we find a websocket upgrade we classify it as a websocket request. If
//! it's otherwise a `/status` request, we return that we're supposed to get the
//! status, and finally after all that if it doesn't match we return an error.
//!
//! This is basically a "poor man's" HTTP router and while it should be good
//! enough for now it should probably be extended/refactored in the future!
//!
//! Note that also to implement this we buffer the request that we read in
//! memory and then attach that to a socket once we've classified what kind of
//! socket this is. That's done to replay the bytes we read again for the
//! tungstenite library, which'll duplicate header parsing but we don't have
//! many other options for now!

use bytes::BytesMut;
use futures::{Future, Poll};
use httparse;
use tokio_core::net::TcpStream;
use tokio_io::AsyncRead;

use errors::*;
use server::webpush_io::WebpushIo;

pub struct Dispatch {
    socket: Option<TcpStream>,
    data: BytesMut,
}

pub enum RequestType {
    Websocket,
    Status,
}

impl Dispatch {
    pub fn new(socket: TcpStream) -> Dispatch {
        Dispatch {
            socket: Some(socket),
            data: BytesMut::new(),
        }
    }
}

impl Future for Dispatch {
    type Item = (WebpushIo, RequestType);
    type Error = Error;

    fn poll(&mut self) -> Poll<(WebpushIo, RequestType), Error> {
        loop {
            if self.data.len() == self.data.capacity() {
                self.data.reserve(16); // get some extra space
            }
            if try_ready!(self.socket.as_mut().unwrap().read_buf(&mut self.data)) == 0 {
                return Err("early eof".into())
            }
            let ty = {
                let mut headers = [httparse::EMPTY_HEADER; 16];
                let mut req = httparse::Request::new(&mut headers);
                match req.parse(&self.data)? {
                    httparse::Status::Complete(_) => {}
                    httparse::Status::Partial => continue,
                }

                if req.headers.iter().any(|h| h.name == "Upgrade") {
                    RequestType::Websocket
                } else {
                    match req.path {
                        Some(ref path) if path.starts_with("/status") => RequestType::Status,
                        _ => {
                            debug!("unknown http request {:?}", req);
                            return Err("unknown http request".into())
                        }
                    }
                }
            };

            let tcp = self.socket.take().unwrap();
            return Ok((WebpushIo::new(tcp, self.data.take()), ty).into())
        }
    }
}
<|endoftext|>


<commit_before>#![feature(unboxed_closures)]
#![feature(plugin)]
#![feature(box_syntax)]
#![feature(std_misc, collections, path_ext)]
#[macro_use]
#[macro_use]
extern crate wtftw;
extern crate wtftw_contrib;

use std::fs::PathExt;
use std::io::Write;
use std::path::Path;
use std::env;
use std::ops::Deref;
//use std::ffi::AsOsStr;
use wtftw::window_system::*;
use wtftw::window_manager::*;
use wtftw::handlers::default::*;
use wtftw::config::*;
use wtftw::util::*;
use wtftw::layout::Direction;
use wtftw::layout::LayoutMessage;
use wtftw_contrib::layout::{ AvoidStrutsLayout, LayoutCollection, BinarySpacePartition, GapLayout, MirrorLayout, NoBordersLayout, FullLayout };

#[no_mangle]
pub extern fn configure(_: &mut WindowManager, w: &WindowSystem, config: &mut Config) {
    let modm = MOD1MASK;

    config.general.mod_mask = modm;
    config.general.border_color = 0x404040;
    config.general.focus_border_color = 0xebebeb;
    config.general.border_width = 1;
    config.general.terminal = (String::from("xterm"), String::from(""));
    config.general.layout = LayoutCollection::new(vec!(
            GapLayout::new(0, AvoidStrutsLayout::new(vec!(Direction::Up, Direction::Down), BinarySpacePartition::new())),
            GapLayout::new(0, AvoidStrutsLayout::new(vec!(Direction::Up, Direction::Down), MirrorLayout::new(BinarySpacePartition::new()))),
            NoBordersLayout::new(box FullLayout)));

    config.general.tags = (vec!("一: ターミナル", "二: ウェブ", "三: コード",
                                "四: メディア", "五: スチーム", "六: ラテック",
                                "七: 音楽", "八: im", "九: 残り"))
        .into_iter().map(String::from).collect();

    // Register key handlers

    // Some standard key handlers for starting, restarting, etc.
    add_key_handler_str!(config, w, "q",      modm | SHIFTMASK, exit);
    add_key_handler_str!(config, w, "q",      modm,             restart);
    add_key_handler_str!(config, w, "Return", modm | SHIFTMASK, start_terminal);
    add_key_handler_str!(config, w, "p",      modm,             start_launcher);

    // Focus and window movement
    add_key_handler_str!(config, w, "j", modm,             |m, w, c| m.windows(w.deref(), c, &|x| x.focus_down()));
    add_key_handler_str!(config, w, "k", modm,             |m, w, c| m.windows(w.deref(), c, &|x| x.focus_up()));
    add_key_handler_str!(config, w, "j", modm | SHIFTMASK, |m, w, c| m.windows(w.deref(), c, &|x| x.swap_down()));
    add_key_handler_str!(config, w, "k", modm | SHIFTMASK, |m, w, c| m.windows(w.deref(), c, &|x| x.swap_up()));
    add_key_handler_str!(config, w, "Return", modm,        |m, w, c| m.windows(w.deref(), c, &|x| x.swap_master()));
    add_key_handler_str!(config, w, "c", modm, |m, w, c| m.kill_window(w.deref()).windows(w.deref(), c, &|x| x.clone()));

    add_key_handler_str!(config, w, "t", modm, |m, w, c| {
        match m.workspaces.peek() {
            Some(window) => m.windows(w.deref(), c, &|x| x.sink(window)),
            None => m.clone()
        }
    });

    // Layout messages
    add_key_handler_str!(config, w, "h",      modm,             send_layout_message!(LayoutMessage::Decrease));
    add_key_handler_str!(config, w, "l",      modm,             send_layout_message!(LayoutMessage::Increase));
    add_key_handler_str!(config, w, "z",      modm,             send_layout_message!(LayoutMessage::DecreaseSlave));
    add_key_handler_str!(config, w, "a",      modm,             send_layout_message!(LayoutMessage::IncreaseSlave));
    add_key_handler_str!(config, w, "x",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::IncreaseGap));
    add_key_handler_str!(config, w, "s",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::DecreaseGap));
    add_key_handler_str!(config, w, "comma",  modm,             send_layout_message!(LayoutMessage::IncreaseMaster));
    add_key_handler_str!(config, w, "period", modm,             send_layout_message!(LayoutMessage::DecreaseMaster));
    add_key_handler_str!(config, w, "space",  modm,             send_layout_message!(LayoutMessage::Next));
    add_key_handler_str!(config, w, "space",  modm | SHIFTMASK, send_layout_message!(LayoutMessage::Prev));
    add_key_handler_str!(config, w, "r",      modm,             send_layout_message!(LayoutMessage::TreeRotate));
    add_key_handler_str!(config, w, "s",      modm,             send_layout_message!(LayoutMessage::TreeSwap));
    add_key_handler_str!(config, w, "u",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Left)));
    add_key_handler_str!(config, w, "p",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Right)));
    add_key_handler_str!(config, w, "i",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Down)));
    add_key_handler_str!(config, w, "o",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Up)));
    add_key_handler_str!(config, w, "u",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Left)));
    add_key_handler_str!(config, w, "p",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Right)));
    add_key_handler_str!(config, w, "i",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Down)));
    add_key_handler_str!(config, w, "o",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Up)));

    // Workspace switching and moving
    for i in (1usize..10) {
        add_key_handler_str!(config, w, &i.to_string(), modm,
        move |m, w, c| switch_to_workspace(m, w, c, i - 1));

        add_key_handler_str!(config, w, &i.to_string(), modm | SHIFTMASK,
        move |m, w, c| move_window_to_workspace(m, w, c, i - 1));
    }

    // Media keys
    add_key_handler_str!(config, w, "j", modm | CONTROLMASK, run!("amixer", "-q set Master 5%-"));
    add_key_handler_str!(config, w, "k", modm | CONTROLMASK, run!("amixer", "-q set Master 5%+"));

    add_key_handler_code!(config, 0x1008ff11, NONEMASK, run!("amixer", "-q set Master 5%-"));
    add_key_handler_code!(config, 0x1008ff13, NONEMASK, run!("amixer", "-q set Master 5%+"));

    add_key_handler_code!(config, 0x1008ff02, NONEMASK, run!("xbacklight", "+10"));
    add_key_handler_code!(config, 0x1008ff03, NONEMASK, run!("xbacklight", "-10"));

    add_mouse_handler!(config, BUTTON1, modm,
                       |m, w, c, s| {
                           m.focus(s, w.deref(), c).mouse_move_window(w.deref(), c, s).windows(w.deref(), c, &|x| x.shift_master())
                       });
    add_mouse_handler!(config, BUTTON3, modm,
                       |m, w, c, s| {
                           m.focus(s, w.deref(), c).mouse_resize_window(w.deref(), c, s).windows(w.deref(), c, &|x| x.shift_master())
                       });

    // Xmobar handling and formatting
    let home_dir = env::home_dir().unwrap();
    let home = home_dir.as_os_str().to_str().unwrap();
    let xmobar_config = format!("{}/.xmonad/xmobar1.hs", home);

    if Path::new(&xmobar_config.clone()).is_file() {
        let mut xmobar = spawn_pipe(config, "/home/rootnode/.cabal/bin/xmobar",
                                    vec!(xmobar_config));
        let tags = config.general.tags.clone();
        config.set_log_hook(box move |m, w| {
            let p = &mut xmobar;
            let tags = &tags;
            let workspaces = tags.clone().iter()
                .enumerate()
                .map(|(i, _)| if i as u32 == m.workspaces.current.workspace.id {
                    format!("<fc=#fb4934>■</fc>")
                } else if m.workspaces.visible.iter().any(|w| w.workspace.id == i as u32) {
                    format!("<fc=#fabd2f>■</fc>")
                } else {
                    format!("■")
                })
            .fold(String::from(""), |a, x| {
                let mut r = a.clone();
                r.push_str(&x);
                r
            });

            let name = match m.workspaces.peek() {
                None => String::from(""),
                Some(window) => w.get_window_name(window)
            };
            let content = format!("{} {} {}\n", workspaces, m.workspaces.current.workspace.layout.description(), name);
            match p.write().unwrap().stdin.as_mut() {
                Some(pin) => {
                    match pin.write(content.as_bytes()) {
                        _      => ()
                    };
                    pin.flush().unwrap();
                },
                _ => ()
            }
        });
    };
}
<commit_msg>Config without unstable features<commit_after>#[macro_use]
#[macro_use]
extern crate wtftw;
extern crate wtftw_contrib;

use std::io::Write;
use std::fs::metadata;
use std::env;
use std::ops::Deref;
//use std::ffi::AsOsStr;
use wtftw::window_system::*;
use wtftw::window_manager::*;
use wtftw::handlers::default::*;
use wtftw::config::*;
use wtftw::util::*;
use wtftw::layout::Direction;
use wtftw::layout::LayoutMessage;
use wtftw_contrib::layout::{ AvoidStrutsLayout, LayoutCollection, BinarySpacePartition, GapLayout, MirrorLayout, NoBordersLayout, FullLayout };

#[no_mangle]
pub extern fn configure(_: &mut WindowManager, w: &WindowSystem, config: &mut Config) {
    let modm = MOD1MASK;

    config.general.mod_mask = modm;
    config.general.border_color = 0x404040;
    config.general.focus_border_color = 0xebebeb;
    config.general.border_width = 1;
    config.general.terminal = (String::from("xterm"), String::from(""));
    config.general.layout = LayoutCollection::new(vec!(
            GapLayout::new(0, AvoidStrutsLayout::new(vec!(Direction::Up, Direction::Down), BinarySpacePartition::new())),
            GapLayout::new(0, AvoidStrutsLayout::new(vec!(Direction::Up, Direction::Down), MirrorLayout::new(BinarySpacePartition::new()))),
            NoBordersLayout::new(Box::new(FullLayout))));

    config.general.tags = (vec!("一: ターミナル", "二: ウェブ", "三: コード",
                                "四: メディア", "五: スチーム", "六: ラテック",
                                "七: 音楽", "八: im", "九: 残り"))
        .into_iter().map(String::from).collect();

    // Register key handlers

    // Some standard key handlers for starting, restarting, etc.
    add_key_handler_str!(config, w, "q",      modm | SHIFTMASK, exit);
    add_key_handler_str!(config, w, "q",      modm,             restart);
    add_key_handler_str!(config, w, "Return", modm | SHIFTMASK, start_terminal);
    add_key_handler_str!(config, w, "p",      modm,             start_launcher);

    // Focus and window movement
    add_key_handler_str!(config, w, "j", modm,             |m, w, c| m.windows(w.deref(), c, &|x| x.focus_down()));
    add_key_handler_str!(config, w, "k", modm,             |m, w, c| m.windows(w.deref(), c, &|x| x.focus_up()));
    add_key_handler_str!(config, w, "j", modm | SHIFTMASK, |m, w, c| m.windows(w.deref(), c, &|x| x.swap_down()));
    add_key_handler_str!(config, w, "k", modm | SHIFTMASK, |m, w, c| m.windows(w.deref(), c, &|x| x.swap_up()));
    add_key_handler_str!(config, w, "Return", modm,        |m, w, c| m.windows(w.deref(), c, &|x| x.swap_master()));
    add_key_handler_str!(config, w, "c", modm, |m, w, c| m.kill_window(w.deref()).windows(w.deref(), c, &|x| x.clone()));

    add_key_handler_str!(config, w, "t", modm, |m, w, c| {
        match m.workspaces.peek() {
            Some(window) => m.windows(w.deref(), c, &|x| x.sink(window)),
            None => m.clone()
        }
    });

    // Layout messages
    add_key_handler_str!(config, w, "h",      modm,             send_layout_message!(LayoutMessage::Decrease));
    add_key_handler_str!(config, w, "l",      modm,             send_layout_message!(LayoutMessage::Increase));
    add_key_handler_str!(config, w, "z",      modm,             send_layout_message!(LayoutMessage::DecreaseSlave));
    add_key_handler_str!(config, w, "a",      modm,             send_layout_message!(LayoutMessage::IncreaseSlave));
    add_key_handler_str!(config, w, "x",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::IncreaseGap));
    add_key_handler_str!(config, w, "s",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::DecreaseGap));
    add_key_handler_str!(config, w, "comma",  modm,             send_layout_message!(LayoutMessage::IncreaseMaster));
    add_key_handler_str!(config, w, "period", modm,             send_layout_message!(LayoutMessage::DecreaseMaster));
    add_key_handler_str!(config, w, "space",  modm,             send_layout_message!(LayoutMessage::Next));
    add_key_handler_str!(config, w, "space",  modm | SHIFTMASK, send_layout_message!(LayoutMessage::Prev));
    add_key_handler_str!(config, w, "r",      modm,             send_layout_message!(LayoutMessage::TreeRotate));
    add_key_handler_str!(config, w, "s",      modm,             send_layout_message!(LayoutMessage::TreeSwap));
    add_key_handler_str!(config, w, "u",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Left)));
    add_key_handler_str!(config, w, "p",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Right)));
    add_key_handler_str!(config, w, "i",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Down)));
    add_key_handler_str!(config, w, "o",      modm | SHIFTMASK, send_layout_message!(LayoutMessage::TreeExpandTowards(Direction::Up)));
    add_key_handler_str!(config, w, "u",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Left)));
    add_key_handler_str!(config, w, "p",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Right)));
    add_key_handler_str!(config, w, "i",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Down)));
    add_key_handler_str!(config, w, "o",      modm | CONTROLMASK, send_layout_message!(LayoutMessage::TreeShrinkFrom(Direction::Up)));

    // Workspace switching and moving
    for i in (1usize..10) {
        add_key_handler_str!(config, w, &i.to_string(), modm,
        move |m, w, c| switch_to_workspace(m, w, c, i - 1));

        add_key_handler_str!(config, w, &i.to_string(), modm | SHIFTMASK,
        move |m, w, c| move_window_to_workspace(m, w, c, i - 1));
    }

    // Media keys
    add_key_handler_str!(config, w, "j", modm | CONTROLMASK, run!("amixer", "-q set Master 5%-"));
    add_key_handler_str!(config, w, "k", modm | CONTROLMASK, run!("amixer", "-q set Master 5%+"));

    add_key_handler_code!(config, 0x1008ff11, NONEMASK, run!("amixer", "-q set Master 5%-"));
    add_key_handler_code!(config, 0x1008ff13, NONEMASK, run!("amixer", "-q set Master 5%+"));

    add_key_handler_code!(config, 0x1008ff02, NONEMASK, run!("xbacklight", "+10"));
    add_key_handler_code!(config, 0x1008ff03, NONEMASK, run!("xbacklight", "-10"));

    add_mouse_handler!(config, BUTTON1, modm,
                       |m, w, c, s| {
                           m.focus(s, w.deref(), c).mouse_move_window(w.deref(), c, s).windows(w.deref(), c, &|x| x.shift_master())
                       });
    add_mouse_handler!(config, BUTTON3, modm,
                       |m, w, c, s| {
                           m.focus(s, w.deref(), c).mouse_resize_window(w.deref(), c, s).windows(w.deref(), c, &|x| x.shift_master())
                       });

    // Xmobar handling and formatting
    let home_dir = env::home_dir().unwrap();
    let home = home_dir.as_os_str().to_str().unwrap();
    let xmobar_config = format!("{}/.xmonad/xmobar1.hs", home);

    if metadata(&xmobar_config).map(|s| s.is_file()).unwrap_or(false) {
        let mut xmobar = spawn_pipe(config, "/home/rootnode/.cabal/bin/xmobar",
                                    vec!(xmobar_config));
        let tags = config.general.tags.clone();
        config.set_log_hook(Box::new(move |m, w| {
            let p = &mut xmobar;
            let tags = &tags;
            let workspaces = tags.clone().iter()
                .enumerate()
                .map(|(i, _)| if i as u32 == m.workspaces.current.workspace.id {
                    format!("<fc=#fb4934>■</fc>")
                } else if m.workspaces.visible.iter().any(|w| w.workspace.id == i as u32) {
                    format!("<fc=#fabd2f>■</fc>")
                } else {
                    format!("■")
                })
            .fold(String::from(""), |a, x| {
                let mut r = a.clone();
                r.push_str(&x);
                r
            });

            let name = match m.workspaces.peek() {
                None => String::from(""),
                Some(window) => w.get_window_name(window)
            };
            let content = format!("{} {} {}\n", workspaces, m.workspaces.current.workspace.layout.description(), name);
            match p.write().unwrap().stdin.as_mut() {
                Some(pin) => {
                    match pin.write(content.as_bytes()) {
                        _      => ()
                    };
                    pin.flush().unwrap();
                },
                _ => ()
            }
        }));
    };
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;

use serialize::*;
use serialize_packed::*;

pub fn expectPacksTo(unpacked : &[u8],
                     packed : &[u8]) {

    use std::rt::io::Writer;
    use std::rt::io::mem::MemWriter;
    use std::rt::io::extensions::ReaderUtil;

    // --------
    // write

    // XXX this is broken and does not compile
    let bytes = do std::rt::io::mem::with_mem_writer |writer| {
        let mut packedOutputStream = PackedOutputStream {inner : writer};
        packedOutputStream.write(unpacked);
    };

    assert!(bytes.slice(0, bytes.len()).equals(&packed));

    // --------
    // read

    let mut reader = std::rt::io::mem::BufReader::new(packed);
    let mut packedInputStream =
        PackedInputStream {inner : reader};

    let bytes = packedInputStream.read_to_end();

    assert!(bytes.slice(0, bytes.len()).equals(&unpacked));

}

#[test]
pub fn simplePacking() {
    expectPacksTo([], []);
    expectPacksTo([0,0,0,0,0,0,0,0], [0,0]);
    expectPacksTo([0,0,12,0,0,34,0,0], [0x24,12,34]);
    expectPacksTo([1,3,2,4,5,7,6,8], [0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([0,0,0,0,0,0,0,0,1,3,2,4,5,7,6,8], [0,0,0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([0,0,12,0,0,34,0,0,1,3,2,4,5,7,6,8], [0x24,12,34,0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([1,3,2,4,5,7,6,8,8,6,7,4,5,2,3,1], [0xff,1,3,2,4,5,7,6,8,1,8,6,7,4,5,2,3,1]);

    expectPacksTo(
        [1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        [0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8,
         0xd6,2,4,9,5,1]);
    expectPacksTo(
        [1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        [0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8,
         0xd6,2,4,9,5,1]);

    expectPacksTo(
        [8,0,100,6,0,1,1,2, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,1,0,2,0,3,1],
        [0xed,8,100,6,1,1,2, 0,2, 0xd4,1,2,3,1]);
}
<commit_msg>use (unimplemented) BufWriter in serialized_packed_test<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;

use serialize_packed::*;

pub fn expectPacksTo(unpacked : &[u8],
                     packed : &[u8]) {

    use std::rt::io::Writer;
    use std::rt::io::extensions::ReaderUtil;

    // --------
    // write

    let mut bytes : ~[u8] = std::vec::from_elem(packed.len(), 0u8);

    {
        let writer = std::rt::io::mem::BufWriter::new(bytes);
        let mut packedOutputStream = PackedOutputStream {inner : writer};
        packedOutputStream.write(unpacked);
    }

    assert!(bytes.slice(0, bytes.len()).equals(&packed));

    // --------
    // read

    let reader = std::rt::io::mem::BufReader::new(packed);
    let mut packedInputStream =
        PackedInputStream {inner : reader};

    let bytes = packedInputStream.read_to_end();

    assert!(bytes.slice(0, bytes.len()).equals(&unpacked));

}

#[test]
pub fn simplePacking() {
    expectPacksTo([], []);
    expectPacksTo([0,0,0,0,0,0,0,0], [0,0]);
    expectPacksTo([0,0,12,0,0,34,0,0], [0x24,12,34]);
    expectPacksTo([1,3,2,4,5,7,6,8], [0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([0,0,0,0,0,0,0,0,1,3,2,4,5,7,6,8], [0,0,0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([0,0,12,0,0,34,0,0,1,3,2,4,5,7,6,8], [0x24,12,34,0xff,1,3,2,4,5,7,6,8,0]);
    expectPacksTo([1,3,2,4,5,7,6,8,8,6,7,4,5,2,3,1], [0xff,1,3,2,4,5,7,6,8,1,8,6,7,4,5,2,3,1]);

    expectPacksTo(
        [1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        [0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8,
         0xd6,2,4,9,5,1]);
    expectPacksTo(
        [1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        [0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8,
         0xd6,2,4,9,5,1]);

    expectPacksTo(
        [8,0,100,6,0,1,1,2, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,1,0,2,0,3,1],
        [0xed,8,100,6,1,1,2, 0,2, 0xd4,1,2,3,1]);
}
<|endoftext|>


<commit_before>extern crate websocket;
extern crate futures;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate serde_derive;

use std::fmt::Debug;

use websocket::message::{Message, OwnedMessage};
use websocket::server::InvalidConnection;
use websocket::client::async::Framed;
use websocket::async::{Server, MessageCodec};
use websocket::WebSocketError;

use tokio_core::net::TcpStream;
use tokio_core::reactor::{Handle, Core};

use futures::{Future, Sink, Stream};
use futures::future::{self, Loop};
use std::sync::{Mutex, Arc};

fn main() {
    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let server = Server::bind("localhost:8081", &handle).unwrap();

    let f = server.incoming()
        // we don't wanna save the stream if it drops
        .map_err(|InvalidConnection { error, .. }| error)
        .for_each(|(upgrade, addr)| {
            println!("Got a connection from: {}", addr);
            if !upgrade.protocols().iter().any(|s| s == "rust-websocket") {
                spawn_future(upgrade.reject(), "Upgrade Rejection", &handle);
                return Ok(());
            }

            let f = upgrade
                .use_protocol("rust-websocket")
                .accept()
                .and_then(|(server, _)| server.send(Message::text("Hello World!").into()))
                .and_then(|server| {
                    future::loop_fn(server, |stream| {
                        stream.into_future()
                            .or_else(|(err, stream)| {
                                println!("Could not send message: {:?}", err);
                                stream.send(OwnedMessage::Close(None)).map(|s| (None, s))
                            })
                            .and_then(|(msg, stream)|{
                                let mut state = State::new();
                                handle_incomming(&mut state, msg);
                                send(&state, stream)
                            })
                    })
                });
            spawn_future(f, "Client Status", &handle);
            Ok(())
        });

    core.run(f).unwrap();
}

fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
    where F: Future<Item = I, Error = E> + 'static,
          E: Debug
{
    handle.spawn(f.map_err(move |e| println!("{}: '{:?}'", desc, e))
        .map(move |_| println!("{}: Finished.", desc)));
}

fn handle_incomming(state: &mut State, msg: Option<OwnedMessage>) {
    /*match msg {
        Some(OwnedMessage::Text(txt)) => {
            stream.send(OwnedMessage::Text(txt))
                .map(|s| Loop::Continue(s))
                .boxed()
        }
        Some(OwnedMessage::Binary(bin)) => {
            stream.send(OwnedMessage::Binary(bin))
                .map(|s| Loop::Continue(s))
                .boxed()
        }
        Some(OwnedMessage::Ping(data)) => {
            stream.send(OwnedMessage::Pong(data))
                .map(|s| Loop::Continue(s))
                .boxed()
        }
        Some(OwnedMessage::Close(_)) => {
            stream.send(OwnedMessage::Close(None))
                .map(|_| Loop::Break(()))
                .boxed()
        }
        Some(OwnedMessage::Pong(_)) => {
            future::ok(Loop::Continue(stream)).boxed()
        }
        None => {
            future::ok(Loop::Break(())).boxed()
        },
    }*/
}

type FramedStream = Framed<TcpStream, MessageCodec<OwnedMessage>>;

fn send(state: &State, stream: FramedStream) -> Box<Future<Item=Loop<(), FramedStream>,Error=WebSocketError>> {
    stream
        .send(OwnedMessage::Text(serde_json::to_string(&state).unwrap()))
        .map(|s| {
            if shutdown() {
                Loop::Break(())
            } else {
                Loop::Continue(s)
            }
        })
        .boxed()

}

fn shutdown() -> bool {
    false
}

#[derive(Serialize)]
struct State {
    msg: String,
}

impl State {
    fn new() -> State {
        State {
            msg: "git gud".to_string()
        }
    }
}

<commit_msg>Core: Simplify things<commit_after>extern crate websocket;
extern crate futures;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate serde_derive;

use std::fmt::Debug;

use websocket::message::{Message, OwnedMessage};
use websocket::server::InvalidConnection;
use websocket::client::async::Framed;
use websocket::async::{Server, MessageCodec};
use websocket::WebSocketError;

use tokio_core::net::TcpStream;
use tokio_core::reactor::{Handle, Core};

use futures::{Future, Sink, Stream};
use futures::future::{self, Loop};
use std::sync::{Mutex, Arc};

fn main() {
    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let server = Server::bind("localhost:8081", &handle).unwrap();

    let f = server.incoming()
        // we don't wanna save the stream if it drops
        .map_err(|InvalidConnection { error, .. }| error)
        .for_each(|(upgrade, addr)| {
            println!("Got a connection from: {}", addr);
            if !upgrade.protocols().iter().any(|s| s == "rust-websocket") {
                spawn_future(upgrade.reject(), "Upgrade Rejection", &handle);
                return Ok(());
            }

            let client = upgrade
                .use_protocol("rust-websocket")
                .accept();
            let f = client
                .and_then(|(stream, _)| {
                 future::loop_fn(stream, |stream| {
                        stream
                            .into_future()
                            .or_else(|(err, stream)| {
                                println!("Could not send message: {:?}", err);
                                stream.send(OwnedMessage::Close(None)).map(|s| (None, s))
                            })
                            .and_then(|(msg, stream)|{
                                let mut state = State::new();
                                handle_incomming(&mut state, msg);
                                Ok(stream)
                            }).map(|stream| {
                                if shutdown() {
                                    Loop::Break(())
                                } else {
                                    Loop::Continue(stream)
                                }
                            })
                            .boxed()
                    })
            });
            spawn_future(f, "Client Status", &handle);
            Ok(())
        });

    core.run(f).unwrap();
}

fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
    where F: Future<Item = I, Error = E> + 'static,
          E: Debug
{
    handle.spawn(f.map_err(move |e| println!("{}: '{:?}'", desc, e))
        .map(move |_| println!("{}: Finished.", desc)));
}

fn handle_incomming(state: &mut State, msg: Option<OwnedMessage>) {
    match msg {
        Some(OwnedMessage::Text(txt)) => {
            println!("Received message: {}", txt);
        }
        _ => {}
    }
}

type FramedStream = Framed<TcpStream, MessageCodec<OwnedMessage>>;

fn send(state: &State, stream: FramedStream) -> Box<Future<Item=Loop<(), FramedStream>,Error=WebSocketError>> {
    let msg = serde_json::to_string(&state).unwrap();
    println!("Sending message: {}", msg);
    stream
        .send(OwnedMessage::Text(msg))
        .map(|s| {
            if shutdown() {
                Loop::Break(())
            } else {
                Loop::Continue(s)
            }
        })
        .boxed()

}

fn shutdown() -> bool {
    false
}

#[derive(Serialize)]
struct State {
    msg: String,
}

impl State {
    fn new() -> State {
        State {
            msg: "git gud".to_string()
        }
    }
}

<|endoftext|>


<commit_before>//! SCIP generator

use std::{
    collections::{HashMap, HashSet},
    time::Instant,
};

use crate::line_index::{LineEndings, LineIndex, OffsetEncoding};
use hir::Name;
use ide::{
    LineCol, MonikerDescriptorKind, MonikerResult, StaticIndex, StaticIndexedFile, TextRange,
    TokenId,
};
use ide_db::LineIndexDatabase;
use project_model::{CargoConfig, ProjectManifest, ProjectWorkspace};
use scip::types as scip_types;
use std::env;

use crate::cli::{
    flags,
    load_cargo::{load_workspace, LoadCargoConfig},
    Result,
};

impl flags::Scip {
    pub fn run(self) -> Result<()> {
        eprintln!("Generating SCIP start...");
        let now = Instant::now();
        let cargo_config = CargoConfig::default();

        let no_progress = &|s| (eprintln!("rust-analyzer: Loading {}", s));
        let load_cargo_config = LoadCargoConfig {
            load_out_dirs_from_check: true,
            with_proc_macro: true,
            prefill_caches: true,
        };
        let path = vfs::AbsPathBuf::assert(env::current_dir()?.join(&self.path));
        let rootpath = path.normalize();
        let manifest = ProjectManifest::discover_single(&path)?;

        let workspace = ProjectWorkspace::load(manifest, &cargo_config, no_progress)?;

        let (host, vfs, _) =
            load_workspace(workspace, &cargo_config.extra_env, &load_cargo_config)?;
        let db = host.raw_database();
        let analysis = host.analysis();

        let si = StaticIndex::compute(&analysis);

        let mut index = scip_types::Index {
            metadata: Some(scip_types::Metadata {
                version: scip_types::ProtocolVersion::UnspecifiedProtocolVersion.into(),
                tool_info: Some(scip_types::ToolInfo {
                    name: "rust-analyzer".to_owned(),
                    version: "0.1".to_owned(),
                    arguments: vec![],
                    ..Default::default()
                })
                .into(),
                project_root: format!(
                    "file://{}",
                    path.normalize()
                        .as_os_str()
                        .to_str()
                        .ok_or(anyhow::anyhow!("Unable to normalize project_root path"))?
                        .to_string()
                ),
                text_document_encoding: scip_types::TextEncoding::UTF8.into(),
                ..Default::default()
            })
            .into(),
            ..Default::default()
        };

        let mut symbols_emitted: HashSet<TokenId> = HashSet::default();
        let mut tokens_to_symbol: HashMap<TokenId, String> = HashMap::new();

        for StaticIndexedFile { file_id, tokens, .. } in si.files {
            let mut local_count = 0;
            let mut new_local_symbol = || {
                let new_symbol = scip::types::Symbol::new_local(local_count);
                local_count += 1;

                new_symbol
            };

            let relative_path = match get_relative_filepath(&vfs, &rootpath, file_id) {
                Some(relative_path) => relative_path,
                None => continue,
            };

            let line_index = LineIndex {
                index: db.line_index(file_id),
                encoding: OffsetEncoding::Utf8,
                endings: LineEndings::Unix,
            };

            let mut doc = scip_types::Document {
                relative_path,
                language: "rust".to_string(),
                ..Default::default()
            };

            tokens.into_iter().for_each(|(range, id)| {
                let token = si.tokens.get(id).unwrap();

                let mut occurrence = scip_types::Occurrence::default();
                occurrence.range = text_range_to_scip_range(&line_index, range);
                occurrence.symbol = tokens_to_symbol
                    .entry(id)
                    .or_insert_with(|| {
                        let symbol = match &token.moniker {
                            Some(moniker) => moniker_to_symbol(&moniker),
                            None => new_local_symbol(),
                        };
                        scip::symbol::format_symbol(symbol)
                    })
                    .clone();

                if let Some(def) = token.definition {
                    if def.range == range {
                        occurrence.symbol_roles |= scip_types::SymbolRole::Definition as i32;
                    }

                    if symbols_emitted.insert(id) {
                        let mut symbol_info = scip_types::SymbolInformation::default();
                        symbol_info.symbol = occurrence.symbol.clone();
                        if let Some(hover) = &token.hover {
                            if !hover.markup.as_str().is_empty() {
                                symbol_info.documentation = vec![hover.markup.as_str().to_string()];
                            }
                        }

                        doc.symbols.push(symbol_info)
                    }
                }

                doc.occurrences.push(occurrence);
            });

            if doc.occurrences.is_empty() {
                continue;
            }

            index.documents.push(doc);
        }

        scip::write_message_to_file("index.scip", index)
            .map_err(|err| anyhow::anyhow!("Failed to write scip to file: {}", err))?;

        eprintln!("Generating SCIP finished {:?}", now.elapsed());
        Ok(())
    }
}

fn get_relative_filepath(
    vfs: &vfs::Vfs,
    rootpath: &vfs::AbsPathBuf,
    file_id: ide::FileId,
) -> Option<String> {
    Some(vfs.file_path(file_id).as_path()?.strip_prefix(&rootpath)?.as_ref().to_str()?.to_string())
}

// SCIP Ranges have a (very large) optimization that ranges if they are on the same line
// only encode as a vector of [start_line, start_col, end_col].
//
// This transforms a line index into the optimized SCIP Range.
fn text_range_to_scip_range(line_index: &LineIndex, range: TextRange) -> Vec<i32> {
    let LineCol { line: start_line, col: start_col } = line_index.index.line_col(range.start());
    let LineCol { line: end_line, col: end_col } = line_index.index.line_col(range.end());

    if start_line == end_line {
        vec![start_line as i32, start_col as i32, end_col as i32]
    } else {
        vec![start_line as i32, start_col as i32, end_line as i32, end_col as i32]
    }
}

fn new_descriptor_str(
    name: &str,
    suffix: scip_types::descriptor::Suffix,
) -> scip_types::Descriptor {
    scip_types::Descriptor {
        name: name.to_string(),
        disambiguator: "".to_string(),
        suffix: suffix.into(),
        ..Default::default()
    }
}

fn new_descriptor(name: Name, suffix: scip_types::descriptor::Suffix) -> scip_types::Descriptor {
    let mut name = name.to_string();
    if name.contains("'") {
        name = format!("`{}`", name);
    }

    new_descriptor_str(name.as_str(), suffix)
}

/// Loosely based on `def_to_moniker`
///
/// Only returns a Symbol when it's a non-local symbol.
///     So if the visibility isn't outside of a document, then it will return None
fn moniker_to_symbol(moniker: &MonikerResult) -> scip_types::Symbol {
    use scip_types::descriptor::Suffix::*;

    let package_name = moniker.package_information.name.clone();
    let version = moniker.package_information.version.clone();
    let descriptors = moniker
        .identifier
        .description
        .iter()
        .map(|desc| {
            new_descriptor(
                desc.name.clone(),
                match desc.desc {
                    MonikerDescriptorKind::Namespace => Namespace,
                    MonikerDescriptorKind::Type => Type,
                    MonikerDescriptorKind::Term => Term,
                    MonikerDescriptorKind::Method => Method,
                    MonikerDescriptorKind::TypeParameter => TypeParameter,
                    MonikerDescriptorKind::Parameter => Parameter,
                    MonikerDescriptorKind::Macro => Macro,
                    MonikerDescriptorKind::Meta => Meta,
                },
            )
        })
        .collect();

    scip_types::Symbol {
        scheme: "rust-analyzer".into(),
        package: Some(scip_types::Package {
            manager: "cargo".to_string(),
            name: package_name,
            version,
            ..Default::default()
        })
        .into(),
        descriptors,
        ..Default::default()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use hir::Semantics;
    use ide::{AnalysisHost, FilePosition};
    use ide_db::defs::IdentClass;
    use ide_db::{base_db::fixture::ChangeFixture, helpers::pick_best_token};
    use scip::symbol::format_symbol;
    use syntax::SyntaxKind::*;
    use syntax::{AstNode, T};

    fn position(ra_fixture: &str) -> (AnalysisHost, FilePosition) {
        let mut host = AnalysisHost::default();
        let change_fixture = ChangeFixture::parse(ra_fixture);
        host.raw_database_mut().apply_change(change_fixture.change);
        let (file_id, range_or_offset) =
            change_fixture.file_position.expect("expected a marker ($0)");
        let offset = range_or_offset.expect_offset();
        (host, FilePosition { file_id, offset })
    }

    /// If expected == "", then assert that there are no symbols (this is basically local symbol)
    #[track_caller]
    fn check_symbol(ra_fixture: &str, expected: &str) {
        let (host, position) = position(ra_fixture);

        let FilePosition { file_id, offset } = position;

        let db = host.raw_database();
        let sema = &Semantics::new(db);
        let file = sema.parse(file_id).syntax().clone();
        let original_token = pick_best_token(file.token_at_offset(offset), |kind| match kind {
            IDENT
            | INT_NUMBER
            | LIFETIME_IDENT
            | T![self]
            | T![super]
            | T![crate]
            | T![Self]
            | COMMENT => 2,
            kind if kind.is_trivia() => 0,
            _ => 1,
        })
        .expect("OK OK");

        let navs = sema
            .descend_into_macros(original_token.clone())
            .into_iter()
            .filter_map(|token| {
                IdentClass::classify_token(sema, &token).map(IdentClass::definitions).map(|it| {
                    it.into_iter().flat_map(|def| {
                        let module = def.module(db).unwrap();
                        let current_crate = module.krate();

                        match MonikerResult::from_def(sema.db, def, current_crate) {
                            Some(moniker_result) => Some(moniker_to_symbol(&moniker_result)),
                            None => None,
                        }
                    })
                })
            })
            .flatten()
            .collect::<Vec<_>>();

        if expected == "" {
            assert_eq!(0, navs.len(), "must have no symbols {:?}", navs);
            return;
        }

        assert_eq!(1, navs.len(), "must have one symbol {:?}", navs);

        let res = navs.get(0).unwrap();
        let formatted = format_symbol(res.clone());
        assert_eq!(formatted, expected);
    }

    #[test]
    fn basic() {
        check_symbol(
            r#"
//- /lib.rs crate:main deps:foo
use foo::example_mod::func;
fn main() {
    func$0();
}
//- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
pub mod example_mod {
    pub fn func() {}
}
"#,
            "rust-analyzer cargo foo 0.1.0 example_mod/func().",
        );
    }

    #[test]
    fn symbol_for_trait() {
        check_symbol(
            r#"
//- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
pub mod module {
    pub trait MyTrait {
        pub fn func$0() {}
    }
}
"#,
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#func().",
        );
    }

    #[test]
    fn symbol_for_trait_constant() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            const MY_CONST$0: u8;
        }
    }
    "#,
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#MY_CONST.",
        );
    }

    #[test]
    fn symbol_for_trait_type() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            type MyType$0;
        }
    }
    "#,
            // "foo::module::MyTrait::MyType",
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#[MyType]",
        );
    }

    #[test]
    fn symbol_for_trait_impl_function() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            pub fn func() {}
        }

        struct MyStruct {}

        impl MyTrait for MyStruct {
            pub fn func$0() {}
        }
    }
    "#,
            // "foo::module::MyStruct::MyTrait::func",
            "rust-analyzer cargo foo 0.1.0 module/MyStruct#MyTrait#func().",
        );
    }

    #[test]
    fn symbol_for_field() {
        check_symbol(
            r#"
    //- /lib.rs crate:main deps:foo
    use foo::St;
    fn main() {
        let x = St { a$0: 2 };
    }
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub struct St {
        pub a: i32,
    }
    "#,
            "rust-analyzer cargo foo 0.1.0 St#a.",
        );
    }

    #[test]
    fn local_symbol_for_local() {
        check_symbol(
            r#"
    //- /lib.rs crate:main deps:foo
    use foo::module::func;
    fn main() {
        func();
    }
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub fn func() {
            let x$0 = 2;
        }
    }
    "#,
            "",
        );
    }
}
<commit_msg>scip: Rewrite tests to be closer to what we actually do.<commit_after>//! SCIP generator

use std::{
    collections::{HashMap, HashSet},
    time::Instant,
};

use crate::line_index::{LineEndings, LineIndex, OffsetEncoding};
use hir::Name;
use ide::{
    LineCol, MonikerDescriptorKind, StaticIndex, StaticIndexedFile, TextRange, TokenId,
    TokenStaticData,
};
use ide_db::LineIndexDatabase;
use project_model::{CargoConfig, ProjectManifest, ProjectWorkspace};
use scip::types as scip_types;
use std::env;

use crate::cli::{
    flags,
    load_cargo::{load_workspace, LoadCargoConfig},
    Result,
};

impl flags::Scip {
    pub fn run(self) -> Result<()> {
        eprintln!("Generating SCIP start...");
        let now = Instant::now();
        let cargo_config = CargoConfig::default();

        let no_progress = &|s| (eprintln!("rust-analyzer: Loading {}", s));
        let load_cargo_config = LoadCargoConfig {
            load_out_dirs_from_check: true,
            with_proc_macro: true,
            prefill_caches: true,
        };
        let path = vfs::AbsPathBuf::assert(env::current_dir()?.join(&self.path));
        let rootpath = path.normalize();
        let manifest = ProjectManifest::discover_single(&path)?;

        let workspace = ProjectWorkspace::load(manifest, &cargo_config, no_progress)?;

        let (host, vfs, _) =
            load_workspace(workspace, &cargo_config.extra_env, &load_cargo_config)?;
        let db = host.raw_database();
        let analysis = host.analysis();

        let si = StaticIndex::compute(&analysis);

        let mut index = scip_types::Index {
            metadata: Some(scip_types::Metadata {
                version: scip_types::ProtocolVersion::UnspecifiedProtocolVersion.into(),
                tool_info: Some(scip_types::ToolInfo {
                    name: "rust-analyzer".to_owned(),
                    version: "0.1".to_owned(),
                    arguments: vec![],
                    ..Default::default()
                })
                .into(),
                project_root: format!(
                    "file://{}",
                    path.normalize()
                        .as_os_str()
                        .to_str()
                        .ok_or(anyhow::anyhow!("Unable to normalize project_root path"))?
                        .to_string()
                ),
                text_document_encoding: scip_types::TextEncoding::UTF8.into(),
                ..Default::default()
            })
            .into(),
            ..Default::default()
        };

        let mut symbols_emitted: HashSet<TokenId> = HashSet::default();
        let mut tokens_to_symbol: HashMap<TokenId, String> = HashMap::new();

        for StaticIndexedFile { file_id, tokens, .. } in si.files {
            let mut local_count = 0;
            let mut new_local_symbol = || {
                let new_symbol = scip::types::Symbol::new_local(local_count);
                local_count += 1;

                new_symbol
            };

            let relative_path = match get_relative_filepath(&vfs, &rootpath, file_id) {
                Some(relative_path) => relative_path,
                None => continue,
            };

            let line_index = LineIndex {
                index: db.line_index(file_id),
                encoding: OffsetEncoding::Utf8,
                endings: LineEndings::Unix,
            };

            let mut doc = scip_types::Document {
                relative_path,
                language: "rust".to_string(),
                ..Default::default()
            };

            tokens.into_iter().for_each(|(range, id)| {
                let token = si.tokens.get(id).unwrap();

                let mut occurrence = scip_types::Occurrence::default();
                occurrence.range = text_range_to_scip_range(&line_index, range);
                occurrence.symbol = tokens_to_symbol
                    .entry(id)
                    .or_insert_with(|| {
                        let symbol = token_to_symbol(&token).unwrap_or_else(&mut new_local_symbol);
                        scip::symbol::format_symbol(symbol)
                    })
                    .clone();

                if let Some(def) = token.definition {
                    if def.range == range {
                        occurrence.symbol_roles |= scip_types::SymbolRole::Definition as i32;
                    }

                    if symbols_emitted.insert(id) {
                        let mut symbol_info = scip_types::SymbolInformation::default();
                        symbol_info.symbol = occurrence.symbol.clone();
                        if let Some(hover) = &token.hover {
                            if !hover.markup.as_str().is_empty() {
                                symbol_info.documentation = vec![hover.markup.as_str().to_string()];
                            }
                        }

                        doc.symbols.push(symbol_info)
                    }
                }

                doc.occurrences.push(occurrence);
            });

            if doc.occurrences.is_empty() {
                continue;
            }

            index.documents.push(doc);
        }

        scip::write_message_to_file("index.scip", index)
            .map_err(|err| anyhow::anyhow!("Failed to write scip to file: {}", err))?;

        eprintln!("Generating SCIP finished {:?}", now.elapsed());
        Ok(())
    }
}

fn get_relative_filepath(
    vfs: &vfs::Vfs,
    rootpath: &vfs::AbsPathBuf,
    file_id: ide::FileId,
) -> Option<String> {
    Some(vfs.file_path(file_id).as_path()?.strip_prefix(&rootpath)?.as_ref().to_str()?.to_string())
}

// SCIP Ranges have a (very large) optimization that ranges if they are on the same line
// only encode as a vector of [start_line, start_col, end_col].
//
// This transforms a line index into the optimized SCIP Range.
fn text_range_to_scip_range(line_index: &LineIndex, range: TextRange) -> Vec<i32> {
    let LineCol { line: start_line, col: start_col } = line_index.index.line_col(range.start());
    let LineCol { line: end_line, col: end_col } = line_index.index.line_col(range.end());

    if start_line == end_line {
        vec![start_line as i32, start_col as i32, end_col as i32]
    } else {
        vec![start_line as i32, start_col as i32, end_line as i32, end_col as i32]
    }
}

fn new_descriptor_str(
    name: &str,
    suffix: scip_types::descriptor::Suffix,
) -> scip_types::Descriptor {
    scip_types::Descriptor {
        name: name.to_string(),
        disambiguator: "".to_string(),
        suffix: suffix.into(),
        ..Default::default()
    }
}

fn new_descriptor(name: Name, suffix: scip_types::descriptor::Suffix) -> scip_types::Descriptor {
    let mut name = name.to_string();
    if name.contains("'") {
        name = format!("`{}`", name);
    }

    new_descriptor_str(name.as_str(), suffix)
}

/// Loosely based on `def_to_moniker`
///
/// Only returns a Symbol when it's a non-local symbol.
///     So if the visibility isn't outside of a document, then it will return None
fn token_to_symbol(token: &TokenStaticData) -> Option<scip_types::Symbol> {
    use scip_types::descriptor::Suffix::*;

    let moniker = token.moniker.as_ref()?;

    let package_name = moniker.package_information.name.clone();
    let version = moniker.package_information.version.clone();
    let descriptors = moniker
        .identifier
        .description
        .iter()
        .map(|desc| {
            new_descriptor(
                desc.name.clone(),
                match desc.desc {
                    MonikerDescriptorKind::Namespace => Namespace,
                    MonikerDescriptorKind::Type => Type,
                    MonikerDescriptorKind::Term => Term,
                    MonikerDescriptorKind::Method => Method,
                    MonikerDescriptorKind::TypeParameter => TypeParameter,
                    MonikerDescriptorKind::Parameter => Parameter,
                    MonikerDescriptorKind::Macro => Macro,
                    MonikerDescriptorKind::Meta => Meta,
                },
            )
        })
        .collect();

    Some(scip_types::Symbol {
        scheme: "rust-analyzer".into(),
        package: Some(scip_types::Package {
            manager: "cargo".to_string(),
            name: package_name,
            version,
            ..Default::default()
        })
        .into(),
        descriptors,
        ..Default::default()
    })
}

#[cfg(test)]
mod test {
    use super::*;
    use ide::{AnalysisHost, FilePosition, StaticIndex, TextSize};
    use ide_db::base_db::fixture::ChangeFixture;
    use scip::symbol::format_symbol;

    fn position(ra_fixture: &str) -> (AnalysisHost, FilePosition) {
        let mut host = AnalysisHost::default();
        let change_fixture = ChangeFixture::parse(ra_fixture);
        host.raw_database_mut().apply_change(change_fixture.change);
        let (file_id, range_or_offset) =
            change_fixture.file_position.expect("expected a marker ($0)");
        let offset = range_or_offset.expect_offset();
        (host, FilePosition { file_id, offset })
    }

    /// If expected == "", then assert that there are no symbols (this is basically local symbol)
    #[track_caller]
    fn check_symbol(ra_fixture: &str, expected: &str) {
        let (host, position) = position(ra_fixture);

        let analysis = host.analysis();
        let si = StaticIndex::compute(&analysis);

        let FilePosition { file_id, offset } = position;

        let mut found_symbol = None;
        for file in &si.files {
            if file.file_id != file_id {
                continue;
            }
            for &(range, id) in &file.tokens {
                if range.contains(offset - TextSize::from(1)) {
                    let token = si.tokens.get(id).unwrap();
                    found_symbol = token_to_symbol(token);
                    break;
                }
            }
        }

        if expected == "" {
            assert!(found_symbol.is_none(), "must have no symbols {:?}", found_symbol);
            return;
        }

        assert!(found_symbol.is_some(), "must have one symbol {:?}", found_symbol);
        let res = found_symbol.unwrap();
        let formatted = format_symbol(res);
        assert_eq!(formatted, expected);
    }

    #[test]
    fn basic() {
        check_symbol(
            r#"
//- /lib.rs crate:main deps:foo
use foo::example_mod::func;
fn main() {
    func$0();
}
//- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
pub mod example_mod {
    pub fn func() {}
}
"#,
            "rust-analyzer cargo foo 0.1.0 example_mod/func().",
        );
    }

    #[test]
    fn symbol_for_trait() {
        check_symbol(
            r#"
//- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
pub mod module {
    pub trait MyTrait {
        pub fn func$0() {}
    }
}
"#,
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#func().",
        );
    }

    #[test]
    fn symbol_for_trait_constant() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            const MY_CONST$0: u8;
        }
    }
    "#,
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#MY_CONST.",
        );
    }

    #[test]
    fn symbol_for_trait_type() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            type MyType$0;
        }
    }
    "#,
            // "foo::module::MyTrait::MyType",
            "rust-analyzer cargo foo 0.1.0 module/MyTrait#[MyType]",
        );
    }

    #[test]
    fn symbol_for_trait_impl_function() {
        check_symbol(
            r#"
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub trait MyTrait {
            pub fn func() {}
        }

        struct MyStruct {}

        impl MyTrait for MyStruct {
            pub fn func$0() {}
        }
    }
    "#,
            // "foo::module::MyStruct::MyTrait::func",
            "rust-analyzer cargo foo 0.1.0 module/MyStruct#MyTrait#func().",
        );
    }

    #[test]
    fn symbol_for_field() {
        check_symbol(
            r#"
    //- /lib.rs crate:main deps:foo
    use foo::St;
    fn main() {
        let x = St { a$0: 2 };
    }
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub struct St {
        pub a: i32,
    }
    "#,
            "rust-analyzer cargo foo 0.1.0 St#a.",
        );
    }

    #[test]
    fn local_symbol_for_local() {
        check_symbol(
            r#"
    //- /lib.rs crate:main deps:foo
    use foo::module::func;
    fn main() {
        func();
    }
    //- /foo/lib.rs crate:foo@CratesIo:0.1.0,https://a.b/foo.git
    pub mod module {
        pub fn func() {
            let x$0 = 2;
        }
    }
    "#,
            "",
        );
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::path::Path;

use memmap::Mmap;

/// Controls the strategy used for determining when to use memory maps.
///
/// If a searcher is called in circumstances where it is possible to use memory
/// maps, and memory maps are enabled, then it will attempt to do so if it
/// believes it will make the search faster.
///
/// By default, memory maps are disabled.
#[derive(Clone, Debug)]
pub struct MmapChoice(MmapChoiceImpl);

#[derive(Clone, Debug)]
enum MmapChoiceImpl {
    Auto,
    Never,
}

impl Default for MmapChoice {
    fn default() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Never)
    }
}

impl MmapChoice {
    /// Use memory maps when they are believed to be advantageous.
    ///
    /// The heuristics used to determine whether to use a memory map or not
    /// may depend on many things, including but not limited to, file size
    /// and platform.
    ///
    /// If memory maps are unavailable or cannot be used for a specific input,
    /// then normal OS read calls are used instead.
    ///
    /// # Safety
    ///
    /// This constructor is not safe because there is no obvious way to
    /// encapsulate the safety of file backed memory maps on all platforms
    /// without simultaneously negating some or all of their benefits.
    ///
    /// The specific contract the caller is required to uphold isn't precise,
    /// but it basically amounts to something like, "the caller guarantees that
    /// the underlying file won't be mutated." This, of course, isn't feasible
    /// in many environments. However, command line tools may still decide to
    /// take the risk of, say, a `SIGBUS` occurring while attempting to read a
    /// memory map.
    pub unsafe fn auto() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Auto)
    }

    /// Never use memory maps, no matter what. This is the default.
    pub fn never() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Never)
    }

    /// Return a memory map if memory maps are enabled and if creating a
    /// memory from the given file succeeded and if memory maps are believed
    /// to be advantageous for performance.
    ///
    /// If this does attempt to open a memory map and it fails, then `None`
    /// is returned and the corresponding error (along with the file path, if
    /// present) is logged at the debug level.
    pub(crate) fn open(
        &self,
        file: &File,
        path: Option<&Path>,
    ) -> Option<Mmap> {
        if !self.is_enabled() {
            return None;
        }
        if cfg!(target_os = "macos") {
            // I guess memory maps on macOS aren't great. Should re-evaluate.
            return None;
        }
        // SAFETY: This is acceptable because the only way `MmapChoiceImpl` can
        // be `Auto` is if the caller invoked the `auto` constructor, which
        // is itself not safe. Thus, this is a propagation of the caller's
        // assertion that using memory maps is safe.
        match unsafe { Mmap::map(file) } {
            Ok(mmap) => Some(mmap),
            Err(err) => {
                if let Some(path) = path {
                    log::debug!(
                        "{}: failed to open memory map: {}",
                        path.display(),
                        err
                    );
                } else {
                    log::debug!("failed to open memory map: {}", err);
                }
                None
            }
        }
    }

    /// Whether this strategy may employ memory maps or not.
    pub(crate) fn is_enabled(&self) -> bool {
        match self.0 {
            MmapChoiceImpl::Auto => true,
            MmapChoiceImpl::Never => false,
        }
    }
}
<commit_msg>searcher: disable mmap searching on non-64 bit<commit_after>use std::fs::File;
use std::path::Path;

use memmap::Mmap;

/// Controls the strategy used for determining when to use memory maps.
///
/// If a searcher is called in circumstances where it is possible to use memory
/// maps, and memory maps are enabled, then it will attempt to do so if it
/// believes it will make the search faster.
///
/// By default, memory maps are disabled.
#[derive(Clone, Debug)]
pub struct MmapChoice(MmapChoiceImpl);

#[derive(Clone, Debug)]
enum MmapChoiceImpl {
    Auto,
    Never,
}

impl Default for MmapChoice {
    fn default() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Never)
    }
}

impl MmapChoice {
    /// Use memory maps when they are believed to be advantageous.
    ///
    /// The heuristics used to determine whether to use a memory map or not
    /// may depend on many things, including but not limited to, file size
    /// and platform.
    ///
    /// If memory maps are unavailable or cannot be used for a specific input,
    /// then normal OS read calls are used instead.
    ///
    /// # Safety
    ///
    /// This constructor is not safe because there is no obvious way to
    /// encapsulate the safety of file backed memory maps on all platforms
    /// without simultaneously negating some or all of their benefits.
    ///
    /// The specific contract the caller is required to uphold isn't precise,
    /// but it basically amounts to something like, "the caller guarantees that
    /// the underlying file won't be mutated." This, of course, isn't feasible
    /// in many environments. However, command line tools may still decide to
    /// take the risk of, say, a `SIGBUS` occurring while attempting to read a
    /// memory map.
    pub unsafe fn auto() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Auto)
    }

    /// Never use memory maps, no matter what. This is the default.
    pub fn never() -> MmapChoice {
        MmapChoice(MmapChoiceImpl::Never)
    }

    /// Return a memory map if memory maps are enabled and if creating a
    /// memory from the given file succeeded and if memory maps are believed
    /// to be advantageous for performance.
    ///
    /// If this does attempt to open a memory map and it fails, then `None`
    /// is returned and the corresponding error (along with the file path, if
    /// present) is logged at the debug level.
    pub(crate) fn open(
        &self,
        file: &File,
        path: Option<&Path>,
    ) -> Option<Mmap> {
        if !self.is_enabled() {
            return None;
        }
        if !cfg!(target_pointer_width = "64") {
            // For 32-bit systems, it looks like mmap will succeed even if it
            // can't address the entire file. This seems to happen at least on
            // Windows, even though it uses to work prior to ripgrep 13. The
            // only Windows-related change in ripgrep 13, AFAIK, was statically
            // linking vcruntime. So maybe that's related? But I'm not sure.
            //
            // See: https://github.com/BurntSushi/ripgrep/issues/1911
            return None;
        }
        if cfg!(target_os = "macos") {
            // I guess memory maps on macOS aren't great. Should re-evaluate.
            return None;
        }
        // SAFETY: This is acceptable because the only way `MmapChoiceImpl` can
        // be `Auto` is if the caller invoked the `auto` constructor, which
        // is itself not safe. Thus, this is a propagation of the caller's
        // assertion that using memory maps is safe.
        match unsafe { Mmap::map(file) } {
            Ok(mmap) => Some(mmap),
            Err(err) => {
                if let Some(path) = path {
                    log::debug!(
                        "{}: failed to open memory map: {}",
                        path.display(),
                        err
                    );
                } else {
                    log::debug!("failed to open memory map: {}", err);
                }
                None
            }
        }
    }

    /// Whether this strategy may employ memory maps or not.
    pub(crate) fn is_enabled(&self) -> bool {
        match self.0 {
            MmapChoiceImpl::Auto => true,
            MmapChoiceImpl::Never => false,
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use num_complex::Complex;
use num_traits::Zero;

use common::{FFTnum, verify_length_inline, verify_length_minimum};

use ::{Length, IsInverse, FftInline};

/// Implementation of Bluestein's Algorithm
///
/// This algorithm computes an arbitrary-sized FFT in O(nlogn) time. It does this by converting this size n FFT into a
/// size M FFT, where M >= 2N - 1. M is usually a power of two, although that isn't a requirement.
///
/// It requires a large scratch space, so it's probably inconvenient to use as an inner FFT to other algorithms.
///
/// ~~~
/// // Computes a forward FFT of size 1201 (prime number), using Bluestein's Algorithm
/// use rustfft::algorithm::RadersAlgorithm;
/// use rustfft::{FFT, FFTplanner};
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 1201];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 1201];
///
/// // plan a FFT of size n - 1 = 1200
/// let mut planner = FFTplanner::new(false);
/// let inner_fft = planner.plan_fft(1200);
///
/// let fft = RadersAlgorithm::new(1201, inner_fft);
/// fft.process(&mut input, &mut output);
/// ~~~
///
/// Rader's Algorithm is relatively expensive compared to other FFT algorithms. Benchmarking shows that it is up to
/// an order of magnitude slower than similar composite sizes. In the example size above of 1201, benchmarking shows
/// that it takes 2.5x more time to compute than a FFT of size 1200.

pub struct BluesteinsAlgorithm<T> {
    inner_fft: Arc<FftInline<T>>,

    inner_fft_multiplier: Box<[Complex<T>]>,
    twiddles: Box<[Complex<T>]>,

    len: usize,
}

impl<T: FFTnum> BluesteinsAlgorithm<T> {
    fn compute_bluesteins_twiddle(index: usize, size: usize, inverse: bool) -> Complex<T> {
        let index_multiplier = core::f64::consts::PI / size as f64;

        let index_float = index as f64;
        let index_squared = index_float * index_float;

        let theta = index_squared * index_multiplier;
        let result = Complex::new(
            T::from_f64(theta.cos()).unwrap(),
            T::from_f64(theta.sin()).unwrap(),
        );

        if inverse {
            result.conj()
        } else {
            result
        }
    }

    /// Creates a FFT instance which will process inputs/outputs of size `len`. `inner_fft.len()` must be >= `len * 2 - 1`
    ///
    /// Note that this constructor is quite expensive to run; This algorithm must run a FFT of size inner_fft.len() within the
    /// constructor. This further underlines the fact that Bluesteins Algorithm is more expensive to run than other
    /// FFT algorithms
    pub fn new(len: usize, inner_fft: Arc<FftInline<T>>) -> Self {
        let inner_fft_len = inner_fft.len();
        assert!(len * 2 - 1 <= inner_fft_len, "Bluestein's algorithm requires inner_fft.len() >= self.len() * 2 - 1. Expected >= {}, got {}", len * 2 - 1, inner_fft_len);

        // when computing FFTs, we're going to run our inner FFT, multiply pairise by some precomputed data, then run an inverse inner FFT. We need to precompute that inner data here
        let inner_len_float = T::from_usize(inner_fft_len).unwrap();
        let inverse = inner_fft.is_inverse();

        // Compute twiddle factors that we'll run our inner FFT on
        let mut inner_fft_input = vec![Complex::zero(); inner_fft_len];
        for i in 0..len {
            inner_fft_input[i] = Self::compute_bluesteins_twiddle(i, len, inverse) / inner_len_float;
        }
        for i in 1..len {
            inner_fft_input[inner_fft_len - i] = inner_fft_input[i];
        }

        //Compute the inner fft
        let mut inner_fft_scratch = vec![Complex::zero(); inner_fft.get_required_scratch_len()];
        inner_fft.process_inline(&mut inner_fft_input, &mut inner_fft_scratch);

        // also compute some more mundane twiddle factors to start and end with
        let twiddles : Vec<_> = (0..len).map(|i| Self::compute_bluesteins_twiddle(i, len, !inverse)).collect();

        Self {
            inner_fft: inner_fft,

            inner_fft_multiplier: inner_fft_input.into_boxed_slice(),
            twiddles: twiddles.into_boxed_slice(),

            len,
        }
    }

    #[target_feature(enable = "avx", enable = "fma")]
    unsafe fn perform_bluestein_fft(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
        let (inner_input, inner_scratch) = scratch.split_at_mut(self.inner_fft_multiplier.len());

        // Copy the buffer into our inner FFT input. the buffer will only fill part of the FFT input, so zero fill the rest
        for ((buffer_entry, inner_entry), twiddle) in buffer.iter().zip(inner_input.iter_mut()).zip(self.twiddles.iter()) {
            *inner_entry = *buffer_entry * *twiddle ;
        }
        for inner in inner_input.iter_mut().skip(buffer.len()) {
            *inner = Complex::zero();
        }

        // run our inner forward FFT
        self.inner_fft.process_inline(inner_input, inner_scratch);

        // Multiply our inner FFT output by our precomputed data. Then, conjugate the result to set up for an inverse FFT
        for (inner, multiplier) in inner_input.iter_mut().zip(self.inner_fft_multiplier.iter()) {
            *inner = (*inner * *multiplier).conj();
        }

        // inverse FFT. we're computing a forward FFT, but we're massaging it into an inverse by conjugating the inputs and outputs
        self.inner_fft.process_inline(inner_input, inner_scratch);

        // copy our data back to the buffer, applying twiddle factors again as we go. Also conjugate inner_input to complete the inverse FFT
        for ((buffer_entry, inner_entry), twiddle) in buffer.iter_mut().zip(inner_input.iter()).zip(self.twiddles.iter()) {
            *buffer_entry = inner_entry.conj() * twiddle;
        }
    }
}

impl<T: FFTnum> FftInline<T> for BluesteinsAlgorithm<T> {
    fn process_inline(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
        verify_length_inline(buffer, self.len());
        verify_length_minimum(scratch, self.get_required_scratch_len());

        unsafe { self.perform_bluestein_fft(buffer, scratch); }
    }
    fn get_required_scratch_len(&self) -> usize {
        self.inner_fft_multiplier.len() + self.inner_fft.get_required_scratch_len()
    }
}
impl<T> Length for BluesteinsAlgorithm<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.len
    }
}
impl<T> IsInverse for BluesteinsAlgorithm<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inner_fft.is_inverse()
    }
}

#[cfg(test)]
mod unit_tests {
    use super::*;
    use std::sync::Arc;
    use test_utils::check_inline_fft_algorithm;
    use algorithm::DFT;

    #[test]
    fn test_bluesteins() {
        for &len in &[3,5,7,11,13] {
            test_bluesteins_with_length(len, false);
            test_bluesteins_with_length(len, true);
        }
    }

    fn test_bluesteins_with_length(len: usize, inverse: bool) {
        let inner_fft = Arc::new(DFT::new((len *2 - 1).checked_next_power_of_two().unwrap(), inverse));
        let fft = BluesteinsAlgorithm::new(len, inner_fft);

        check_inline_fft_algorithm(&fft, len, inverse);
    }
}
<commit_msg>Removed accidental unsafe<commit_after>use std::sync::Arc;

use num_complex::Complex;
use num_traits::Zero;

use common::{FFTnum, verify_length_inline, verify_length_minimum};

use ::{Length, IsInverse, FftInline};

/// Implementation of Bluestein's Algorithm
///
/// This algorithm computes an arbitrary-sized FFT in O(nlogn) time. It does this by converting this size n FFT into a
/// size M FFT, where M >= 2N - 1. M is usually a power of two, although that isn't a requirement.
///
/// It requires a large scratch space, so it's probably inconvenient to use as an inner FFT to other algorithms.
///
/// ~~~
/// // Computes a forward FFT of size 1201 (prime number), using Bluestein's Algorithm
/// use rustfft::algorithm::RadersAlgorithm;
/// use rustfft::{FFT, FFTplanner};
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 1201];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 1201];
///
/// // plan a FFT of size n - 1 = 1200
/// let mut planner = FFTplanner::new(false);
/// let inner_fft = planner.plan_fft(1200);
///
/// let fft = RadersAlgorithm::new(1201, inner_fft);
/// fft.process(&mut input, &mut output);
/// ~~~
///
/// Rader's Algorithm is relatively expensive compared to other FFT algorithms. Benchmarking shows that it is up to
/// an order of magnitude slower than similar composite sizes. In the example size above of 1201, benchmarking shows
/// that it takes 2.5x more time to compute than a FFT of size 1200.

pub struct BluesteinsAlgorithm<T> {
    inner_fft: Arc<FftInline<T>>,

    inner_fft_multiplier: Box<[Complex<T>]>,
    twiddles: Box<[Complex<T>]>,

    len: usize,
}

impl<T: FFTnum> BluesteinsAlgorithm<T> {
    fn compute_bluesteins_twiddle(index: usize, size: usize, inverse: bool) -> Complex<T> {
        let index_multiplier = core::f64::consts::PI / size as f64;

        let index_float = index as f64;
        let index_squared = index_float * index_float;

        let theta = index_squared * index_multiplier;
        let result = Complex::new(
            T::from_f64(theta.cos()).unwrap(),
            T::from_f64(theta.sin()).unwrap(),
        );

        if inverse {
            result.conj()
        } else {
            result
        }
    }

    /// Creates a FFT instance which will process inputs/outputs of size `len`. `inner_fft.len()` must be >= `len * 2 - 1`
    ///
    /// Note that this constructor is quite expensive to run; This algorithm must run a FFT of size inner_fft.len() within the
    /// constructor. This further underlines the fact that Bluesteins Algorithm is more expensive to run than other
    /// FFT algorithms
    pub fn new(len: usize, inner_fft: Arc<FftInline<T>>) -> Self {
        let inner_fft_len = inner_fft.len();
        assert!(len * 2 - 1 <= inner_fft_len, "Bluestein's algorithm requires inner_fft.len() >= self.len() * 2 - 1. Expected >= {}, got {}", len * 2 - 1, inner_fft_len);

        // when computing FFTs, we're going to run our inner FFT, multiply pairise by some precomputed data, then run an inverse inner FFT. We need to precompute that inner data here
        let inner_len_float = T::from_usize(inner_fft_len).unwrap();
        let inverse = inner_fft.is_inverse();

        // Compute twiddle factors that we'll run our inner FFT on
        let mut inner_fft_input = vec![Complex::zero(); inner_fft_len];
        for i in 0..len {
            inner_fft_input[i] = Self::compute_bluesteins_twiddle(i, len, inverse) / inner_len_float;
        }
        for i in 1..len {
            inner_fft_input[inner_fft_len - i] = inner_fft_input[i];
        }

        //Compute the inner fft
        let mut inner_fft_scratch = vec![Complex::zero(); inner_fft.get_required_scratch_len()];
        inner_fft.process_inline(&mut inner_fft_input, &mut inner_fft_scratch);

        // also compute some more mundane twiddle factors to start and end with
        let twiddles : Vec<_> = (0..len).map(|i| Self::compute_bluesteins_twiddle(i, len, !inverse)).collect();

        Self {
            inner_fft: inner_fft,

            inner_fft_multiplier: inner_fft_input.into_boxed_slice(),
            twiddles: twiddles.into_boxed_slice(),

            len,
        }
    }

    fn perform_bluestein_fft(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
        let (inner_input, inner_scratch) = scratch.split_at_mut(self.inner_fft_multiplier.len());

        // Copy the buffer into our inner FFT input. the buffer will only fill part of the FFT input, so zero fill the rest
        for ((buffer_entry, inner_entry), twiddle) in buffer.iter().zip(inner_input.iter_mut()).zip(self.twiddles.iter()) {
            *inner_entry = *buffer_entry * *twiddle ;
        }
        for inner in inner_input.iter_mut().skip(buffer.len()) {
            *inner = Complex::zero();
        }

        // run our inner forward FFT
        self.inner_fft.process_inline(inner_input, inner_scratch);

        // Multiply our inner FFT output by our precomputed data. Then, conjugate the result to set up for an inverse FFT
        for (inner, multiplier) in inner_input.iter_mut().zip(self.inner_fft_multiplier.iter()) {
            *inner = (*inner * *multiplier).conj();
        }

        // inverse FFT. we're computing a forward FFT, but we're massaging it into an inverse by conjugating the inputs and outputs
        self.inner_fft.process_inline(inner_input, inner_scratch);

        // copy our data back to the buffer, applying twiddle factors again as we go. Also conjugate inner_input to complete the inverse FFT
        for ((buffer_entry, inner_entry), twiddle) in buffer.iter_mut().zip(inner_input.iter()).zip(self.twiddles.iter()) {
            *buffer_entry = inner_entry.conj() * twiddle;
        }
    }
}

impl<T: FFTnum> FftInline<T> for BluesteinsAlgorithm<T> {
    fn process_inline(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
        verify_length_inline(buffer, self.len());
        verify_length_minimum(scratch, self.get_required_scratch_len());

        self.perform_bluestein_fft(buffer, scratch);
    }
    fn get_required_scratch_len(&self) -> usize {
        self.inner_fft_multiplier.len() + self.inner_fft.get_required_scratch_len()
    }
}
impl<T> Length for BluesteinsAlgorithm<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.len
    }
}
impl<T> IsInverse for BluesteinsAlgorithm<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inner_fft.is_inverse()
    }
}

#[cfg(test)]
mod unit_tests {
    use super::*;
    use std::sync::Arc;
    use test_utils::check_inline_fft_algorithm;
    use algorithm::DFT;

    #[test]
    fn test_bluesteins() {
        for &len in &[3,5,7,11,13] {
            test_bluesteins_with_length(len, false);
            test_bluesteins_with_length(len, true);
        }
    }

    fn test_bluesteins_with_length(len: usize, inverse: bool) {
        let inner_fft = Arc::new(DFT::new((len *2 - 1).checked_next_power_of_two().unwrap(), inverse));
        let fft = BluesteinsAlgorithm::new(len, inner_fft);

        check_inline_fft_algorithm(&fft, len, inverse);
    }
}
<|endoftext|>


<commit_before>use crate::config::SslMode;
use crate::tls::TlsConnect;
use crate::{connect_tls, Error};
use futures::future;
use postgres_protocol::message::frontend;
use std::pin::Pin;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};

pub async fn cancel_query_raw<S, T>(
    stream: S,
    mode: SslMode,
    tls: T,
    process_id: i32,
    secret_key: i32,
) -> Result<(), Error>
where
    S: AsyncRead + AsyncWrite + Unpin,
    T: TlsConnect<S>,
{
    let (mut stream, _) = connect_tls::connect_tls(stream, mode, tls).await?;

    let mut buf = vec![];
    frontend::cancel_request(process_id, secret_key, &mut buf);

    stream.write_all(&buf).await.map_err(Error::io)?;
    future::poll_fn(|cx| Pin::new(&mut stream).poll_flush(cx))
        .await
        .map_err(Error::io)?;
    future::poll_fn(|cx| Pin::new(&mut stream).poll_shutdown(cx))
        .await
        .map_err(Error::io)?;

    Ok(())
}
<commit_msg>Use flush combinator<commit_after>use crate::config::SslMode;
use crate::tls::TlsConnect;
use crate::{connect_tls, Error};
use futures::future;
use postgres_protocol::message::frontend;
use std::pin::Pin;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};

pub async fn cancel_query_raw<S, T>(
    stream: S,
    mode: SslMode,
    tls: T,
    process_id: i32,
    secret_key: i32,
) -> Result<(), Error>
where
    S: AsyncRead + AsyncWrite + Unpin,
    T: TlsConnect<S>,
{
    let (mut stream, _) = connect_tls::connect_tls(stream, mode, tls).await?;

    let mut buf = vec![];
    frontend::cancel_request(process_id, secret_key, &mut buf);

    stream.write_all(&buf).await.map_err(Error::io)?;
    stream.flush().await.map_err(Error::io)?;
    future::poll_fn(|cx| Pin::new(&mut stream).poll_shutdown(cx))
        .await
        .map_err(Error::io)?;

    Ok(())
}
<|endoftext|>