<commit_before>//! Contains the `SwitchPacket` structure, storing the packet header
//! used by the Switch, as defined by
//! https://github.com/cjdelisle/cjdns/blob/cjdns-v18/doc/Whitepaper.md#in-memory-representation

use byteorder::BigEndian;
use byteorder::ByteOrder;

use operation::{switch, reverse_label, Director, RoutingDecision, Label, ForwardPath, BackwardPath, LABEL_LENGTH};
use packets::control::ControlPacket;
use session_manager::SessionHandle;

#[derive(Debug)]
pub enum Payload {
    Control(ControlPacket),
    CryptoAuthHello(Vec<u8>),
    CryptoAuthKey(Vec<u8>),
    CryptoAuthData(SessionHandle, Vec<u8>),
}

#[derive(Debug)]
pub struct SwitchPacket {
    pub raw: Vec<u8>,
}

impl SwitchPacket {
    /// Returns a new packet, constructed from its route and its payload.
    #[cfg(not(feature="sfcp"))]
    pub fn new(path: ForwardPath, payload: Payload) -> SwitchPacket {
        let mut raw = vec![0u8; 12];
        raw[0..8].copy_from_slice(&Label::from(path));
        match payload {
            Payload::Control(msg) => {
                raw.append(&mut vec![0xff, 0xff, 0xff, 0xff]);
                raw.append(&mut msg.encode());
            },
            Payload::CryptoAuthHello(mut msg) | Payload::CryptoAuthKey(mut msg) => {
                let session_state = BigEndian::read_u32(&msg[0..4]);
                assert!(session_state < 4);
                assert!(session_state != 0xffffffff);
                raw.append(&mut msg);
            },
            Payload::CryptoAuthData(session_handle, mut msg) => {
                assert!(session_handle.0 >= 4);
                assert!(session_handle.0 != 0xffffffff);
                let mut raw_handle = vec![0u8; 4];
                BigEndian::write_u32(&mut raw_handle, session_handle.0);
                raw.append(&mut raw_handle);
                raw.append(&mut msg);
            },
        }
        SwitchPacket { raw: raw }
    }

    /// Returns a new packet, constructed from its route and its payload.
    #[cfg(feature="sfcp")]
    pub fn new(path: ForwardPath, payload: Payload) -> SwitchPacket {
        let mut raw = vec![0u8; 20];
        raw[0..16].copy_from_slice(&Label::from(path));
        match payload {
            Payload::Control(msg) => {
                raw.append(&mut vec![0xff, 0xff, 0xff, 0xff]);
                raw.append(&mut msg.encode());
            },
            Payload::CryptoAuthHello(mut msg) | Payload::CryptoAuthKey(mut msg) => {
                let session_state = BigEndian::read_u32(&msg[0..4]);
                assert!(session_state < 4);
                assert!(session_state != 0xffffffff);
                raw.append(&mut msg);
            },
            Payload::CryptoAuthData(session_handle, mut msg) => {
                assert!(session_handle.0 >= 4);
                assert!(session_handle.0 != 0xffffffff);
                let mut raw_handle = vec![0u8; 4];
                BigEndian::write_u32(&mut raw_handle, session_handle.0);
                raw.append(&mut raw_handle);
                raw.append(&mut msg);
            },
        }
        SwitchPacket { raw: raw }
    }

    /// Returns a new packet, constructed as a reply of a received one.
    pub fn new_reply(received: &SwitchPacket, payload: Payload) -> SwitchPacket {
        let mut path = BackwardPath::from(received.label()).reverse();
        SwitchPacket::new(path, payload)
    }

    /// Returns the address label of the packet.
    pub fn label(&self) -> Label {
        let mut label = [0u8; LABEL_LENGTH];
        label.copy_from_slice(&self.raw[0..LABEL_LENGTH]);
        label
    }

    #[cfg(not(feature="sfcp"))]
    pub fn congest(&self) -> u8 {
        self.raw[8] >> 1
    }

    #[cfg(not(feature="sfcp"))]
    pub fn suppress_errors(&self) -> bool {
        self.raw[8] & 0b00000001 == 1
    }

    #[cfg(not(feature="sfcp"))]
    pub fn version(&self) -> u8 {
        self.raw[9] >> 6
    }

    #[cfg(not(feature="sfcp"))]
    pub fn label_shift(&self) -> u8 {
        self.raw[9] & 0b00111111
    }

    #[cfg(not(feature="sfcp"))]
    pub fn penalty(&self) -> [u8; 2] {
        let mut a = [0u8; 2];
        a.copy_from_slice(&self.raw[10..12]);
        a
    }

    /// Returns a reference to the content of the packet.
    #[cfg(not(feature="sfcp"))]
    pub fn payload(&self) -> Option<Payload> {
        match BigEndian::read_u32(&self.raw[12..16]) {
            0xffffffff => ControlPacket::decode(&self.raw[16..].to_vec()).map(Payload::Control),
            0 | 1 => Some(Payload::CryptoAuthHello(self.raw[12..].to_vec())),
            2 | 3 => Some(Payload::CryptoAuthKey(self.raw[12..].to_vec())),
            handle => Some(Payload::CryptoAuthData(SessionHandle(handle), self.raw[16..].to_vec())),
        }
    }

    /// Returns a reference to the content of the packet.
    #[cfg(feature="sfcp")]
    pub fn payload(&self) -> Option<Payload> {
        match BigEndian::read_u32(&self.raw[20..24]) {
            0xffffffff => ControlPacket::decode(&self.raw[24..].to_vec()).map(Payload::Control),
            0 | 1 => Some(Payload::CryptoAuthHello(self.raw[20..].to_vec())),
            2 | 3 => Some(Payload::CryptoAuthKey(self.raw[20..].to_vec())),
            handle => Some(Payload::CryptoAuthData(SessionHandle(handle), self.raw[20..].to_vec())),
        }
    }

    /// Make this packet advance one logical hop.
    ///
    /// Using the Director Length of this switch, determines what interface
    /// the packet will go next, update the label of the packet to take that
    /// switching into account (pops the interface from the path, and puts it
    /// on the reverse path), then returns the interface.
    ///
    /// See the doc of `fcp::operation::switch` for more details.
    pub fn switch(&mut self, director_length: u8, reversed_origin_iface: &Director) -> RoutingDecision {
        let (new_label, decision) = switch(&self.label(), director_length, reversed_origin_iface);
        self.raw[0..LABEL_LENGTH].copy_from_slice(&new_label);
        decision
    }

    /// Inverses the path and the return path.
    pub fn reverse_label(&mut self) {
        // TODO: do this in-place/no-copy.
        let mut label = [0u8; LABEL_LENGTH];
        reverse_label(&mut label);
        self.raw[0..LABEL_LENGTH].copy_from_slice(&label);
    }

}

#[cfg(test)]
mod test {
    use hex::FromHex;
    use super::*;
    use super::super::super::operation::RoutingDecision;
    use super::super::super::packets::control::ControlPacket;

    #[test]
    #[cfg(not(feature="sfcp"))]
    fn switch_and_reply() {
        let mut received = SwitchPacket { raw: Vec::from_hex("800000000000000100440000ffffffff9986000309f9110200000011467c6febbde26264a38cd12e").unwrap() };
        let decision = received.switch(4, &0b1100);
        let opaque_data = match decision {
            RoutingDecision::SelfInterface(_) => {
                match received.payload() {
                    Some(Payload::Control(ControlPacket::Ping { opaque_data, .. })) => {
                        opaque_data
                    },
                    _ => panic!("parsed as non-Ping."),
                }
            }
            _ => panic!("routed to non-self interface."),
        };
        let control_response = ControlPacket::Pong { version: 17, opaque_data: opaque_data };
        let mut response = SwitchPacket::new_reply(&received, Payload::Control(control_response));
        let decision = response.switch(4, &0b1000);
        assert_eq!(decision, RoutingDecision::Forward(0b0011));
        assert_eq!(response.raw, Vec::from_hex("800000000000000100000000ffffffff33b000049d74e35b00000011467c6febbde26264a38cd12e").unwrap());
    }
}
<commit_msg>Parse the labelShift in SFCP.<commit_after>//! Contains the `SwitchPacket` structure, storing the packet header
//! used by the Switch, as defined by
//! https://github.com/cjdelisle/cjdns/blob/cjdns-v18/doc/Whitepaper.md#in-memory-representation

use byteorder::BigEndian;
use byteorder::ByteOrder;

use operation::{switch, reverse_label, Director, RoutingDecision, Label, ForwardPath, BackwardPath, LABEL_LENGTH};
use packets::control::ControlPacket;
use session_manager::SessionHandle;

#[derive(Debug)]
pub enum Payload {
    Control(ControlPacket),
    CryptoAuthHello(Vec<u8>),
    CryptoAuthKey(Vec<u8>),
    CryptoAuthData(SessionHandle, Vec<u8>),
}

#[derive(Debug)]
pub struct SwitchPacket {
    pub raw: Vec<u8>,
}

impl SwitchPacket {
    /// Returns a new packet, constructed from its route and its payload.
    #[cfg(not(feature="sfcp"))]
    pub fn new(path: ForwardPath, payload: Payload) -> SwitchPacket {
        let mut raw = vec![0u8; 12];
        raw[0..8].copy_from_slice(&Label::from(path));
        match payload {
            Payload::Control(msg) => {
                raw.append(&mut vec![0xff, 0xff, 0xff, 0xff]);
                raw.append(&mut msg.encode());
            },
            Payload::CryptoAuthHello(mut msg) | Payload::CryptoAuthKey(mut msg) => {
                let session_state = BigEndian::read_u32(&msg[0..4]);
                assert!(session_state < 4);
                assert!(session_state != 0xffffffff);
                raw.append(&mut msg);
            },
            Payload::CryptoAuthData(session_handle, mut msg) => {
                assert!(session_handle.0 >= 4);
                assert!(session_handle.0 != 0xffffffff);
                let mut raw_handle = vec![0u8; 4];
                BigEndian::write_u32(&mut raw_handle, session_handle.0);
                raw.append(&mut raw_handle);
                raw.append(&mut msg);
            },
        }
        SwitchPacket { raw: raw }
    }

    /// Returns a new packet, constructed from its route and its payload.
    #[cfg(feature="sfcp")]
    pub fn new(path: ForwardPath, payload: Payload) -> SwitchPacket {
        let mut raw = vec![0u8; 20];
        raw[0..16].copy_from_slice(&Label::from(path));
        match payload {
            Payload::Control(msg) => {
                raw.append(&mut vec![0xff, 0xff, 0xff, 0xff]);
                raw.append(&mut msg.encode());
            },
            Payload::CryptoAuthHello(mut msg) | Payload::CryptoAuthKey(mut msg) => {
                let session_state = BigEndian::read_u32(&msg[0..4]);
                assert!(session_state < 4);
                assert!(session_state != 0xffffffff);
                raw.append(&mut msg);
            },
            Payload::CryptoAuthData(session_handle, mut msg) => {
                assert!(session_handle.0 >= 4);
                assert!(session_handle.0 != 0xffffffff);
                let mut raw_handle = vec![0u8; 4];
                BigEndian::write_u32(&mut raw_handle, session_handle.0);
                raw.append(&mut raw_handle);
                raw.append(&mut msg);
            },
        }
        SwitchPacket { raw: raw }
    }

    /// Returns a new packet, constructed as a reply of a received one.
    pub fn new_reply(received: &SwitchPacket, payload: Payload) -> SwitchPacket {
        let mut path = BackwardPath::from(received.label()).reverse();
        SwitchPacket::new(path, payload)
    }

    /// Returns the address label of the packet.
    pub fn label(&self) -> Label {
        let mut label = [0u8; LABEL_LENGTH];
        label.copy_from_slice(&self.raw[0..LABEL_LENGTH]);
        label
    }

    #[cfg(not(feature="sfcp"))]
    pub fn congest(&self) -> u8 {
        self.raw[8] >> 1
    }

    #[cfg(not(feature="sfcp"))]
    pub fn suppress_errors(&self) -> bool {
        self.raw[8] & 0b00000001 == 1
    }

    #[cfg(feature="sfcp")]
    pub fn suppress_errors(&self) -> bool {
        self.raw[17] & 0b10000000 == 1
    }

    #[cfg(not(feature="sfcp"))]
    pub fn version(&self) -> u8 {
        self.raw[9] >> 6
    }

    #[cfg(not(feature="sfcp"))]
    pub fn label_shift(&self) -> u8 {
        self.raw[9] & 0b00111111
    }

    #[cfg(feature="sfcp")]
    pub fn label_shift(&self) -> u8 {
        self.raw[17] & 0b01111111
    }

    #[cfg(not(feature="sfcp"))]
    pub fn penalty(&self) -> [u8; 2] {
        let mut a = [0u8; 2];
        a.copy_from_slice(&self.raw[10..12]);
        a
    }

    /// Returns a reference to the content of the packet.
    #[cfg(not(feature="sfcp"))]
    pub fn payload(&self) -> Option<Payload> {
        match BigEndian::read_u32(&self.raw[12..16]) {
            0xffffffff => ControlPacket::decode(&self.raw[16..].to_vec()).map(Payload::Control),
            0 | 1 => Some(Payload::CryptoAuthHello(self.raw[12..].to_vec())),
            2 | 3 => Some(Payload::CryptoAuthKey(self.raw[12..].to_vec())),
            handle => Some(Payload::CryptoAuthData(SessionHandle(handle), self.raw[16..].to_vec())),
        }
    }

    /// Returns a reference to the content of the packet.
    #[cfg(feature="sfcp")]
    pub fn payload(&self) -> Option<Payload> {
        match BigEndian::read_u32(&self.raw[20..24]) {
            0xffffffff => ControlPacket::decode(&self.raw[24..].to_vec()).map(Payload::Control),
            0 | 1 => Some(Payload::CryptoAuthHello(self.raw[20..].to_vec())),
            2 | 3 => Some(Payload::CryptoAuthKey(self.raw[20..].to_vec())),
            handle => Some(Payload::CryptoAuthData(SessionHandle(handle), self.raw[20..].to_vec())),
        }
    }

    /// Make this packet advance one logical hop.
    ///
    /// Using the Director Length of this switch, determines what interface
    /// the packet will go next, update the label of the packet to take that
    /// switching into account (pops the interface from the path, and puts it
    /// on the reverse path), then returns the interface.
    ///
    /// See the doc of `fcp::operation::switch` for more details.
    pub fn switch(&mut self, director_length: u8, reversed_origin_iface: &Director) -> RoutingDecision {
        let (new_label, decision) = switch(&self.label(), director_length, reversed_origin_iface);
        self.raw[0..LABEL_LENGTH].copy_from_slice(&new_label);
        decision
    }

    /// Inverses the path and the return path.
    pub fn reverse_label(&mut self) {
        // TODO: do this in-place/no-copy.
        let mut label = [0u8; LABEL_LENGTH];
        reverse_label(&mut label);
        self.raw[0..LABEL_LENGTH].copy_from_slice(&label);
    }

}

#[cfg(test)]
mod test {
    use hex::FromHex;
    use super::*;
    use super::super::super::operation::RoutingDecision;
    use super::super::super::packets::control::ControlPacket;

    #[test]
    #[cfg(not(feature="sfcp"))]
    fn switch_and_reply() {
        let mut received = SwitchPacket { raw: Vec::from_hex("800000000000000100440000ffffffff9986000309f9110200000011467c6febbde26264a38cd12e").unwrap() };
        let decision = received.switch(4, &0b1100);
        let opaque_data = match decision {
            RoutingDecision::SelfInterface(_) => {
                match received.payload() {
                    Some(Payload::Control(ControlPacket::Ping { opaque_data, .. })) => {
                        opaque_data
                    },
                    _ => panic!("parsed as non-Ping."),
                }
            }
            _ => panic!("routed to non-self interface."),
        };
        let control_response = ControlPacket::Pong { version: 17, opaque_data: opaque_data };
        let mut response = SwitchPacket::new_reply(&received, Payload::Control(control_response));
        let decision = response.switch(4, &0b1000);
        assert_eq!(decision, RoutingDecision::Forward(0b0011));
        assert_eq!(response.raw, Vec::from_hex("800000000000000100000000ffffffff33b000049d74e35b00000011467c6febbde26264a38cd12e").unwrap());
    }
}
<|endoftext|>


<commit_before>// This module implements the http server support for our application.
// Copyright (c) 2014 by Shipeng Feng.
// Licensed under the BSD License, see LICENSE for more details.

use std::io::net::ip::{SocketAddr, Ipv4Addr};

use http;
use http::server::{Server, Request, ResponseWriter};
use http::headers::content_type::MediaType;

use app::Pencil;

/// The pencil server.
#[deriving(Clone)]
struct PencilServer {
    app: Pencil,
}

impl PencilServer {
    /// Create a `PencilServer`.
    pub fn new(application: Pencil) -> PencilServer {
        PencilServer { app: application }
    }
}

impl Server for PencilServer {
    fn get_config(&self) -> http::server::Config {
        http::server::Config { bind_address: SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 8000 } }
    }

    fn handle_request(&self, r: Request, w: &mut ResponseWriter) {
        // let request = r;
        let response = self.app.handle_request(r);

        w.headers.content_type = Some(MediaType {
            type_ : String::from_str("text"),
            subtype: String::from_str("plain"),
            parameters: vec!((String::from_str("charset"), String::from_str("UTF-8")))
        });
        w.headers.server = Some(String::from_str("Pencil"));
        w.write(response.body.as_bytes()).unwrap();

        self.app.do_teardown_request();
    }
}

/// Run the `Pencil` application.
pub fn run_server(application: Pencil) {
    let pencil_server = PencilServer::new(application);
    pencil_server.serve_forever();
}
<commit_msg>Use text/html for serving<commit_after>// This module implements the http server support for our application.
// Copyright (c) 2014 by Shipeng Feng.
// Licensed under the BSD License, see LICENSE for more details.

use std::io::net::ip::{SocketAddr, Ipv4Addr};

use http;
use http::server::{Server, Request, ResponseWriter};
use http::headers::content_type::MediaType;

use app::Pencil;

/// The pencil server.
#[deriving(Clone)]
struct PencilServer {
    app: Pencil,
}

impl PencilServer {
    /// Create a `PencilServer`.
    pub fn new(application: Pencil) -> PencilServer {
        PencilServer { app: application }
    }
}

impl Server for PencilServer {
    fn get_config(&self) -> http::server::Config {
        http::server::Config { bind_address: SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 8000 } }
    }

    fn handle_request(&self, r: Request, w: &mut ResponseWriter) {
        // let request = r;
        let response = self.app.handle_request(r);

        w.headers.content_type = Some(MediaType {
            type_ : String::from_str("text"),
            subtype: String::from_str("html"),
            parameters: vec!((String::from_str("charset"), String::from_str("UTF-8")))
        });
        w.headers.server = Some(String::from_str("Pencil"));
        w.write(response.body.as_bytes()).unwrap();

        self.app.do_teardown_request();
    }
}

/// Run the `Pencil` application.
pub fn run_server(application: Pencil) {
    let pencil_server = PencilServer::new(application);
    pencil_server.serve_forever();
}
<|endoftext|>


<commit_before>use syntax::ast::{self, AstNode};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    if ctx.frange.range.is_empty() {
        return None;
    }

    let node = ctx.find_node_at_range::<ast::Type>()?;
    let insert = ctx.find_node_at_offset::<ast::Item>()?.syntax().text_range().start();
    let target = node.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            builder.edit_file(ctx.frange.file_id);
            builder.replace(target, "Type");
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(cap, insert, format!("type $0Type = {};\n\n", node));
                }
                None => {
                    builder.insert(insert, format!("type Type = {};\n\n", node));
                }
            }
        },
    )
}

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

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }
}
<commit_msg>Place extracted type alias outside of impl<commit_after>use syntax::ast::{self, AstNode};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    if ctx.frange.range.is_empty() {
        return None;
    }

    let node = ctx.find_node_at_range::<ast::Type>()?;
    let insert = ctx
        .find_node_at_offset::<ast::Impl>()
        .map(|imp| imp.syntax().clone())
        .or_else(|| ctx.find_node_at_offset::<ast::Item>().map(|item| item.syntax().clone()))?
        .text_range()
        .start();
    let target = node.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            builder.edit_file(ctx.frange.file_id);
            builder.replace(target, "Type");
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(cap, insert, format!("type $0Type = {};\n\n", node));
                }
                None => {
                    builder.insert(insert, format!("type Type = {};\n\n", node));
                }
            }
        },
    )
}

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

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }

    #[test]
    fn extract_from_impl() {
        // When invoked in an impl, extracted type alias should be placed next to the impl, not
        // inside.
        check_assist(
            extract_type_alias,
            r#"
impl S {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

impl S {
    fn f() -> Type {}
}
            "#,
        );
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::sync::mpsc::{self, channel, TryRecvError};
use std::io::{Read, ErrorKind};
use std::collections::HashMap;
use std::time::Duration;

use mio::*;
use mio::tcp::TcpListener;

use ladspa::{PluginDescriptor, Plugin, PortConnection};

use super::BASE_PORT;
use super::packet::{BYTE_BUFFER_SIZE, Packet};

type ClientPacket = (u64, Packet);

const SERVER: Token = Token(0);

pub struct Receiver {
    sample_rate: u64,
    channel: u16,
    packet_rx: Option<mpsc::Receiver<ClientPacket>>,
    active_packets: Vec<ClientPacket>,
    notify_tx: Option<Sender<<PacketReceiver as Handler>::Message>>,
    client_time_map: HashMap<u64, u64>,
}

impl Receiver {
    pub fn new(_: &PluginDescriptor, sample_rate: u64) -> Box<Plugin + Send> {
        println!("receiver::new");
        Box::new(Receiver {
            sample_rate: sample_rate,
            channel: 0,
            packet_rx: None,
            active_packets: Vec::new(),
            notify_tx: None,
            client_time_map: HashMap::new(),
        })
    }

    fn init_server(&mut self) {
        let (data_tx, data_rx) = channel();
        self.packet_rx = Some(data_rx);

        let channel = self.channel;
        let mut event_loop = EventLoop::new().unwrap();
        self.notify_tx = Some(event_loop.channel());
        thread::spawn(move || {
            let addr = format!("127.0.0.1:{}", BASE_PORT + channel).parse().unwrap();
            let server;
            loop {
                match TcpListener::bind(&addr) {
                    Ok(s) => {
                        server = s;
                        break;
                    }
                    Err(_) => {}
                }
            }
            event_loop.register(&server, SERVER).unwrap();
            event_loop.run(&mut PacketReceiver {
                          server: server,
                          data_tx: data_tx,
                          client_id: 0,
                      })
                      .unwrap();
        });
    }

    fn kill_server(&mut self) {
        self.notify_tx.as_ref().unwrap().send(()).unwrap();
    }
}

impl Plugin for Receiver {
    fn run<'a>(&mut self, sample_count: usize, ports: &[&'a PortConnection<'a>]) {
        let inputl = ports[0].unwrap_audio();
        let inputr = ports[1].unwrap_audio();
        let mut outputl = ports[2].unwrap_audio_mut();
        let mut outputr = ports[3].unwrap_audio_mut();

        let channel = *ports[4].unwrap_control() as u16;

        if channel != self.channel {
            self.channel = channel;
            println!("set channel {}", self.channel);
            self.kill_server();
            self.init_server();
            return;
        }

        loop {
            let packet = match self.packet_rx.as_ref().unwrap().try_recv() {
                Ok(packet) => packet,
                Err(TryRecvError::Disconnected) => {
                    println!("ladspa packet receive failed, dead channel!");
                    self.kill_server();
                    self.init_server();
                    break;
                }
                Err(TryRecvError::Empty) => {
                    break;
                }
            };
            self.active_packets.push(packet);
        }

        for i in 0..sample_count {
            outputl[i] = inputl[i];
            outputr[i] = inputr[i];
        }

        let mut read_clients = Vec::new();
        for &mut (ref client_id, ref mut packet) in &mut self.active_packets {
            let client_time = self.client_time_map.get(client_id).map(|x| *x).unwrap_or(0);
            for i in 0..sample_count {
                let (l, r) = packet.read(client_time + i as u64);
                outputl[i] += l;
                outputr[i] += r;
            }
            read_clients.push(*client_id);
        }
        for client_id in read_clients {
            let client_time = self.client_time_map.get(&client_id).map(|x| *x).unwrap_or(0);
            self.client_time_map.insert(client_id, client_time + sample_count as u64);
        }

        let mut i = 0;
        while i < self.active_packets.len() {
            let complete = {
                let (ref client_id, ref packet) = self.active_packets[i];
                let client_time = self.client_time_map[client_id];
                packet.complete(client_time)
            };
            if complete {
                self.active_packets.remove(i);
            } else {
                i += 1;
            }
        }
    }

    fn activate(&mut self) {
        println!("activate {}", self.channel);
        self.client_time_map.clear();
        self.init_server();
    }

    fn deactivate(&mut self) {
        println!("deactivate {}", self.channel);
        self.kill_server();
    }
}

struct PacketReceiver {
    server: TcpListener,
    data_tx: mpsc::Sender<ClientPacket>,
    client_id: u64,
}

impl Handler for PacketReceiver {
    type Timeout = ();
        type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        match token {
            SERVER => {
                // Only receive readable events
                assert!(events.is_readable());

                println!("server wait");
                match self.server.accept() {
                    Ok(Some(mut socket)) => {
                        socket.set_nodelay(true).unwrap();
                        let tx = self.data_tx.clone();
                        let mut buf = [0; BYTE_BUFFER_SIZE];
                        let mut buf_pos = 0;
                        let client_id = self.client_id;
                        println!("server accept client {}", client_id);
                        self.client_id += 1;
                        loop {
                            let res = socket.read(&mut buf[buf_pos..]);
                            match res {
                                Ok(num_read) => {
                                    //println!("server read {}", num_read);
                                    // if we got a length zero read, the connection is done.
                                    if num_read == 0 {
                                        println!("read zero bytes");
                                        return;
                                    }

                                    // check if we've filled the buffer
                                    buf_pos += num_read;
                                    if buf_pos != BYTE_BUFFER_SIZE {
                                        continue;
                                    }
                                    buf_pos = 0;
                                }
                                Err(e) => {
                                    if e.kind() == ErrorKind::WouldBlock {
                                        //TODO this is a quick hack to reduce CPU usage
                                        thread::sleep(Duration::from_millis(10));
                                        continue;
                                    }
                                    panic!(e);
                                }
                            }
                            let packet = Packet::parse(&buf[..]);
                            if let Err(_) = tx.send((client_id, packet)) {
                                println!("send packet to ladspa error! channel is dead.");
                                return;
                            }
                        }
                    }
                    Ok(None) => {
                        println!("the server socket wasn't actually ready");
                    }
                    Err(e) => {
                        println!("listener.accept() errored: {}", e);
                        // event_loop.shutdown();
                    }
                }
            }
            _ => panic!("Received unknown token"),
        }
    }
    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Self::Message) {
        event_loop.shutdown();
    }
}
<commit_msg>probably performance and latency improvements<commit_after>use std::thread;
use std::sync::mpsc::{self, channel, TryRecvError};
use std::io::{Read, ErrorKind};
use std::collections::HashMap;
use std::time::Duration;

use mio::*;
use mio::tcp::TcpListener;

use ladspa::{PluginDescriptor, Plugin, PortConnection};

use super::BASE_PORT;
use super::packet::{BYTE_BUFFER_SIZE, Packet};

type ClientPacket = (u64, Packet);

const SERVER: Token = Token(0);

pub struct Receiver {
    sample_rate: u64,
    channel: u16,
    packet_rx: Option<mpsc::Receiver<ClientPacket>>,
    active_packets: Vec<ClientPacket>,
    notify_tx: Option<Sender<<PacketReceiver as Handler>::Message>>,
    client_time_map: HashMap<u64, u64>,
}

impl Receiver {
    pub fn new(_: &PluginDescriptor, sample_rate: u64) -> Box<Plugin + Send> {
        println!("receiver::new");
        Box::new(Receiver {
            sample_rate: sample_rate,
            channel: 0,
            packet_rx: None,
            active_packets: Vec::new(),
            notify_tx: None,
            client_time_map: HashMap::new(),
        })
    }

    fn init_server(&mut self) {
        let (data_tx, data_rx) = channel();
        self.packet_rx = Some(data_rx);

        let channel = self.channel;
        let mut event_loop = EventLoop::new().unwrap();
        self.notify_tx = Some(event_loop.channel());
        thread::spawn(move || {
            let addr = format!("127.0.0.1:{}", BASE_PORT + channel).parse().unwrap();
            let server;
            loop {
                match TcpListener::bind(&addr) {
                    Ok(s) => {
                        server = s;
                        break;
                    }
                    Err(_) => {}
                }
            }
            event_loop.register(&server, SERVER).unwrap();
            event_loop.run(&mut PacketReceiver {
                          server: server,
                          data_tx: data_tx,
                          client_id: 0,
                      })
                      .unwrap();
        });
    }

    fn kill_server(&mut self) {
        self.notify_tx.as_ref().unwrap().send(()).unwrap();
    }
}

impl Plugin for Receiver {
    fn run<'a>(&mut self, sample_count: usize, ports: &[&'a PortConnection<'a>]) {
        let inputl = ports[0].unwrap_audio();
        let inputr = ports[1].unwrap_audio();
        let mut outputl = ports[2].unwrap_audio_mut();
        let mut outputr = ports[3].unwrap_audio_mut();

        let channel = *ports[4].unwrap_control() as u16;

        if channel != self.channel {
            self.channel = channel;
            println!("set channel {}", self.channel);
            self.kill_server();
            self.init_server();
            return;
        }

        loop {
            let packet = match self.packet_rx.as_ref().unwrap().try_recv() {
                Ok(packet) => packet,
                Err(TryRecvError::Disconnected) => {
                    println!("ladspa packet receive failed, dead channel!");
                    self.kill_server();
                    self.init_server();
                    break;
                }
                Err(TryRecvError::Empty) => {
                    break;
                }
            };
            self.active_packets.push(packet);
        }

        for i in 0..sample_count {
            outputl[i] = inputl[i];
            outputr[i] = inputr[i];
        }

        let mut read_clients = Vec::new();
        for &(ref client_id, ref packet) in &self.active_packets {
            let client_time = self.client_time_map.get(client_id).map(|x| *x).unwrap_or(0);
            for i in 0..sample_count {
                let (l, r) = packet.read(client_time + i as u64);
                outputl[i] += l;
                outputr[i] += r;
            }
            read_clients.push(*client_id);
        }
        for client_id in read_clients {
            let client_time = self.client_time_map.get(&client_id).map(|x| *x).unwrap_or(0);
            self.client_time_map.insert(client_id, client_time + sample_count as u64);
        }

        let mut i = self.active_packets.len();
        while i > 0 {
            i -= 1;
            let complete = {
                let (ref client_id, ref packet) = self.active_packets[i];
                let client_time = self.client_time_map[client_id];
                packet.complete(client_time)
            };
            if complete {
                self.active_packets.remove(i);
            }
        }
    }

    fn activate(&mut self) {
        println!("activate {}", self.channel);
        self.client_time_map.clear();
        self.init_server();
    }

    fn deactivate(&mut self) {
        println!("deactivate {}", self.channel);
        self.kill_server();
    }
}

struct PacketReceiver {
    server: TcpListener,
    data_tx: mpsc::Sender<ClientPacket>,
    client_id: u64,
}

impl Handler for PacketReceiver {
    type Timeout = ();
        type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        match token {
            SERVER => {
                // Only receive readable events
                assert!(events.is_readable());

                println!("server wait");
                match self.server.accept() {
                    Ok(Some(mut socket)) => {
                        socket.set_nodelay(true).unwrap();
                        let tx = self.data_tx.clone();
                        let mut buf = [0; BYTE_BUFFER_SIZE];
                        let mut buf_pos = 0;
                        let client_id = self.client_id;
                        println!("server accept client {}", client_id);
                        self.client_id += 1;
                        loop {
                            let res = socket.read(&mut buf[buf_pos..]);
                            match res {
                                Ok(num_read) => {
                                    //println!("server read {}", num_read);
                                    // if we got a length zero read, the connection is done.
                                    if num_read == 0 {
                                        println!("read zero bytes");
                                        return;
                                    }

                                    // check if we've filled the buffer
                                    buf_pos += num_read;
                                    if buf_pos != BYTE_BUFFER_SIZE {
                                        continue;
                                    }
                                    buf_pos = 0;
                                }
                                Err(e) => {
                                    if e.kind() == ErrorKind::WouldBlock {
                                        //TODO this is a quick hack to reduce CPU usage
                                        thread::sleep(Duration::from_millis(1));
                                        continue;
                                    }
                                    panic!(e);
                                }
                            }
                            let packet = Packet::parse(&buf[..]);
                            if let Err(_) = tx.send((client_id, packet)) {
                                println!("send packet to ladspa error! channel is dead.");
                                return;
                            }
                        }
                    }
                    Ok(None) => {
                        println!("the server socket wasn't actually ready");
                    }
                    Err(e) => {
                        println!("listener.accept() errored: {}", e);
                        // event_loop.shutdown();
                    }
                }
            }
            _ => panic!("Received unknown token"),
        }
    }
    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Self::Message) {
        event_loop.shutdown();
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{self, Read};
use std::thread::sleep;
use std::time::{Duration, Instant};

use rand::{thread_rng, Rng};
use redis::Value::Okay;
use redis::{Client, IntoConnectionInfo, RedisResult, Value};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
#[derive(Debug, Clone)]
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

pub struct RedLockGuard<'a> {
    pub lock: Lock<'a>,
}

impl Drop for RedLockGuard<'_> {
    fn drop(&mut self) {
        self.lock.lock_manager.unlock(&self.lock);
    }
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new<T: AsRef<str> + IntoConnectionInfo>(uris: Vec<T>) -> RedLock {
        let quorum = (uris.len() as u32) / 2 + 1;

        let servers: Vec<Client> = uris
            .into_iter()
            .map(|uri| Client::open(uri).unwrap())
            .collect();

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open("/dev/urandom")?;
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&mut con);
        match result {
            Ok(Okay) => true,
            Ok(_) | Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let mut rng = thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = rng.gen_range(0..self.retry_delay);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    /// Acquire the lock for the given resource and the requested TTL. \
    /// Will wait and yield current task (tokio runtime) until the lock \
    /// is acquired
    ///
    /// Returns a `RedLockGuard` instance which is a RAII wrapper for \
    /// the old `Lock` object
    #[cfg(feature = "async")]
    pub async fn acquire_async(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => {
                    lock = l;
                    break;
                }
                None => tokio::task::yield_now().await,
            }
        }
        RedLockGuard { lock }
    }

    pub fn acquire(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            if let Some(l) = self.lock(resource, ttl) {
                lock = l;
                break;
            }
        }
        RedLockGuard { lock }
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&mut con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use once_cell::sync::Lazy;
    use testcontainers::clients::Cli;
    use testcontainers::images::redis::Redis;
    use testcontainers::{Container, Docker};

    use super::*;

    static DOCKER: Lazy<Cli> = Lazy::new(Cli::default);
    static CONTAINERS: Lazy<Vec<Container<Cli, Redis>>> = Lazy::new(|| {
        (0..3)
            .map(|_| DOCKER.run(Redis::default().with_tag("6-alpine")))
            .collect()
    });
    static ADDRESSES: Lazy<Vec<String>> = Lazy::new(|| match std::env::var("ADDRESSES") {
        Ok(addresses) => addresses.split(',').map(String::from).collect(),
        Err(_) => CONTAINERS
            .iter()
            .map(|c| format!("redis://localhost:{}", c.get_host_port(6379).unwrap()))
            .collect(),
    });

    #[test]
    fn test_redlock_get_unique_id() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());
        assert_eq!(rl.get_unique_lock_id()?.len(), 20);
        Ok(())
    }

    #[test]
    fn test_redlock_get_unique_id_uniqueness() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());

        let id1 = rl.get_unique_lock_id()?;
        let id2 = rl.get_unique_lock_id()?;

        assert_eq!(20, id1.len());
        assert_eq!(20, id2.len());
        assert_ne!(id1, id2);
        Ok(())
    }

    #[test]
    fn test_redlock_valid_instance() {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        assert_eq!(3, rl.servers.len());
        assert_eq!(2, rl.quorum);
    }

    #[test]
    fn test_redlock_direct_unlock_fails() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        assert!(!rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_unlock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        redis::cmd("SET").arg(&*key).arg(&*val).execute(&mut con);

        assert!(rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_lock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;

        redis::cmd("DEL").arg(&*key).execute(&mut con);
        assert!(rl.lock_instance(&rl.servers[0], &*key, &*val, 1000));
        Ok(())
    }

    #[test]
    fn test_redlock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        let _: () = redis::cmd("SET")
            .arg(&*key)
            .arg(&*val)
            .query(&mut con)
            .unwrap();

        let lock = Lock {
            lock_manager: &rl,
            resource: key,
            val,
            validity_time: 0,
        };
        rl.unlock(&lock);
        Ok(())
    }

    #[test]
    fn test_redlock_lock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        match rl.lock(&key, 1000) {
            Some(lock) => {
                assert_eq!(key, lock.resource);
                assert_eq!(20, lock.val.len());
                assert!(lock.validity_time > 900);
                assert!(
                    lock.validity_time > 900,
                    "validity time: {}",
                    lock.validity_time
                );
            }
            None => panic!("Lock failed"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;

        let lock = rl.lock(&key, 1000).unwrap();
        assert!(
            lock.validity_time > 900,
            "validity time: {}",
            lock.validity_time
        );

        if let Some(_l) = rl2.lock(&key, 1000) {
            panic!("Lock acquired, even though it should be locked")
        }

        rl.unlock(&lock);

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock_raii() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        {
            let lock_guard = rl.acquire(&key, 1000);
            let lock = &lock_guard.lock;
            assert!(
                lock.validity_time > 900,
                "validity time: {}",
                lock.validity_time
            );

            if let Some(_l) = rl2.lock(&key, 1000) {
                panic!("Lock acquired, even though it should be locked")
            }
        }

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }
}
<commit_msg>Reuse existing redis connections (#19)<commit_after>use std::fs::File;
use std::io::{self, Read};
use std::thread::sleep;
use std::time::{Duration, Instant};

use rand::{thread_rng, Rng};
use redis::Value::Okay;
use redis::{Client, IntoConnectionInfo, RedisResult, Value};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
#[derive(Debug, Clone)]
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

pub struct RedLockGuard<'a> {
    pub lock: Lock<'a>,
}

impl Drop for RedLockGuard<'_> {
    fn drop(&mut self) {
        self.lock.lock_manager.unlock(&self.lock);
    }
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new<T: AsRef<str> + IntoConnectionInfo>(uris: Vec<T>) -> RedLock {
        let servers: Vec<Client> = uris
            .into_iter()
            .map(|uri| Client::open(uri).unwrap())
            .collect();

        Self::with_clients(servers)
    }

    pub fn with_client(server: Client) -> RedLock {
        Self::with_clients(vec![server])
    }

    /// Create a new lock manager instance, defined by the given Redis client instance.
    /// Quorum is defined to be N/2+1, with N being the number of given client instances.
    pub fn with_clients(servers: Vec<Client>) -> RedLock {
        let quorum = (servers.len() as u32) / 2 + 1;

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open("/dev/urandom")?;
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&mut con);
        match result {
            Ok(Okay) => true,
            Ok(_) | Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let mut rng = thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = rng.gen_range(0..self.retry_delay);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    /// Acquire the lock for the given resource and the requested TTL. \
    /// Will wait and yield current task (tokio runtime) until the lock \
    /// is acquired
    ///
    /// Returns a `RedLockGuard` instance which is a RAII wrapper for \
    /// the old `Lock` object
    #[cfg(feature = "async")]
    pub async fn acquire_async(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => {
                    lock = l;
                    break;
                }
                None => tokio::task::yield_now().await,
            }
        }
        RedLockGuard { lock }
    }

    pub fn acquire(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            if let Some(l) = self.lock(resource, ttl) {
                lock = l;
                break;
            }
        }
        RedLockGuard { lock }
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&mut con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use once_cell::sync::Lazy;
    use testcontainers::clients::Cli;
    use testcontainers::images::redis::Redis;
    use testcontainers::{Container, Docker};

    use super::*;

    static DOCKER: Lazy<Cli> = Lazy::new(Cli::default);
    static CONTAINERS: Lazy<Vec<Container<Cli, Redis>>> = Lazy::new(|| {
        (0..3)
            .map(|_| DOCKER.run(Redis::default().with_tag("6-alpine")))
            .collect()
    });
    static ADDRESSES: Lazy<Vec<String>> = Lazy::new(|| match std::env::var("ADDRESSES") {
        Ok(addresses) => addresses.split(',').map(String::from).collect(),
        Err(_) => CONTAINERS
            .iter()
            .map(|c| format!("redis://localhost:{}", c.get_host_port(6379).unwrap()))
            .collect(),
    });

    #[test]
    fn test_redlock_get_unique_id() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());
        assert_eq!(rl.get_unique_lock_id()?.len(), 20);
        Ok(())
    }

    #[test]
    fn test_redlock_get_unique_id_uniqueness() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());

        let id1 = rl.get_unique_lock_id()?;
        let id2 = rl.get_unique_lock_id()?;

        assert_eq!(20, id1.len());
        assert_eq!(20, id2.len());
        assert_ne!(id1, id2);
        Ok(())
    }

    #[test]
    fn test_redlock_valid_instance() {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        assert_eq!(3, rl.servers.len());
        assert_eq!(2, rl.quorum);
    }

    #[test]
    fn test_redlock_direct_unlock_fails() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        assert!(!rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_unlock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        redis::cmd("SET").arg(&*key).arg(&*val).execute(&mut con);

        assert!(rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_lock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;

        redis::cmd("DEL").arg(&*key).execute(&mut con);
        assert!(rl.lock_instance(&rl.servers[0], &*key, &*val, 1000));
        Ok(())
    }

    #[test]
    fn test_redlock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        let _: () = redis::cmd("SET")
            .arg(&*key)
            .arg(&*val)
            .query(&mut con)
            .unwrap();

        let lock = Lock {
            lock_manager: &rl,
            resource: key,
            val,
            validity_time: 0,
        };
        rl.unlock(&lock);
        Ok(())
    }

    #[test]
    fn test_redlock_lock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        match rl.lock(&key, 1000) {
            Some(lock) => {
                assert_eq!(key, lock.resource);
                assert_eq!(20, lock.val.len());
                assert!(lock.validity_time > 900);
                assert!(
                    lock.validity_time > 900,
                    "validity time: {}",
                    lock.validity_time
                );
            }
            None => panic!("Lock failed"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;

        let lock = rl.lock(&key, 1000).unwrap();
        assert!(
            lock.validity_time > 900,
            "validity time: {}",
            lock.validity_time
        );

        if let Some(_l) = rl2.lock(&key, 1000) {
            panic!("Lock acquired, even though it should be locked")
        }

        rl.unlock(&lock);

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock_raii() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        {
            let lock_guard = rl.acquire(&key, 1000);
            let lock = &lock_guard.lock;
            assert!(
                lock.validity_time > 900,
                "validity time: {}",
                lock.validity_time
            );

            if let Some(_l) = rl2.lock(&key, 1000) {
                panic!("Lock acquired, even though it should be locked")
            }
        }

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Reorder items.
//!
//! `mod`, `extern crate` and `use` declarations are reorderd in alphabetical
//! order. Trait items are reordered in pre-determined order (associated types
//! and constants comes before methods).

// TODO(#2455): Reorder trait items.

use config::{lists::*, Config};
use syntax::{ast, attr, codemap::Span};

use attr::filter_inline_attrs;
use codemap::LineRangeUtils;
use comment::combine_strs_with_missing_comments;
use imports::{merge_use_trees, UseTree};
use items::{is_mod_decl, rewrite_extern_crate, rewrite_mod};
use lists::{itemize_list, write_list, ListFormatting, ListItem};
use rewrite::{Rewrite, RewriteContext};
use shape::Shape;
use spanned::Spanned;
use utils::mk_sp;
use visitor::FmtVisitor;

use std::cmp::{Ord, Ordering};

/// Choose the ordering between the given two items.
fn compare_items(a: &ast::Item, b: &ast::Item) -> Ordering {
    match (&a.node, &b.node) {
        (&ast::ItemKind::Mod(..), &ast::ItemKind::Mod(..)) => {
            a.ident.name.as_str().cmp(&b.ident.name.as_str())
        }
        (&ast::ItemKind::ExternCrate(ref a_name), &ast::ItemKind::ExternCrate(ref b_name)) => {
            // `extern crate foo as bar;`
            //               ^^^ Comparing this.
            let a_orig_name =
                a_name.map_or_else(|| a.ident.name.as_str(), |symbol| symbol.as_str());
            let b_orig_name =
                b_name.map_or_else(|| b.ident.name.as_str(), |symbol| symbol.as_str());
            let result = a_orig_name.cmp(&b_orig_name);
            if result != Ordering::Equal {
                return result;
            }

            // `extern crate foo as bar;`
            //                      ^^^ Comparing this.
            match (a_name, b_name) {
                (Some(..), None) => Ordering::Greater,
                (None, Some(..)) => Ordering::Less,
                (None, None) => Ordering::Equal,
                (Some(..), Some(..)) => a.ident.name.as_str().cmp(&b.ident.name.as_str()),
            }
        }
        _ => unreachable!(),
    }
}

fn wrap_reorderable_items(
    context: &RewriteContext,
    list_items: &[ListItem],
    shape: Shape,
) -> Option<String> {
    let fmt = ListFormatting {
        tactic: DefinitiveListTactic::Vertical,
        separator: "",
        trailing_separator: SeparatorTactic::Never,
        separator_place: SeparatorPlace::Back,
        shape,
        ends_with_newline: true,
        preserve_newline: false,
        config: context.config,
    };

    write_list(list_items, &fmt)
}

fn rewrite_reorderable_item(
    context: &RewriteContext,
    item: &ast::Item,
    shape: Shape,
) -> Option<String> {
    let attrs = filter_inline_attrs(&item.attrs, item.span());
    let attrs_str = attrs.rewrite(context, shape)?;

    let missed_span = if attrs.is_empty() {
        mk_sp(item.span.lo(), item.span.lo())
    } else {
        mk_sp(attrs.last().unwrap().span.hi(), item.span.lo())
    };

    let item_str = match item.node {
        ast::ItemKind::ExternCrate(..) => rewrite_extern_crate(context, item)?,
        ast::ItemKind::Mod(..) => rewrite_mod(item),
        _ => return None,
    };

    combine_strs_with_missing_comments(context, &attrs_str, &item_str, missed_span, shape, false)
}

/// Rewrite a list of items with reordering. Every item in `items` must have
/// the same `ast::ItemKind`.
fn rewrite_reorderable_items(
    context: &RewriteContext,
    reorderable_items: &[&ast::Item],
    shape: Shape,
    span: Span,
) -> Option<String> {
    match reorderable_items[0].node {
        // FIXME: Remove duplicated code.
        ast::ItemKind::Use(..) => {
            let mut normalized_items: Vec<_> = reorderable_items
                .iter()
                .filter_map(|item| UseTree::from_ast_with_normalization(context, item))
                .collect();
            let cloned = normalized_items.clone();
            // Add comments before merging.
            let list_items = itemize_list(
                context.snippet_provider,
                cloned.iter(),
                "",
                ";",
                |item| item.span().lo(),
                |item| item.span().hi(),
                |_item| Some("".to_owned()),
                span.lo(),
                span.hi(),
                false,
            );
            for (item, list_item) in normalized_items.iter_mut().zip(list_items) {
                item.list_item = Some(list_item.clone());
            }
            if context.config.merge_imports() {
                normalized_items = merge_use_trees(normalized_items);
            }
            normalized_items.sort();

            // 4 = "use ", 1 = ";"
            let nested_shape = shape.offset_left(4)?.sub_width(1)?;
            let item_vec: Vec<_> = normalized_items
                .into_iter()
                .map(|use_tree| ListItem {
                    item: use_tree.rewrite_top_level(context, nested_shape),
                    ..use_tree.list_item.unwrap_or_else(ListItem::empty)
                })
                .collect();

            wrap_reorderable_items(context, &item_vec, nested_shape)
        }
        _ => {
            let list_items = itemize_list(
                context.snippet_provider,
                reorderable_items.iter(),
                "",
                ";",
                |item| item.span().lo(),
                |item| item.span().hi(),
                |item| rewrite_reorderable_item(context, item, shape),
                span.lo(),
                span.hi(),
                false,
            );

            let mut item_pair_vec: Vec<_> = list_items.zip(reorderable_items.iter()).collect();
            item_pair_vec.sort_by(|a, b| compare_items(a.1, b.1));
            let item_vec: Vec<_> = item_pair_vec.into_iter().map(|pair| pair.0).collect();

            wrap_reorderable_items(context, &item_vec, shape)
        }
    }
}

fn contains_macro_use_attr(item: &ast::Item) -> bool {
    attr::contains_name(&filter_inline_attrs(&item.attrs, item.span()), "macro_use")
}

/// A simplified version of `ast::ItemKind`.
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
enum ReorderableItemKind {
    ExternCrate,
    Mod,
    Use,
    /// An item that cannot be reordered. Either has an unreorderable item kind
    /// or an `macro_use` attribute.
    Other,
}

impl ReorderableItemKind {
    pub fn from(item: &ast::Item) -> Self {
        match item.node {
            _ if contains_macro_use_attr(item) => ReorderableItemKind::Other,
            ast::ItemKind::ExternCrate(..) => ReorderableItemKind::ExternCrate,
            ast::ItemKind::Mod(..) if is_mod_decl(item) => ReorderableItemKind::Mod,
            ast::ItemKind::Use(..) => ReorderableItemKind::Use,
            _ => ReorderableItemKind::Other,
        }
    }

    pub fn is_same_item_kind(&self, item: &ast::Item) -> bool {
        ReorderableItemKind::from(item) == *self
    }

    pub fn is_reorderable(&self, config: &Config) -> bool {
        match *self {
            ReorderableItemKind::ExternCrate => config.reorder_imports(),
            ReorderableItemKind::Mod => config.reorder_modules(),
            ReorderableItemKind::Use => config.reorder_imports(),
            ReorderableItemKind::Other => false,
        }
    }

    pub fn in_group(&self) -> bool {
        match *self {
            ReorderableItemKind::ExternCrate
            | ReorderableItemKind::Mod
            | ReorderableItemKind::Use => true,
            ReorderableItemKind::Other => false,
        }
    }
}

impl<'b, 'a: 'b> FmtVisitor<'a> {
    /// Format items with the same item kind and reorder them. If `in_group` is
    /// `true`, then the items separated by an empty line will not be reordered
    /// together.
    fn walk_reorderable_items(
        &mut self,
        items: &[&ast::Item],
        item_kind: ReorderableItemKind,
        in_group: bool,
    ) -> usize {
        let mut last = self.codemap.lookup_line_range(items[0].span());
        let item_length = items
            .iter()
            .take_while(|ppi| {
                item_kind.is_same_item_kind(&***ppi) && (!in_group || {
                    let current = self.codemap.lookup_line_range(ppi.span());
                    let in_same_group = current.lo < last.hi + 2;
                    last = current;
                    in_same_group
                })
            })
            .count();
        let items = &items[..item_length];

        let at_least_one_in_file_lines = items
            .iter()
            .any(|item| !out_of_file_lines_range!(self, item.span));

        if at_least_one_in_file_lines && !items.is_empty() {
            let lo = items.first().unwrap().span().lo();
            let hi = items.last().unwrap().span().hi();
            let span = mk_sp(lo, hi);
            let rw = rewrite_reorderable_items(&self.get_context(), items, self.shape(), span);
            self.push_rewrite(span, rw);
        } else {
            for item in items {
                self.push_rewrite(item.span, None);
            }
        }

        item_length
    }

    /// Visit and format the given items. Items are reordered If they are
    /// consecutive and reorderable.
    pub fn visit_items_with_reordering(&mut self, mut items: &[&ast::Item]) {
        while !items.is_empty() {
            // If the next item is a `use`, `extern crate` or `mod`, then extract it and any
            // subsequent items that have the same item kind to be reordered within
            // `walk_reorderable_items`. Otherwise, just format the next item for output.
            let item_kind = ReorderableItemKind::from(items[0]);
            if item_kind.is_reorderable(self.config) {
                let visited_items_num =
                    self.walk_reorderable_items(items, item_kind, item_kind.in_group());
                let (_, rest) = items.split_at(visited_items_num);
                items = rest;
            } else {
                // Reaching here means items were not reordered. There must be at least
                // one item left in `items`, so calling `unwrap()` here is safe.
                let (item, rest) = items.split_first().unwrap();
                self.visit_item(item);
                items = rest;
            }
        }
    }
}
<commit_msg>Remove some unnecessary `pub`s<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Reorder items.
//!
//! `mod`, `extern crate` and `use` declarations are reorderd in alphabetical
//! order. Trait items are reordered in pre-determined order (associated types
//! and constants comes before methods).

// TODO(#2455): Reorder trait items.

use config::{lists::*, Config};
use syntax::{ast, attr, codemap::Span};

use attr::filter_inline_attrs;
use codemap::LineRangeUtils;
use comment::combine_strs_with_missing_comments;
use imports::{merge_use_trees, UseTree};
use items::{is_mod_decl, rewrite_extern_crate, rewrite_mod};
use lists::{itemize_list, write_list, ListFormatting, ListItem};
use rewrite::{Rewrite, RewriteContext};
use shape::Shape;
use spanned::Spanned;
use utils::mk_sp;
use visitor::FmtVisitor;

use std::cmp::{Ord, Ordering};

/// Choose the ordering between the given two items.
fn compare_items(a: &ast::Item, b: &ast::Item) -> Ordering {
    match (&a.node, &b.node) {
        (&ast::ItemKind::Mod(..), &ast::ItemKind::Mod(..)) => {
            a.ident.name.as_str().cmp(&b.ident.name.as_str())
        }
        (&ast::ItemKind::ExternCrate(ref a_name), &ast::ItemKind::ExternCrate(ref b_name)) => {
            // `extern crate foo as bar;`
            //               ^^^ Comparing this.
            let a_orig_name =
                a_name.map_or_else(|| a.ident.name.as_str(), |symbol| symbol.as_str());
            let b_orig_name =
                b_name.map_or_else(|| b.ident.name.as_str(), |symbol| symbol.as_str());
            let result = a_orig_name.cmp(&b_orig_name);
            if result != Ordering::Equal {
                return result;
            }

            // `extern crate foo as bar;`
            //                      ^^^ Comparing this.
            match (a_name, b_name) {
                (Some(..), None) => Ordering::Greater,
                (None, Some(..)) => Ordering::Less,
                (None, None) => Ordering::Equal,
                (Some(..), Some(..)) => a.ident.name.as_str().cmp(&b.ident.name.as_str()),
            }
        }
        _ => unreachable!(),
    }
}

fn wrap_reorderable_items(
    context: &RewriteContext,
    list_items: &[ListItem],
    shape: Shape,
) -> Option<String> {
    let fmt = ListFormatting {
        tactic: DefinitiveListTactic::Vertical,
        separator: "",
        trailing_separator: SeparatorTactic::Never,
        separator_place: SeparatorPlace::Back,
        shape,
        ends_with_newline: true,
        preserve_newline: false,
        config: context.config,
    };

    write_list(list_items, &fmt)
}

fn rewrite_reorderable_item(
    context: &RewriteContext,
    item: &ast::Item,
    shape: Shape,
) -> Option<String> {
    let attrs = filter_inline_attrs(&item.attrs, item.span());
    let attrs_str = attrs.rewrite(context, shape)?;

    let missed_span = if attrs.is_empty() {
        mk_sp(item.span.lo(), item.span.lo())
    } else {
        mk_sp(attrs.last().unwrap().span.hi(), item.span.lo())
    };

    let item_str = match item.node {
        ast::ItemKind::ExternCrate(..) => rewrite_extern_crate(context, item)?,
        ast::ItemKind::Mod(..) => rewrite_mod(item),
        _ => return None,
    };

    combine_strs_with_missing_comments(context, &attrs_str, &item_str, missed_span, shape, false)
}

/// Rewrite a list of items with reordering. Every item in `items` must have
/// the same `ast::ItemKind`.
fn rewrite_reorderable_items(
    context: &RewriteContext,
    reorderable_items: &[&ast::Item],
    shape: Shape,
    span: Span,
) -> Option<String> {
    match reorderable_items[0].node {
        // FIXME: Remove duplicated code.
        ast::ItemKind::Use(..) => {
            let mut normalized_items: Vec<_> = reorderable_items
                .iter()
                .filter_map(|item| UseTree::from_ast_with_normalization(context, item))
                .collect();
            let cloned = normalized_items.clone();
            // Add comments before merging.
            let list_items = itemize_list(
                context.snippet_provider,
                cloned.iter(),
                "",
                ";",
                |item| item.span().lo(),
                |item| item.span().hi(),
                |_item| Some("".to_owned()),
                span.lo(),
                span.hi(),
                false,
            );
            for (item, list_item) in normalized_items.iter_mut().zip(list_items) {
                item.list_item = Some(list_item.clone());
            }
            if context.config.merge_imports() {
                normalized_items = merge_use_trees(normalized_items);
            }
            normalized_items.sort();

            // 4 = "use ", 1 = ";"
            let nested_shape = shape.offset_left(4)?.sub_width(1)?;
            let item_vec: Vec<_> = normalized_items
                .into_iter()
                .map(|use_tree| ListItem {
                    item: use_tree.rewrite_top_level(context, nested_shape),
                    ..use_tree.list_item.unwrap_or_else(ListItem::empty)
                })
                .collect();

            wrap_reorderable_items(context, &item_vec, nested_shape)
        }
        _ => {
            let list_items = itemize_list(
                context.snippet_provider,
                reorderable_items.iter(),
                "",
                ";",
                |item| item.span().lo(),
                |item| item.span().hi(),
                |item| rewrite_reorderable_item(context, item, shape),
                span.lo(),
                span.hi(),
                false,
            );

            let mut item_pair_vec: Vec<_> = list_items.zip(reorderable_items.iter()).collect();
            item_pair_vec.sort_by(|a, b| compare_items(a.1, b.1));
            let item_vec: Vec<_> = item_pair_vec.into_iter().map(|pair| pair.0).collect();

            wrap_reorderable_items(context, &item_vec, shape)
        }
    }
}

fn contains_macro_use_attr(item: &ast::Item) -> bool {
    attr::contains_name(&filter_inline_attrs(&item.attrs, item.span()), "macro_use")
}

/// A simplified version of `ast::ItemKind`.
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
enum ReorderableItemKind {
    ExternCrate,
    Mod,
    Use,
    /// An item that cannot be reordered. Either has an unreorderable item kind
    /// or an `macro_use` attribute.
    Other,
}

impl ReorderableItemKind {
    fn from(item: &ast::Item) -> Self {
        match item.node {
            _ if contains_macro_use_attr(item) => ReorderableItemKind::Other,
            ast::ItemKind::ExternCrate(..) => ReorderableItemKind::ExternCrate,
            ast::ItemKind::Mod(..) if is_mod_decl(item) => ReorderableItemKind::Mod,
            ast::ItemKind::Use(..) => ReorderableItemKind::Use,
            _ => ReorderableItemKind::Other,
        }
    }

    fn is_same_item_kind(&self, item: &ast::Item) -> bool {
        ReorderableItemKind::from(item) == *self
    }

    fn is_reorderable(&self, config: &Config) -> bool {
        match *self {
            ReorderableItemKind::ExternCrate => config.reorder_imports(),
            ReorderableItemKind::Mod => config.reorder_modules(),
            ReorderableItemKind::Use => config.reorder_imports(),
            ReorderableItemKind::Other => false,
        }
    }

    fn in_group(&self) -> bool {
        match *self {
            ReorderableItemKind::ExternCrate
            | ReorderableItemKind::Mod
            | ReorderableItemKind::Use => true,
            ReorderableItemKind::Other => false,
        }
    }
}

impl<'b, 'a: 'b> FmtVisitor<'a> {
    /// Format items with the same item kind and reorder them. If `in_group` is
    /// `true`, then the items separated by an empty line will not be reordered
    /// together.
    fn walk_reorderable_items(
        &mut self,
        items: &[&ast::Item],
        item_kind: ReorderableItemKind,
        in_group: bool,
    ) -> usize {
        let mut last = self.codemap.lookup_line_range(items[0].span());
        let item_length = items
            .iter()
            .take_while(|ppi| {
                item_kind.is_same_item_kind(&***ppi) && (!in_group || {
                    let current = self.codemap.lookup_line_range(ppi.span());
                    let in_same_group = current.lo < last.hi + 2;
                    last = current;
                    in_same_group
                })
            })
            .count();
        let items = &items[..item_length];

        let at_least_one_in_file_lines = items
            .iter()
            .any(|item| !out_of_file_lines_range!(self, item.span));

        if at_least_one_in_file_lines && !items.is_empty() {
            let lo = items.first().unwrap().span().lo();
            let hi = items.last().unwrap().span().hi();
            let span = mk_sp(lo, hi);
            let rw = rewrite_reorderable_items(&self.get_context(), items, self.shape(), span);
            self.push_rewrite(span, rw);
        } else {
            for item in items {
                self.push_rewrite(item.span, None);
            }
        }

        item_length
    }

    /// Visit and format the given items. Items are reordered If they are
    /// consecutive and reorderable.
    pub fn visit_items_with_reordering(&mut self, mut items: &[&ast::Item]) {
        while !items.is_empty() {
            // If the next item is a `use`, `extern crate` or `mod`, then extract it and any
            // subsequent items that have the same item kind to be reordered within
            // `walk_reorderable_items`. Otherwise, just format the next item for output.
            let item_kind = ReorderableItemKind::from(items[0]);
            if item_kind.is_reorderable(self.config) {
                let visited_items_num =
                    self.walk_reorderable_items(items, item_kind, item_kind.in_group());
                let (_, rest) = items.split_at(visited_items_num);
                items = rest;
            } else {
                // Reaching here means items were not reordered. There must be at least
                // one item left in `items`, so calling `unwrap()` here is safe.
                let (item, rest) = items.split_first().unwrap();
                self.visit_item(item);
                items = rest;
            }
        }
    }
}
<|endoftext|>