<commit_before>use crate::common::CodegenCx;
use crate::coverageinfo;
use crate::llvm;

use llvm::coverageinfo::CounterMappingRegion;
use rustc_codegen_ssa::coverageinfo::map::{Counter, CounterExpression};
use rustc_codegen_ssa::traits::ConstMethods;
use rustc_data_structures::fx::FxIndexSet;
use rustc_llvm::RustString;
use rustc_middle::mir::coverage::CodeRegion;

use std::ffi::CString;

use tracing::debug;

/// Generates and exports the Coverage Map.
///
/// This Coverage Map complies with Coverage Mapping Format version 4 (zero-based encoded as 3),
/// as defined at [LLVM Code Coverage Mapping Format](https://github.com/rust-lang/llvm-project/blob/rustc/11.0-2020-10-12/llvm/docs/CoverageMappingFormat.rst#llvm-code-coverage-mapping-format)
/// and published in Rust's current (November 2020) fork of LLVM. This version is supported by the
/// LLVM coverage tools (`llvm-profdata` and `llvm-cov`) bundled with Rust's fork of LLVM.
///
/// Consequently, Rust's bundled version of Clang also generates Coverage Maps compliant with
/// version 3. Clang's implementation of Coverage Map generation was referenced when implementing
/// this Rust version, and though the format documentation is very explicit and detailed, some
/// undocumented details in Clang's implementation (that may or may not be important) were also
/// replicated for Rust's Coverage Map.
pub fn finalize<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
    // Ensure LLVM supports Coverage Map Version 4 (encoded as a zero-based value: 3).
    // If not, the LLVM Version must be less than 11.
    let version = coverageinfo::mapping_version();
    assert_eq!(version, 3, "rustc option `-Z instrument-coverage` requires LLVM 11 or higher.");

    let function_coverage_map = match cx.coverage_context() {
        Some(ctx) => ctx.take_function_coverage_map(),
        None => return,
    };
    if function_coverage_map.is_empty() {
        // This module has no functions with coverage instrumentation
        return;
    }

    let mut mapgen = CoverageMapGenerator::new();

    // Encode coverage mappings and generate function records
    let mut function_data = Vec::new();
    for (instance, function_coverage) in function_coverage_map {
        debug!("Generate coverage map for: {:?}", instance);

        let mangled_function_name = cx.tcx.symbol_name(instance).to_string();
        let function_source_hash = function_coverage.source_hash();
        let (expressions, counter_regions) =
            function_coverage.get_expressions_and_counter_regions();

        let coverage_mapping_buffer = llvm::build_byte_buffer(|coverage_mapping_buffer| {
            mapgen.write_coverage_mapping(expressions, counter_regions, coverage_mapping_buffer);
        });
        debug_assert!(
            coverage_mapping_buffer.len() > 0,
            "Every `FunctionCoverage` should have at least one counter"
        );

        function_data.push((mangled_function_name, function_source_hash, coverage_mapping_buffer));
    }

    // Encode all filenames referenced by counters/expressions in this module
    let filenames_buffer = llvm::build_byte_buffer(|filenames_buffer| {
        coverageinfo::write_filenames_section_to_buffer(&mapgen.filenames, filenames_buffer);
    });

    let filenames_size = filenames_buffer.len();
    let filenames_val = cx.const_bytes(&filenames_buffer[..]);
    let filenames_ref = coverageinfo::hash_bytes(filenames_buffer);

    // Generate the LLVM IR representation of the coverage map and store it in a well-known global
    let cov_data_val = mapgen.generate_coverage_map(cx, version, filenames_size, filenames_val);

    for (mangled_function_name, function_source_hash, coverage_mapping_buffer) in function_data {
        save_function_record(
            cx,
            mangled_function_name,
            function_source_hash,
            filenames_ref,
            coverage_mapping_buffer,
        );
    }

    // Save the coverage data value to LLVM IR
    coverageinfo::save_cov_data_to_mod(cx, cov_data_val);
}

struct CoverageMapGenerator {
    filenames: FxIndexSet<CString>,
}

impl CoverageMapGenerator {
    fn new() -> Self {
        Self { filenames: FxIndexSet::default() }
    }

    /// Using the `expressions` and `counter_regions` collected for the current function, generate
    /// the `mapping_regions` and `virtual_file_mapping`, and capture any new filenames. Then use
    /// LLVM APIs to encode the `virtual_file_mapping`, `expressions`, and `mapping_regions` into
    /// the given `coverage_mapping` byte buffer, compliant with the LLVM Coverage Mapping format.
    fn write_coverage_mapping(
        &mut self,
        expressions: Vec<CounterExpression>,
        counter_regions: impl Iterator<Item = (Counter, &'a CodeRegion)>,
        coverage_mapping_buffer: &RustString,
    ) {
        let mut counter_regions = counter_regions.collect::<Vec<_>>();
        if counter_regions.is_empty() {
            return;
        }

        let mut virtual_file_mapping = Vec::new();
        let mut mapping_regions = Vec::new();
        let mut current_file_name = None;
        let mut current_file_id = 0;

        // Convert the list of (Counter, CodeRegion) pairs to an array of `CounterMappingRegion`, sorted
        // by filename and position. Capture any new files to compute the `CounterMappingRegion`s
        // `file_id` (indexing files referenced by the current function), and construct the
        // function-specific `virtual_file_mapping` from `file_id` to its index in the module's
        // `filenames` array.
        counter_regions.sort_unstable_by_key(|(_counter, region)| *region);
        for (counter, region) in counter_regions {
            let CodeRegion { file_name, start_line, start_col, end_line, end_col } = *region;
            let same_file = current_file_name.as_ref().map_or(false, |p| *p == file_name);
            if !same_file {
                if current_file_name.is_some() {
                    current_file_id += 1;
                }
                current_file_name = Some(file_name);
                let c_filename = CString::new(file_name.to_string())
                    .expect("null error converting filename to C string");
                debug!("  file_id: {} = '{:?}'", current_file_id, c_filename);
                let (filenames_index, _) = self.filenames.insert_full(c_filename);
                virtual_file_mapping.push(filenames_index as u32);
            }
            debug!("Adding counter {:?} to map for {:?}", counter, region);
            mapping_regions.push(CounterMappingRegion::code_region(
                counter,
                current_file_id,
                start_line,
                start_col,
                end_line,
                end_col,
            ));
        }

        // Encode and append the current function's coverage mapping data
        coverageinfo::write_mapping_to_buffer(
            virtual_file_mapping,
            expressions,
            mapping_regions,
            coverage_mapping_buffer,
        );
    }

    /// Construct coverage map header and the array of function records, and combine them into the
    /// coverage map. Save the coverage map data into the LLVM IR as a static global using a
    /// specific, well-known section and name.
    fn generate_coverage_map(
        self,
        cx: &CodegenCx<'ll, 'tcx>,
        version: u32,
        filenames_size: usize,
        filenames_val: &'ll llvm::Value,
    ) -> &'ll llvm::Value {
        debug!("cov map: filenames_size = {}, 0-based version = {}", filenames_size, version);

        // Create the coverage data header (Note, fields 0 and 2 are now always zero,
        // as of `llvm::coverage::CovMapVersion::Version4`.)
        let zero_was_n_records_val = cx.const_u32(0);
        let filenames_size_val = cx.const_u32(filenames_size as u32);
        let zero_was_coverage_size_val = cx.const_u32(0);
        let version_val = cx.const_u32(version);
        let cov_data_header_val = cx.const_struct(
            &[zero_was_n_records_val, filenames_size_val, zero_was_coverage_size_val, version_val],
            /*packed=*/ false,
        );

        // Create the complete LLVM coverage data value to add to the LLVM IR
        cx.const_struct(&[cov_data_header_val, filenames_val], /*packed=*/ false)
    }
}

/// Construct a function record and combine it with the function's coverage mapping data.
/// Save the function record into the LLVM IR as a static global using a
/// specific, well-known section and name.
fn save_function_record(
    cx: &CodegenCx<'ll, 'tcx>,
    mangled_function_name: String,
    function_source_hash: u64,
    filenames_ref: u64,
    coverage_mapping_buffer: Vec<u8>,
) {
    // Concatenate the encoded coverage mappings
    let coverage_mapping_size = coverage_mapping_buffer.len();
    let coverage_mapping_val = cx.const_bytes(&coverage_mapping_buffer[..]);

    let func_name_hash = coverageinfo::hash_str(&mangled_function_name);
    let func_name_hash_val = cx.const_u64(func_name_hash);
    let coverage_mapping_size_val = cx.const_u32(coverage_mapping_size as u32);
    let func_hash_val = cx.const_u64(function_source_hash);
    let filenames_ref_val = cx.const_u64(filenames_ref);
    let func_record_val = cx.const_struct(
        &[
            func_name_hash_val,
            coverage_mapping_size_val,
            func_hash_val,
            filenames_ref_val,
            coverage_mapping_val,
        ],
        /*packed=*/ true,
    );

    // At the present time, the coverage map for Rust assumes every instrumented function `is_used`.
    // Note that Clang marks functions as "unused" in `CodeGenPGO::emitEmptyCounterMapping`. (See:
    // https://github.com/rust-lang/llvm-project/blob/de02a75e398415bad4df27b4547c25b896c8bf3b/clang%2Flib%2FCodeGen%2FCodeGenPGO.cpp#L877-L878
    // for example.)
    //
    // It's not yet clear if or how this may be applied to Rust in the future, but the `is_used`
    // argument is available and handled similarly.
    let is_used = true;
    coverageinfo::save_func_record_to_mod(cx, func_name_hash, func_record_val, is_used);
}
<commit_msg>replace assert with condition and `fatal` error<commit_after>use crate::common::CodegenCx;
use crate::coverageinfo;
use crate::llvm;

use llvm::coverageinfo::CounterMappingRegion;
use rustc_codegen_ssa::coverageinfo::map::{Counter, CounterExpression};
use rustc_codegen_ssa::traits::ConstMethods;
use rustc_data_structures::fx::FxIndexSet;
use rustc_llvm::RustString;
use rustc_middle::mir::coverage::CodeRegion;

use std::ffi::CString;

use tracing::debug;

/// Generates and exports the Coverage Map.
///
/// This Coverage Map complies with Coverage Mapping Format version 4 (zero-based encoded as 3),
/// as defined at [LLVM Code Coverage Mapping Format](https://github.com/rust-lang/llvm-project/blob/rustc/11.0-2020-10-12/llvm/docs/CoverageMappingFormat.rst#llvm-code-coverage-mapping-format)
/// and published in Rust's current (November 2020) fork of LLVM. This version is supported by the
/// LLVM coverage tools (`llvm-profdata` and `llvm-cov`) bundled with Rust's fork of LLVM.
///
/// Consequently, Rust's bundled version of Clang also generates Coverage Maps compliant with
/// version 3. Clang's implementation of Coverage Map generation was referenced when implementing
/// this Rust version, and though the format documentation is very explicit and detailed, some
/// undocumented details in Clang's implementation (that may or may not be important) were also
/// replicated for Rust's Coverage Map.
pub fn finalize<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
    // Ensure LLVM supports Coverage Map Version 4 (encoded as a zero-based value: 3).
    // If not, the LLVM Version must be less than 11.
    let version = coverageinfo::mapping_version();
    if version != 3 {
        cx.tcx.sess.fatal("rustc option `-Z instrument-coverage` requires LLVM 11 or higher.");
    }```

    let function_coverage_map = match cx.coverage_context() {
        Some(ctx) => ctx.take_function_coverage_map(),
        None => return,
    };
    if function_coverage_map.is_empty() {
        // This module has no functions with coverage instrumentation
        return;
    }

    let mut mapgen = CoverageMapGenerator::new();

    // Encode coverage mappings and generate function records
    let mut function_data = Vec::new();
    for (instance, function_coverage) in function_coverage_map {
        debug!("Generate coverage map for: {:?}", instance);

        let mangled_function_name = cx.tcx.symbol_name(instance).to_string();
        let function_source_hash = function_coverage.source_hash();
        let (expressions, counter_regions) =
            function_coverage.get_expressions_and_counter_regions();

        let coverage_mapping_buffer = llvm::build_byte_buffer(|coverage_mapping_buffer| {
            mapgen.write_coverage_mapping(expressions, counter_regions, coverage_mapping_buffer);
        });
        debug_assert!(
            coverage_mapping_buffer.len() > 0,
            "Every `FunctionCoverage` should have at least one counter"
        );

        function_data.push((mangled_function_name, function_source_hash, coverage_mapping_buffer));
    }

    // Encode all filenames referenced by counters/expressions in this module
    let filenames_buffer = llvm::build_byte_buffer(|filenames_buffer| {
        coverageinfo::write_filenames_section_to_buffer(&mapgen.filenames, filenames_buffer);
    });

    let filenames_size = filenames_buffer.len();
    let filenames_val = cx.const_bytes(&filenames_buffer[..]);
    let filenames_ref = coverageinfo::hash_bytes(filenames_buffer);

    // Generate the LLVM IR representation of the coverage map and store it in a well-known global
    let cov_data_val = mapgen.generate_coverage_map(cx, version, filenames_size, filenames_val);

    for (mangled_function_name, function_source_hash, coverage_mapping_buffer) in function_data {
        save_function_record(
            cx,
            mangled_function_name,
            function_source_hash,
            filenames_ref,
            coverage_mapping_buffer,
        );
    }

    // Save the coverage data value to LLVM IR
    coverageinfo::save_cov_data_to_mod(cx, cov_data_val);
}

struct CoverageMapGenerator {
    filenames: FxIndexSet<CString>,
}

impl CoverageMapGenerator {
    fn new() -> Self {
        Self { filenames: FxIndexSet::default() }
    }

    /// Using the `expressions` and `counter_regions` collected for the current function, generate
    /// the `mapping_regions` and `virtual_file_mapping`, and capture any new filenames. Then use
    /// LLVM APIs to encode the `virtual_file_mapping`, `expressions`, and `mapping_regions` into
    /// the given `coverage_mapping` byte buffer, compliant with the LLVM Coverage Mapping format.
    fn write_coverage_mapping(
        &mut self,
        expressions: Vec<CounterExpression>,
        counter_regions: impl Iterator<Item = (Counter, &'a CodeRegion)>,
        coverage_mapping_buffer: &RustString,
    ) {
        let mut counter_regions = counter_regions.collect::<Vec<_>>();
        if counter_regions.is_empty() {
            return;
        }

        let mut virtual_file_mapping = Vec::new();
        let mut mapping_regions = Vec::new();
        let mut current_file_name = None;
        let mut current_file_id = 0;

        // Convert the list of (Counter, CodeRegion) pairs to an array of `CounterMappingRegion`, sorted
        // by filename and position. Capture any new files to compute the `CounterMappingRegion`s
        // `file_id` (indexing files referenced by the current function), and construct the
        // function-specific `virtual_file_mapping` from `file_id` to its index in the module's
        // `filenames` array.
        counter_regions.sort_unstable_by_key(|(_counter, region)| *region);
        for (counter, region) in counter_regions {
            let CodeRegion { file_name, start_line, start_col, end_line, end_col } = *region;
            let same_file = current_file_name.as_ref().map_or(false, |p| *p == file_name);
            if !same_file {
                if current_file_name.is_some() {
                    current_file_id += 1;
                }
                current_file_name = Some(file_name);
                let c_filename = CString::new(file_name.to_string())
                    .expect("null error converting filename to C string");
                debug!("  file_id: {} = '{:?}'", current_file_id, c_filename);
                let (filenames_index, _) = self.filenames.insert_full(c_filename);
                virtual_file_mapping.push(filenames_index as u32);
            }
            debug!("Adding counter {:?} to map for {:?}", counter, region);
            mapping_regions.push(CounterMappingRegion::code_region(
                counter,
                current_file_id,
                start_line,
                start_col,
                end_line,
                end_col,
            ));
        }

        // Encode and append the current function's coverage mapping data
        coverageinfo::write_mapping_to_buffer(
            virtual_file_mapping,
            expressions,
            mapping_regions,
            coverage_mapping_buffer,
        );
    }

    /// Construct coverage map header and the array of function records, and combine them into the
    /// coverage map. Save the coverage map data into the LLVM IR as a static global using a
    /// specific, well-known section and name.
    fn generate_coverage_map(
        self,
        cx: &CodegenCx<'ll, 'tcx>,
        version: u32,
        filenames_size: usize,
        filenames_val: &'ll llvm::Value,
    ) -> &'ll llvm::Value {
        debug!("cov map: filenames_size = {}, 0-based version = {}", filenames_size, version);

        // Create the coverage data header (Note, fields 0 and 2 are now always zero,
        // as of `llvm::coverage::CovMapVersion::Version4`.)
        let zero_was_n_records_val = cx.const_u32(0);
        let filenames_size_val = cx.const_u32(filenames_size as u32);
        let zero_was_coverage_size_val = cx.const_u32(0);
        let version_val = cx.const_u32(version);
        let cov_data_header_val = cx.const_struct(
            &[zero_was_n_records_val, filenames_size_val, zero_was_coverage_size_val, version_val],
            /*packed=*/ false,
        );

        // Create the complete LLVM coverage data value to add to the LLVM IR
        cx.const_struct(&[cov_data_header_val, filenames_val], /*packed=*/ false)
    }
}

/// Construct a function record and combine it with the function's coverage mapping data.
/// Save the function record into the LLVM IR as a static global using a
/// specific, well-known section and name.
fn save_function_record(
    cx: &CodegenCx<'ll, 'tcx>,
    mangled_function_name: String,
    function_source_hash: u64,
    filenames_ref: u64,
    coverage_mapping_buffer: Vec<u8>,
) {
    // Concatenate the encoded coverage mappings
    let coverage_mapping_size = coverage_mapping_buffer.len();
    let coverage_mapping_val = cx.const_bytes(&coverage_mapping_buffer[..]);

    let func_name_hash = coverageinfo::hash_str(&mangled_function_name);
    let func_name_hash_val = cx.const_u64(func_name_hash);
    let coverage_mapping_size_val = cx.const_u32(coverage_mapping_size as u32);
    let func_hash_val = cx.const_u64(function_source_hash);
    let filenames_ref_val = cx.const_u64(filenames_ref);
    let func_record_val = cx.const_struct(
        &[
            func_name_hash_val,
            coverage_mapping_size_val,
            func_hash_val,
            filenames_ref_val,
            coverage_mapping_val,
        ],
        /*packed=*/ true,
    );

    // At the present time, the coverage map for Rust assumes every instrumented function `is_used`.
    // Note that Clang marks functions as "unused" in `CodeGenPGO::emitEmptyCounterMapping`. (See:
    // https://github.com/rust-lang/llvm-project/blob/de02a75e398415bad4df27b4547c25b896c8bf3b/clang%2Flib%2FCodeGen%2FCodeGenPGO.cpp#L877-L878
    // for example.)
    //
    // It's not yet clear if or how this may be applied to Rust in the future, but the `is_used`
    // argument is available and handled similarly.
    let is_used = true;
    coverageinfo::save_func_record_to_mod(cx, func_name_hash, func_record_val, is_used);
}
<|endoftext|>


<commit_before>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{BlockingSend, MultiPeerBackend, Socket, SocketBackend, SocketType};
use futures::channel::oneshot;

use async_trait::async_trait;
use dashmap::DashMap;
use futures::SinkExt;
use futures_codec::FramedWrite;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::sync::Arc;

pub(crate) struct Subscriber {
    pub(crate) subscriptions: Vec<Vec<u8>>,
    pub(crate) send_queue: FramedWrite<Box<dyn FrameableWrite>, ZmqCodec>,
    _subscription_coro_stop: oneshot::Sender<()>,
}

pub(crate) struct PubSocketBackend {
    subscribers: DashMap<PeerIdentity, Subscriber>,
}

impl PubSocketBackend {
    fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let message = match message {
            Message::Message(m) => m,
            _ => return,
        };
        let data: Vec<u8> = message.into();
        if data.is_empty() {
            return;
        }
        match data[0] {
            1 => {
                // Subscribe
                self.subscribers
                    .get_mut(&peer_id)
                    .unwrap()
                    .subscriptions
                    .push(Vec::from(&data[1..]));
            }
            0 => {
                // Unsubscribe
                let mut del_index = None;
                let sub = Vec::from(&data[1..]);
                for (idx, subscription) in self
                    .subscribers
                    .get(&peer_id)
                    .unwrap()
                    .subscriptions
                    .iter()
                    .enumerate()
                {
                    if &sub == subscription {
                        del_index = Some(idx);
                        break;
                    }
                }
                if let Some(index) = del_index {
                    self.subscribers
                        .get_mut(&peer_id)
                        .unwrap()
                        .subscriptions
                        .remove(index);
                }
            }
            _ => (),
        }
    }
}

impl SocketBackend for PubSocketBackend {
    fn socket_type(&self) -> SocketType {
        SocketType::PUB
    }

    fn shutdown(&self) {
        self.subscribers.clear();
    }
}

impl MultiPeerBackend for PubSocketBackend {
    fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (mut recv_queue, send_queue) = io.into_parts();
        // TODO provide handling for recv_queue
        let (sender, mut stop_receiver) = oneshot::channel();
        self.subscribers.insert(
            peer_id.clone(),
            Subscriber {
                subscriptions: vec![],
                send_queue,
                _subscription_coro_stop: sender,
            },
        );
        let backend = self.clone();
        let peer_id = peer_id.clone();
        tokio::spawn(async move {
            use futures::StreamExt;
            loop {
                tokio::select! {
                     _ = &mut stop_receiver => {
                         break;
                     },
                     message = &mut recv_queue.next() => {
                        match message {
                            Some(Ok(m)) => backend.message_received(&peer_id, m),
                            Some(Err(e)) => {
                                dbg!(e);
                                backend.peer_disconnected(&peer_id);
                                break;
                            }
                            None => {
                                backend.peer_disconnected(&peer_id);
                                break
                            }
                        }

                     }
                }
            }
        });
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        log::info!("Client disconnected {:?}", peer_id);
        self.subscribers.remove(peer_id);
    }
}

pub struct PubSocket {
    pub(crate) backend: Arc<PubSocketBackend>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for PubSocket {
    fn drop(&mut self) {
        self.backend.shutdown();
    }
}

#[async_trait]
impl BlockingSend for PubSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        let mut dead_peers = Vec::new();
        for mut subscriber in self.backend.subscribers.iter_mut() {
            for sub_filter in &subscriber.subscriptions {
                if sub_filter.as_slice() == &message.data[0..sub_filter.len()] {
                    let res = subscriber
                        .send_queue
                        .send(Message::Message(message.clone()))
                        .await;
                    match res {
                        Ok(()) => {}
                        Err(CodecError::Io(e)) => {
                            if e.kind() == ErrorKind::BrokenPipe {
                                dead_peers.push(subscriber.key().clone());
                            } else {
                                dbg!(e);
                            }
                        }
                        Err(e) => {
                            dbg!(e);
                            todo!()
                        }
                    }
                    break;
                }
            }
        }
        for peer in dead_peers {
            self.backend.peer_disconnected(&peer);
        }
        Ok(())
    }
}

#[async_trait]
impl Socket for PubSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(PubSocketBackend {
                subscribers: DashMap::new(),
            }),
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::util::tests::{
        test_bind_to_any_port_helper, test_bind_to_unspecified_interface_helper,
    };
    use crate::ZmqResult;
    use std::net::IpAddr;

    #[tokio::test]
    async fn test_bind_to_any_port() -> ZmqResult<()> {
        let s = PubSocket::new();
        test_bind_to_any_port_helper(s).await
    }

    #[tokio::test]
    async fn test_bind_to_any_ipv4_interface() -> ZmqResult<()> {
        let any_ipv4: IpAddr = "0.0.0.0".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv4, s, 4000).await
    }

    #[tokio::test]
    async fn test_bind_to_any_ipv6_interface() -> ZmqResult<()> {
        let any_ipv6: IpAddr = "::".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv6, s, 4010).await
    }
}
<commit_msg>todo marker for PUB issue<commit_after>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{BlockingSend, MultiPeerBackend, Socket, SocketBackend, SocketType};
use futures::channel::oneshot;

use async_trait::async_trait;
use dashmap::DashMap;
use futures::SinkExt;
use futures_codec::FramedWrite;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::sync::Arc;

pub(crate) struct Subscriber {
    pub(crate) subscriptions: Vec<Vec<u8>>,
    pub(crate) send_queue: FramedWrite<Box<dyn FrameableWrite>, ZmqCodec>,
    _subscription_coro_stop: oneshot::Sender<()>,
}

pub(crate) struct PubSocketBackend {
    subscribers: DashMap<PeerIdentity, Subscriber>,
}

impl PubSocketBackend {
    fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let message = match message {
            Message::Message(m) => m,
            _ => return,
        };
        let data: Vec<u8> = message.into();
        if data.is_empty() {
            return;
        }
        match data[0] {
            1 => {
                // Subscribe
                self.subscribers
                    .get_mut(&peer_id)
                    .unwrap()
                    .subscriptions
                    .push(Vec::from(&data[1..]));
            }
            0 => {
                // Unsubscribe
                let mut del_index = None;
                let sub = Vec::from(&data[1..]);
                for (idx, subscription) in self
                    .subscribers
                    .get(&peer_id)
                    .unwrap()
                    .subscriptions
                    .iter()
                    .enumerate()
                {
                    if &sub == subscription {
                        del_index = Some(idx);
                        break;
                    }
                }
                if let Some(index) = del_index {
                    self.subscribers
                        .get_mut(&peer_id)
                        .unwrap()
                        .subscriptions
                        .remove(index);
                }
            }
            _ => (),
        }
    }
}

impl SocketBackend for PubSocketBackend {
    fn socket_type(&self) -> SocketType {
        SocketType::PUB
    }

    fn shutdown(&self) {
        self.subscribers.clear();
    }
}

impl MultiPeerBackend for PubSocketBackend {
    fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (mut recv_queue, send_queue) = io.into_parts();
        // TODO provide handling for recv_queue
        let (sender, mut stop_receiver) = oneshot::channel();
        self.subscribers.insert(
            peer_id.clone(),
            Subscriber {
                subscriptions: vec![],
                send_queue,
                _subscription_coro_stop: sender,
            },
        );
        let backend = self.clone();
        let peer_id = peer_id.clone();
        tokio::spawn(async move {
            use futures::StreamExt;
            loop {
                tokio::select! {
                     _ = &mut stop_receiver => {
                         break;
                     },
                     message = &mut recv_queue.next() => {
                        match message {
                            Some(Ok(m)) => backend.message_received(&peer_id, m),
                            Some(Err(e)) => {
                                dbg!(e);
                                backend.peer_disconnected(&peer_id);
                                break;
                            }
                            None => {
                                backend.peer_disconnected(&peer_id);
                                break
                            }
                        }

                     }
                }
            }
        });
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        log::info!("Client disconnected {:?}", peer_id);
        self.subscribers.remove(peer_id);
    }
}

pub struct PubSocket {
    pub(crate) backend: Arc<PubSocketBackend>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for PubSocket {
    fn drop(&mut self) {
        self.backend.shutdown();
    }
}

#[async_trait]
impl BlockingSend for PubSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        let mut dead_peers = Vec::new();
        for mut subscriber in self.backend.subscribers.iter_mut() {
            for sub_filter in &subscriber.subscriptions {
                if sub_filter.as_slice() == &message.data[0..sub_filter.len()] {
                    // TODO should use try_send here
                    let res = subscriber
                        .send_queue
                        .send(Message::Message(message.clone()))
                        .await;
                    match res {
                        Ok(()) => {}
                        Err(CodecError::Io(e)) => {
                            if e.kind() == ErrorKind::BrokenPipe {
                                dead_peers.push(subscriber.key().clone());
                            } else {
                                dbg!(e);
                            }
                        }
                        Err(e) => {
                            dbg!(e);
                            todo!()
                        }
                    }
                    break;
                }
            }
        }
        for peer in dead_peers {
            self.backend.peer_disconnected(&peer);
        }
        Ok(())
    }
}

#[async_trait]
impl Socket for PubSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(PubSocketBackend {
                subscribers: DashMap::new(),
            }),
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::util::tests::{
        test_bind_to_any_port_helper, test_bind_to_unspecified_interface_helper,
    };
    use crate::ZmqResult;
    use std::net::IpAddr;

    #[tokio::test]
    async fn test_bind_to_any_port() -> ZmqResult<()> {
        let s = PubSocket::new();
        test_bind_to_any_port_helper(s).await
    }

    #[tokio::test]
    async fn test_bind_to_any_ipv4_interface() -> ZmqResult<()> {
        let any_ipv4: IpAddr = "0.0.0.0".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv4, s, 4000).await
    }

    #[tokio::test]
    async fn test_bind_to_any_ipv6_interface() -> ZmqResult<()> {
        let any_ipv6: IpAddr = "::".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv6, s, 4010).await
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The Rouille developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Everything in this module is private, but is still publicly accessible from the outside
//! because of the `try_or_400!` macro.

use std::error::Error;

/// This macro assumes that the current function returns a `Response` and takes a `Result`.
/// If the expression you pass to the macro is an error, then a 400 response is returned.
///
/// # Example
///
/// ```
/// # #[macro_use] extern crate rouille;
/// # fn main() {
/// use rouille::Request;
/// use rouille::Response;
///
/// fn handle_something(request: &Request) -> Response {
///     let data = try_or_400!(post_input!(request, {
///         field1: u32,
///         field2: String,
///     }));
///
///     Response::text("hello")
/// }
/// # }
/// ```
#[macro_export]
macro_rules! try_or_400 {
    ($result:expr) => (
        match $result {
            Ok(r) => r,
            Err(err) => {
                let json = $crate::try_or_400::ErrJson::from_err(&err);
                return $crate::Response::json(&json)
                    .with_status_code(400)
            },
        }
    );
}

#[derive(Serialize)]
pub struct ErrJson<'a> {
    description: &'a str,
    cause: Option<Box<ErrJson<'a>>>,
}

impl<'a> ErrJson<'a> {
    pub fn from_err<E: ?Sized + Error>(err: &'a E) -> ErrJson<'a> {
        let cause = err.cause().map(ErrJson::from_err).map(Box::new);
        ErrJson { description: err.description(), cause }
    }
}
<commit_msg>Error#cause is deprecated<commit_after>// Copyright (c) 2016 The Rouille developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Everything in this module is private, but is still publicly accessible from the outside
//! because of the `try_or_400!` macro.

use std::error::Error;

/// This macro assumes that the current function returns a `Response` and takes a `Result`.
/// If the expression you pass to the macro is an error, then a 400 response is returned.
///
/// # Example
///
/// ```
/// # #[macro_use] extern crate rouille;
/// # fn main() {
/// use rouille::Request;
/// use rouille::Response;
///
/// fn handle_something(request: &Request) -> Response {
///     let data = try_or_400!(post_input!(request, {
///         field1: u32,
///         field2: String,
///     }));
///
///     Response::text("hello")
/// }
/// # }
/// ```
#[macro_export]
macro_rules! try_or_400 {
    ($result:expr) => (
        match $result {
            Ok(r) => r,
            Err(err) => {
                let json = $crate::try_or_400::ErrJson::from_err(&err);
                return $crate::Response::json(&json)
                    .with_status_code(400)
            },
        }
    );
}

#[derive(Serialize)]
pub struct ErrJson<'a> {
    description: &'a str,
    cause: Option<Box<ErrJson<'a>>>,
}

impl<'a> ErrJson<'a> {
    pub fn from_err<E: ?Sized + Error>(err: &'a E) -> ErrJson<'a> {
        let cause = err.source().map(ErrJson::from_err).map(Box::new);
        ErrJson { description: err.description(), cause }
    }
}
<|endoftext|>


<commit_before>use ide_db::{base_db::FileId, defs::Definition, search::FileReference};
use syntax::{
    algo::find_node_at_range,
    ast::{self, ArgListOwner},
    AstNode, SourceFile, SyntaxKind, SyntaxNode, TextRange, T,
};

use SyntaxKind::WHITESPACE;

use crate::{
    assist_context::AssistBuilder, utils::next_prev, AssistContext, AssistId, AssistKind, Assists,
};

// Assist: remove_unused_param
//
// Removes unused function parameter.
//
// ```
// fn frobnicate(x: i32$0) {}
//
// fn main() {
//     frobnicate(92);
// }
// ```
// ->
// ```
// fn frobnicate() {}
//
// fn main() {
//     frobnicate();
// }
// ```
pub(crate) fn remove_unused_param(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let param: ast::Param = ctx.find_node_at_offset()?;
    let ident_pat = match param.pat()? {
        ast::Pat::IdentPat(it) => it,
        _ => return None,
    };
    let func = param.syntax().ancestors().find_map(ast::Fn::cast)?;
    let param_position = func.param_list()?.params().position(|it| it == param)?;

    let fn_def = {
        let func = ctx.sema.to_def(&func)?;
        Definition::ModuleDef(func.into())
    };

    let param_def = {
        let local = ctx.sema.to_def(&ident_pat)?;
        Definition::Local(local)
    };
    if param_def.usages(&ctx.sema).at_least_one() {
        cov_mark::hit!(keep_used);
        return None;
    }
    acc.add(
        AssistId("remove_unused_param", AssistKind::Refactor),
        "Remove unused parameter",
        param.syntax().text_range(),
        |builder| {
            builder.delete(range_to_remove(param.syntax()));
            for (file_id, references) in fn_def.usages(&ctx.sema).all() {
                process_usages(ctx, builder, file_id, references, param_position);
            }
        },
    )
}

fn process_usages(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    file_id: FileId,
    references: Vec<FileReference>,
    arg_to_remove: usize,
) {
    let source_file = ctx.sema.parse(file_id);
    builder.edit_file(file_id);
    for usage in references {
        if let Some(text_range) = process_usage(&source_file, usage, arg_to_remove) {
            builder.delete(text_range);
        }
    }
}

fn process_usage(
    source_file: &SourceFile,
    FileReference { range, .. }: FileReference,
    arg_to_remove: usize,
) -> Option<TextRange> {
    let call_expr: ast::CallExpr = find_node_at_range(source_file.syntax(), range)?;
    let call_expr_range = call_expr.expr()?.syntax().text_range();
    if !call_expr_range.contains_range(range) {
        return None;
    }
    let arg = call_expr.arg_list()?.args().nth(arg_to_remove)?;
    Some(range_to_remove(arg.syntax()))
}

fn range_to_remove(node: &SyntaxNode) -> TextRange {
    let up_to_comma = next_prev().find_map(|dir| {
        node.siblings_with_tokens(dir)
            .filter_map(|it| it.into_token())
            .find(|it| it.kind() == T![,])
            .map(|it| (dir, it))
    });
    if let Some((dir, token)) = up_to_comma {
        if node.next_sibling().is_some() {
            let up_to_space = token
                .siblings_with_tokens(dir)
                .skip(1)
                .take_while(|it| it.kind() == WHITESPACE)
                .last()
                .and_then(|it| it.into_token());
            return node
                .text_range()
                .cover(up_to_space.map_or(token.text_range(), |it| it.text_range()));
        }
        node.text_range().cover(token.text_range())
    } else {
        node.text_range()
    }
}

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

    use super::*;

    #[test]
    fn remove_unused() {
        check_assist(
            remove_unused_param,
            r#"
fn a() { foo(9, 2) }
fn foo(x: i32, $0y: i32) { x; }
fn b() { foo(9, 2,) }
"#,
            r#"
fn a() { foo(9) }
fn foo(x: i32) { x; }
fn b() { foo(9, ) }
"#,
        );
    }

    #[test]
    fn remove_unused_first_param() {
        check_assist(
            remove_unused_param,
            r#"
fn foo($0x: i32, y: i32) { y; }
fn a() { foo(1, 2) }
fn b() { foo(1, 2,) }
"#,
            r#"
fn foo(y: i32) { y; }
fn a() { foo(2) }
fn b() { foo(2,) }
"#,
        );
    }

    #[test]
    fn remove_unused_single_param() {
        check_assist(
            remove_unused_param,
            r#"
fn foo($0x: i32) { 0; }
fn a() { foo(1) }
fn b() { foo(1, ) }
"#,
            r#"
fn foo() { 0; }
fn a() { foo() }
fn b() { foo( ) }
"#,
        );
    }

    #[test]
    fn remove_unused_surrounded_by_parms() {
        check_assist(
            remove_unused_param,
            r#"
fn foo(x: i32, $0y: i32, z: i32) { x; }
fn a() { foo(1, 2, 3) }
fn b() { foo(1, 2, 3,) }
"#,
            r#"
fn foo(x: i32, z: i32) { x; }
fn a() { foo(1, 3) }
fn b() { foo(1, 3,) }
"#,
        );
    }

    #[test]
    fn remove_unused_qualified_call() {
        check_assist(
            remove_unused_param,
            r#"
mod bar { pub fn foo(x: i32, $0y: i32) { x; } }
fn b() { bar::foo(9, 2) }
"#,
            r#"
mod bar { pub fn foo(x: i32) { x; } }
fn b() { bar::foo(9) }
"#,
        );
    }

    #[test]
    fn remove_unused_turbofished_func() {
        check_assist(
            remove_unused_param,
            r#"
pub fn foo<T>(x: T, $0y: i32) { x; }
fn b() { foo::<i32>(9, 2) }
"#,
            r#"
pub fn foo<T>(x: T) { x; }
fn b() { foo::<i32>(9) }
"#,
        );
    }

    #[test]
    fn remove_unused_generic_unused_param_func() {
        check_assist(
            remove_unused_param,
            r#"
pub fn foo<T>(x: i32, $0y: T) { x; }
fn b() { foo::<i32>(9, 2) }
fn b2() { foo(9, 2) }
"#,
            r#"
pub fn foo<T>(x: i32) { x; }
fn b() { foo::<i32>(9) }
fn b2() { foo(9) }
"#,
        );
    }

    #[test]
    fn keep_used() {
        cov_mark::check!(keep_used);
        check_assist_not_applicable(
            remove_unused_param,
            r#"
fn foo(x: i32, $0y: i32) { y; }
fn main() { foo(9, 2) }
"#,
        );
    }

    #[test]
    fn remove_across_files() {
        check_assist(
            remove_unused_param,
            r#"
//- /main.rs
fn foo(x: i32, $0y: i32) { x; }

mod foo;

//- /foo.rs
use super::foo;

fn bar() {
    let _ = foo(1, 2);
}
"#,
            r#"
//- /main.rs
fn foo(x: i32) { x; }

mod foo;

//- /foo.rs
use super::foo;

fn bar() {
    let _ = foo(1);
}
"#,
        )
    }
}
<commit_msg>Don't suggest Remove unused param in trait impls<commit_after>use ide_db::{base_db::FileId, defs::Definition, search::FileReference};
use syntax::{
    algo::find_node_at_range,
    ast::{self, ArgListOwner},
    AstNode, SourceFile, SyntaxKind, SyntaxNode, TextRange, T,
};

use SyntaxKind::WHITESPACE;

use crate::{
    assist_context::AssistBuilder, utils::next_prev, AssistContext, AssistId, AssistKind, Assists,
};

// Assist: remove_unused_param
//
// Removes unused function parameter.
//
// ```
// fn frobnicate(x: i32$0) {}
//
// fn main() {
//     frobnicate(92);
// }
// ```
// ->
// ```
// fn frobnicate() {}
//
// fn main() {
//     frobnicate();
// }
// ```
pub(crate) fn remove_unused_param(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let param: ast::Param = ctx.find_node_at_offset()?;
    let ident_pat = match param.pat()? {
        ast::Pat::IdentPat(it) => it,
        _ => return None,
    };
    let func = param.syntax().ancestors().find_map(ast::Fn::cast)?;

    // check if fn is in impl Trait for ..
    if func
        .syntax()
        .parent() // AssocItemList
        .and_then(|x| x.parent())
        .and_then(ast::Impl::cast)
        .map_or(false, |imp| imp.trait_().is_some())
    {
        cov_mark::hit!(trait_impl);
        return None;
    }

    let param_position = func.param_list()?.params().position(|it| it == param)?;
    let fn_def = {
        let func = ctx.sema.to_def(&func)?;
        Definition::ModuleDef(func.into())
    };

    let param_def = {
        let local = ctx.sema.to_def(&ident_pat)?;
        Definition::Local(local)
    };
    if param_def.usages(&ctx.sema).at_least_one() {
        cov_mark::hit!(keep_used);
        return None;
    }
    acc.add(
        AssistId("remove_unused_param", AssistKind::Refactor),
        "Remove unused parameter",
        param.syntax().text_range(),
        |builder| {
            builder.delete(range_to_remove(param.syntax()));
            for (file_id, references) in fn_def.usages(&ctx.sema).all() {
                process_usages(ctx, builder, file_id, references, param_position);
            }
        },
    )
}

fn process_usages(
    ctx: &AssistContext,
    builder: &mut AssistBuilder,
    file_id: FileId,
    references: Vec<FileReference>,
    arg_to_remove: usize,
) {
    let source_file = ctx.sema.parse(file_id);
    builder.edit_file(file_id);
    for usage in references {
        if let Some(text_range) = process_usage(&source_file, usage, arg_to_remove) {
            builder.delete(text_range);
        }
    }
}

fn process_usage(
    source_file: &SourceFile,
    FileReference { range, .. }: FileReference,
    arg_to_remove: usize,
) -> Option<TextRange> {
    let call_expr: ast::CallExpr = find_node_at_range(source_file.syntax(), range)?;
    let call_expr_range = call_expr.expr()?.syntax().text_range();
    if !call_expr_range.contains_range(range) {
        return None;
    }
    let arg = call_expr.arg_list()?.args().nth(arg_to_remove)?;
    Some(range_to_remove(arg.syntax()))
}

fn range_to_remove(node: &SyntaxNode) -> TextRange {
    let up_to_comma = next_prev().find_map(|dir| {
        node.siblings_with_tokens(dir)
            .filter_map(|it| it.into_token())
            .find(|it| it.kind() == T![,])
            .map(|it| (dir, it))
    });
    if let Some((dir, token)) = up_to_comma {
        if node.next_sibling().is_some() {
            let up_to_space = token
                .siblings_with_tokens(dir)
                .skip(1)
                .take_while(|it| it.kind() == WHITESPACE)
                .last()
                .and_then(|it| it.into_token());
            return node
                .text_range()
                .cover(up_to_space.map_or(token.text_range(), |it| it.text_range()));
        }
        node.text_range().cover(token.text_range())
    } else {
        node.text_range()
    }
}

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

    use super::*;

    #[test]
    fn remove_unused() {
        check_assist(
            remove_unused_param,
            r#"
fn a() { foo(9, 2) }
fn foo(x: i32, $0y: i32) { x; }
fn b() { foo(9, 2,) }
"#,
            r#"
fn a() { foo(9) }
fn foo(x: i32) { x; }
fn b() { foo(9, ) }
"#,
        );
    }

    #[test]
    fn remove_unused_first_param() {
        check_assist(
            remove_unused_param,
            r#"
fn foo($0x: i32, y: i32) { y; }
fn a() { foo(1, 2) }
fn b() { foo(1, 2,) }
"#,
            r#"
fn foo(y: i32) { y; }
fn a() { foo(2) }
fn b() { foo(2,) }
"#,
        );
    }

    #[test]
    fn remove_unused_single_param() {
        check_assist(
            remove_unused_param,
            r#"
fn foo($0x: i32) { 0; }
fn a() { foo(1) }
fn b() { foo(1, ) }
"#,
            r#"
fn foo() { 0; }
fn a() { foo() }
fn b() { foo( ) }
"#,
        );
    }

    #[test]
    fn remove_unused_surrounded_by_parms() {
        check_assist(
            remove_unused_param,
            r#"
fn foo(x: i32, $0y: i32, z: i32) { x; }
fn a() { foo(1, 2, 3) }
fn b() { foo(1, 2, 3,) }
"#,
            r#"
fn foo(x: i32, z: i32) { x; }
fn a() { foo(1, 3) }
fn b() { foo(1, 3,) }
"#,
        );
    }

    #[test]
    fn remove_unused_qualified_call() {
        check_assist(
            remove_unused_param,
            r#"
mod bar { pub fn foo(x: i32, $0y: i32) { x; } }
fn b() { bar::foo(9, 2) }
"#,
            r#"
mod bar { pub fn foo(x: i32) { x; } }
fn b() { bar::foo(9) }
"#,
        );
    }

    #[test]
    fn remove_unused_turbofished_func() {
        check_assist(
            remove_unused_param,
            r#"
pub fn foo<T>(x: T, $0y: i32) { x; }
fn b() { foo::<i32>(9, 2) }
"#,
            r#"
pub fn foo<T>(x: T) { x; }
fn b() { foo::<i32>(9) }
"#,
        );
    }

    #[test]
    fn remove_unused_generic_unused_param_func() {
        check_assist(
            remove_unused_param,
            r#"
pub fn foo<T>(x: i32, $0y: T) { x; }
fn b() { foo::<i32>(9, 2) }
fn b2() { foo(9, 2) }
"#,
            r#"
pub fn foo<T>(x: i32) { x; }
fn b() { foo::<i32>(9) }
fn b2() { foo(9) }
"#,
        );
    }

    #[test]
    fn keep_used() {
        cov_mark::check!(keep_used);
        check_assist_not_applicable(
            remove_unused_param,
            r#"
fn foo(x: i32, $0y: i32) { y; }
fn main() { foo(9, 2) }
"#,
        );
    }

    #[test]
    fn trait_impl() {
        cov_mark::check!(trait_impl);
        check_assist_not_applicable(
            remove_unused_param,
            r#"
trait Trait {
    fn foo(x: i32);
}
impl Trait for () {
    fn foo($0x: i32) {}
}
"#,
        );
    }

    #[test]
    fn remove_across_files() {
        check_assist(
            remove_unused_param,
            r#"
//- /main.rs
fn foo(x: i32, $0y: i32) { x; }

mod foo;

//- /foo.rs
use super::foo;

fn bar() {
    let _ = foo(1, 2);
}
"#,
            r#"
//- /main.rs
fn foo(x: i32) { x; }

mod foo;

//- /foo.rs
use super::foo;

fn bar() {
    let _ = foo(1);
}
"#,
        )
    }
}
<|endoftext|>


<commit_before>pub use self::Side::{Left, Right};
pub use self::Color::{Yellow, Blue};

use std::collections::BTreeMap;
use ::prelude::*;
use ::game;

const FIELD_LENGTH: f32    = 9.010;
const CENTER_DIAMETER: f32 = 1.000;
const DEFENSE_RADIUS: f32  = 1.000;
const ROBOT_RADIUS: f32    = 0.090;
const BALL_RADIUS: f32     = 0.023;
const BALL_FRICT_LOSS: f32 = 0.020;

#[derive(Clone, Debug, PartialEq)]
pub struct Vec2d {
    pub x: f32,
    pub y: f32,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Vec3d {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Ball {
    pub pos: Vec3d,
    pub vel: Vec3d,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Robot {
    pub pos: Vec2d,
    pub vel: Vec2d,
    pub w: f32,
    pub vw: f32,
}

impl Robot {
    #[inline]
    pub fn new() -> Robot {
        Robot { pos: Vec2d { x: 0.0, y: 0.0 }, w: 0.0, vel: Vec2d { x: 0.0, y: 0.0 }, vw: 0.0 }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Color {
    Yellow,
    Blue,
}

impl Color {
    #[inline] pub fn is_yellow(&self) -> bool { *self == Yellow }
    #[inline] pub fn is_blue(&self)   -> bool { *self == Blue }
    #[inline] pub fn yellow(is_yellow: bool) -> Color { if is_yellow { Yellow } else { Blue   } }
    #[inline] pub fn blue(is_blue: bool)     -> Color { if is_blue   { Blue   } else { Yellow } }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Side {
    Left,
    Right
}

impl Side {
    #[inline] pub fn is_right(&self) -> bool { *self == Right }
    #[inline] pub fn is_left(&self)  -> bool { *self == Left }
    #[inline] pub fn right(is_right: bool) -> Side { if is_right { Right } else { Left  } }
    #[inline] pub fn left(is_left: bool)   -> Side { if is_left  { Left  } else { Right } }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct TeamSide(pub Color, pub Side);

impl TeamSide {
    #[inline]
    pub fn yellow_is_left(&self) -> bool {
        match *self {
            TeamSide(Yellow, Left)  => true,
            TeamSide(Yellow, Right) => false,
            TeamSide(Blue,   Right) => true,
            TeamSide(Blue,   Left)  => false,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Id(pub Color, pub u8);

#[derive(Clone, Debug, PartialEq)]
pub struct State {
    pub timestamp: f32,
    pub ball: Ball,
    pub robots: BTreeMap<Id, Robot>,
    pub team_side: TeamSide,
}

trait BTreeMapSimExt {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side);
}

impl BTreeMapSimExt for BTreeMap<Id, Robot> {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side) {
        use std::f32::consts::PI;

        let w_0 = if side.is_right() { 0.0 } else { PI };
        let w_delta = 2.0 * PI / (count as f32) * if side.is_right() { 1.0 } else { -1.0 };
        let x_offset = (CENTER_DIAMETER / 4.0 + FIELD_LENGTH / 4.0 - DEFENSE_RADIUS / 2.0) * if side.is_right() { 1.0 } else { -1.0 };
        //let radius = FIELD_LENGTH / 8.0;
        let radius = (FIELD_LENGTH / 2.0 - CENTER_DIAMETER / 2.0 - DEFENSE_RADIUS) / 3.0 - ROBOT_RADIUS;

        for i in 0..count {
            let robot_id = Id(color, i);
            let robot = self.entry(robot_id).or_insert(Robot::new());
            let w = w_0 + (i as f32) * w_delta;
            robot.pos.x = radius * w.cos() + x_offset;
            robot.pos.y = radius * w.sin();
            robot.w = w + PI;
        }
    }
}

impl State {
    pub fn new(team_side: TeamSide) -> State {
        let mut robots = BTreeMap::new();
        if team_side.yellow_is_left() {
            robots.initial_formation(6, Yellow, Left);
            robots.initial_formation(6, Blue, Right);
        } else {
            robots.initial_formation(6, Yellow, Right);
            robots.initial_formation(6, Blue, Left);
        }
        State {
            timestamp: 0.0,
            //ball: Ball { pos: Vec3d { x: 0.0, y: 0.0, z: 0.0 }, vel: Vec3d { x: 0.0, y: 0.0, z: 0.0 } },
            ball: Ball { pos: Vec3d { x: 0.0, y: 0.05, z: 0.0 }, vel: Vec3d { x: -4.0, y: 0.0, z: 0.0 } },
            robots: robots,
            team_side: team_side,
        }
    }

    pub fn update_game(&self, game_state: &mut game::State) {
        // ball
        {
            let game_ball = game_state.get_ball_mut();
            game_ball.set_x(self.ball.pos.x);
            game_ball.set_y(self.ball.pos.y);
            //game_ball.set_z(self.ball.pos.z);
            game_ball.set_vx(self.ball.vel.x);
            game_ball.set_vy(self.ball.vel.y);
            //game_ball.set_vz(self.ball.vel.z);
        }
        // robots
        for (id, ref robot) in self.robots.iter() {
            let (robot_id, ref mut robots) = match id {
                &Id(Blue, i)   => (i, game_state.get_robots_blue_mut()),
                &Id(Yellow, i) => (i, game_state.get_robots_yellow_mut()),
            };
            let game_robot = robots.entry(robot_id).or_insert(game::Robot::new(robot_id));
            game_robot.set_x(robot.pos.x);
            game_robot.set_y(robot.pos.y);
            game_robot.set_w(robot.w);
            game_robot.set_vx(robot.vel.x);
            game_robot.set_vy(robot.vel.y);
            game_robot.set_vw(robot.vw);
        }
        game_state.inc_counter();
    }

    pub fn step(&mut self, command: game::Command, timestep: f32) {
        self.timestamp += timestep;
        let color = Color::yellow(command.is_yellow);
        let &mut State { ref mut ball, ref mut robots, .. } = self;
        let mut d_time_ball = timestep;

        // XXX: overly simplified physics ahead
        for (id, robot_command) in command.robots {
            let robot_id = &Id(color, id);
            if let Some(mut robot) = robots.get_mut(robot_id) {
                let d_time = timestep;

                let v_tangent = robot_command.v_tangent;
                let v_normal  = robot_command.v_normal;
                let v_angular = robot_command.v_angular;

                let w = robot.w;

                let vx = v_normal * w.sin() + v_tangent * w.cos();
                let vy = v_normal * w.cos() - v_tangent * w.sin();
                let vw = v_angular;

                let dx = d_time * vx;
                let dy = d_time * vy;
                let dw = d_time * vw;

                // detect collision with ball
                {
                    let rr = ROBOT_RADIUS;
                    let rb = BALL_RADIUS;
                    let xr  = robot.pos.x;
                    let yr  = robot.pos.y;
                    let vrx = robot.pos.x;
                    let vry = robot.pos.y;
                    let xb  = ball.pos.x;
                    let yb  = ball.pos.y;
                    let vbx = ball.vel.x;
                    let vby = ball.vel.y;

                    // Bhāskara:
                    let a = (vry - vby) * (vry - vby) + (vrx - vbx) * (vrx - vbx);
                    let b = 2.0 * ((xr - xb) * (vrx - vbx) + (yr - yb) * (vry - vby));
                    let c = (xr - xb) * (xr - xb) + (yr - yb) * (yr - yb) - (rr + rb) * (rr + rb);

                    if c > 0.0 && a != 0.0 {
                        //debug!("there may be a collision");
                        let tc = (-b - (b * b - 4.0 * a * c).sqrt()) * (0.5 / a);
                        if tc >= 0.0 && tc <= timestep {
                            use std::f32::consts::PI;

                            let xbc = xb + tc * vbx;
                            let ybc = yb + tc * vby;
                            let xrc = xr + tc * vrx;
                            let yrc = yr + tc * vry;
                            debug!("there is a collision, with robot {:?} at {},{}_{},{}", robot_id, xrc, yrc, xbc, ybc);
                            // TODO: dummy collision
                            ball.pos.x = xbc;
                            ball.pos.y = ybc;
                            let wc = (ybc - yrc).atan2(xbc - xrc);
                            let wb = vby.atan2(vbx);
                            let v = (vbx * vbx + vby * vby).sqrt();
                            let w = PI + wc - wb;
                            ball.vel.x = v * w.cos() + vrx;
                            ball.vel.y = v * w.sin() + vry;
                            d_time_ball -= tc;
                        }
                    }
                }

                //println!("dx: {}", dx);

                robot.pos.x += dx;
                robot.pos.y += dy;
                robot.w += dw;

                // TODO: effect of robot_command.action
            }
        }

        // update ball pos
        {
            ball.pos.x += d_time_ball * ball.vel.x;
            ball.pos.y += d_time_ball * ball.vel.y;
            ball.pos.z += d_time_ball * ball.vel.z;
            let r = 1.0 - BALL_FRICT_LOSS;
            ball.vel.x *= r;
            ball.vel.y *= r;
            ball.vel.z *= r;
        }
    }
}
<commit_msg>Fix collision input parametrs.<commit_after>pub use self::Side::{Left, Right};
pub use self::Color::{Yellow, Blue};

use std::collections::BTreeMap;
use ::prelude::*;
use ::game;

const FIELD_LENGTH: f32    = 9.010;
const CENTER_DIAMETER: f32 = 1.000;
const DEFENSE_RADIUS: f32  = 1.000;
const ROBOT_RADIUS: f32    = 0.090;
const BALL_RADIUS: f32     = 0.023;
const BALL_FRICT_LOSS: f32 = 1.000;

#[derive(Clone, Debug, PartialEq)]
pub struct Vec2d {
    pub x: f32,
    pub y: f32,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Vec3d {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Ball {
    pub pos: Vec3d,
    pub vel: Vec3d,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Robot {
    pub pos: Vec2d,
    pub vel: Vec2d,
    pub w: f32,
    pub vw: f32,
}

impl Robot {
    #[inline]
    pub fn new() -> Robot {
        Robot { pos: Vec2d { x: 0.0, y: 0.0 }, w: 0.0, vel: Vec2d { x: 0.0, y: 0.0 }, vw: 0.0 }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Color {
    Yellow,
    Blue,
}

impl Color {
    #[inline] pub fn is_yellow(&self) -> bool { *self == Yellow }
    #[inline] pub fn is_blue(&self)   -> bool { *self == Blue }
    #[inline] pub fn yellow(is_yellow: bool) -> Color { if is_yellow { Yellow } else { Blue   } }
    #[inline] pub fn blue(is_blue: bool)     -> Color { if is_blue   { Blue   } else { Yellow } }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Side {
    Left,
    Right
}

impl Side {
    #[inline] pub fn is_right(&self) -> bool { *self == Right }
    #[inline] pub fn is_left(&self)  -> bool { *self == Left }
    #[inline] pub fn right(is_right: bool) -> Side { if is_right { Right } else { Left  } }
    #[inline] pub fn left(is_left: bool)   -> Side { if is_left  { Left  } else { Right } }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct TeamSide(pub Color, pub Side);

impl TeamSide {
    #[inline]
    pub fn yellow_is_left(&self) -> bool {
        match *self {
            TeamSide(Yellow, Left)  => true,
            TeamSide(Yellow, Right) => false,
            TeamSide(Blue,   Right) => true,
            TeamSide(Blue,   Left)  => false,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Id(pub Color, pub u8);

#[derive(Clone, Debug, PartialEq)]
pub struct State {
    pub timestamp: f32,
    pub ball: Ball,
    pub robots: BTreeMap<Id, Robot>,
    pub team_side: TeamSide,
}

trait BTreeMapSimExt {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side);
}

impl BTreeMapSimExt for BTreeMap<Id, Robot> {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side) {
        use std::f32::consts::PI;

        let w_0 = if side.is_right() { 0.0 } else { PI };
        let w_delta = 2.0 * PI / (count as f32) * if side.is_right() { 1.0 } else { -1.0 };
        let x_offset = (CENTER_DIAMETER / 4.0 + FIELD_LENGTH / 4.0 - DEFENSE_RADIUS / 2.0) * if side.is_right() { 1.0 } else { -1.0 };
        //let radius = FIELD_LENGTH / 8.0;
        let radius = (FIELD_LENGTH / 2.0 - CENTER_DIAMETER / 2.0 - DEFENSE_RADIUS) / 3.0 - ROBOT_RADIUS;

        for i in 0..count {
            let robot_id = Id(color, i);
            let robot = self.entry(robot_id).or_insert(Robot::new());
            let w = w_0 + (i as f32) * w_delta;
            robot.pos.x = radius * w.cos() + x_offset;
            robot.pos.y = radius * w.sin();
            robot.w = w + PI;
        }
    }
}

impl State {
    pub fn new(team_side: TeamSide) -> State {
        let mut robots = BTreeMap::new();
        if team_side.yellow_is_left() {
            robots.initial_formation(6, Yellow, Left);
            robots.initial_formation(6, Blue, Right);
        } else {
            robots.initial_formation(6, Yellow, Right);
            robots.initial_formation(6, Blue, Left);
        }
        State {
            timestamp: 0.0,
            //ball: Ball { pos: Vec3d { x: 0.0, y: 0.0, z: 0.0 }, vel: Vec3d { x: 0.0, y: 0.0, z: 0.0 } },
            ball: Ball { pos: Vec3d { x: 0.0, y: 0.05, z: 0.0 }, vel: Vec3d { x: -4.0, y: 0.0, z: 0.0 } },
            robots: robots,
            team_side: team_side,
        }
    }

    pub fn update_game(&self, game_state: &mut game::State) {
        // ball
        {
            let game_ball = game_state.get_ball_mut();
            game_ball.set_x(self.ball.pos.x);
            game_ball.set_y(self.ball.pos.y);
            //game_ball.set_z(self.ball.pos.z);
            game_ball.set_vx(self.ball.vel.x);
            game_ball.set_vy(self.ball.vel.y);
            //game_ball.set_vz(self.ball.vel.z);
        }
        // robots
        for (id, ref robot) in self.robots.iter() {
            let (robot_id, ref mut robots) = match id {
                &Id(Blue, i)   => (i, game_state.get_robots_blue_mut()),
                &Id(Yellow, i) => (i, game_state.get_robots_yellow_mut()),
            };
            let game_robot = robots.entry(robot_id).or_insert(game::Robot::new(robot_id));
            game_robot.set_x(robot.pos.x);
            game_robot.set_y(robot.pos.y);
            game_robot.set_w(robot.w);
            game_robot.set_vx(robot.vel.x);
            game_robot.set_vy(robot.vel.y);
            game_robot.set_vw(robot.vw);
        }
        game_state.inc_counter();
    }

    pub fn step(&mut self, command: game::Command, timestep: f32) {
        self.timestamp += timestep;
        let color = Color::yellow(command.is_yellow);
        let &mut State { ref mut ball, ref mut robots, .. } = self;
        let mut d_time_ball = timestep;

        // XXX: overly simplified physics ahead
        for (id, robot_command) in command.robots {
            let robot_id = &Id(color, id);
            if let Some(mut robot) = robots.get_mut(robot_id) {
                let d_time = timestep;

                let v_tangent = robot_command.v_tangent;
                let v_normal  = robot_command.v_normal;
                let v_angular = robot_command.v_angular;

                let w = robot.w;

                let vx = v_normal * w.sin() + v_tangent * w.cos();
                let vy = v_normal * w.cos() - v_tangent * w.sin();
                let vw = v_angular;

                robot.vel.x = vx;
                robot.vel.y = vy;

                let dx = d_time * vx;
                let dy = d_time * vy;
                let dw = d_time * vw;

                // detect collision with ball
                {
                    let rr = ROBOT_RADIUS;
                    let rb = BALL_RADIUS;
                    let xr  = robot.pos.x;
                    let yr  = robot.pos.y;
                    let vrx = robot.vel.x;
                    let vry = robot.vel.y;
                    let xb  = ball.pos.x;
                    let yb  = ball.pos.y;
                    let vbx = ball.vel.x;
                    let vby = ball.vel.y;

                    // Bhāskara:
                    let a = (vry - vby) * (vry - vby) + (vrx - vbx) * (vrx - vbx);
                    let b = 2.0 * ((xr - xb) * (vrx - vbx) + (yr - yb) * (vry - vby));
                    let c = (xr - xb) * (xr - xb) + (yr - yb) * (yr - yb) - (rr + rb) * (rr + rb);

                    if c > 0.0 && a != 0.0 {
                        //debug!("there may be a collision");
                        let tc = (-b - (b * b - 4.0 * a * c).sqrt()) * (0.5 / a);
                        if tc >= 0.0 && tc <= timestep {
                            use std::f32::consts::PI;

                            let xbc = xb + tc * vbx;
                            let ybc = yb + tc * vby;
                            let xrc = xr + tc * vrx;
                            let yrc = yr + tc * vry;
                            debug!("there is a collision, with robot {:?} at {},{}_{},{}", robot_id, xrc, yrc, xbc, ybc);
                            // TODO: dummy collision
                            ball.pos.x = xbc;
                            ball.pos.y = ybc;
                            let wc = (ybc - yrc).atan2(xbc - xrc);
                            let wb = vby.atan2(vbx);
                            let v = (vbx * vbx + vby * vby).sqrt();
                            let w = PI + wc - wb;
                            ball.vel.x = v * w.cos() + 2.0 * vrx;
                            ball.vel.y = v * w.sin() + 2.0 * vry;
                            d_time_ball -= tc;
                        }
                    }
                }

                //println!("dx: {}", dx);

                robot.pos.x += dx;
                robot.pos.y += dy;
                robot.w += dw;

                // TODO: effect of robot_command.action
            }
        }

        // update ball pos
        {
            ball.pos.x += d_time_ball * ball.vel.x;
            ball.pos.y += d_time_ball * ball.vel.y;
            ball.pos.z += d_time_ball * ball.vel.z;
            let r = 1.0 - BALL_FRICT_LOSS * timestep;
            ball.vel.x *= r;
            ball.vel.y *= r;
            ball.vel.z *= r;
        }
    }
}
<|endoftext|>


<commit_before>trait StrParseExt {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool;
}

impl<'a> StrParseExt for &'a str {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool,
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len()),
        }
    }
}

pub trait XmlStr {
    /// Find the end of the quoted attribute value, not including the quote
    fn end_of_attribute(&self, quote: &str) -> Option<usize>;
    /// Find the end of the direct character data
    fn end_of_char_data(&self) -> Option<usize>;
    /// Find the end of the CData section, not including the ]]>
    fn end_of_cdata(&self) -> Option<usize>;
    /// Find the end of a run of decimal characters
    fn end_of_decimal_chars(&self) -> Option<usize>;
    /// Find the end of a run of hexidecimal characters
    fn end_of_hex_chars(&self) -> Option<usize>;
    /// Find the end of the comment, not including the -->
    fn end_of_comment(&self) -> Option<usize>;
    /// Find the end of the processing instruction, not including the ?>
    fn end_of_pi_value(&self) -> Option<usize>;
    /// Find the end of the [Name](http://www.w3.org/TR/xml/#NT-Name)
    fn end_of_name(&self) -> Option<usize>;
    /// Find the end of the [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)
    fn end_of_ncname(&self) -> Option<usize>;
    /// Find the end of a run of space characters
    fn end_of_space(&self) -> Option<usize>;
    /// Find the end of the starting tag
    fn end_of_start_tag(&self) -> Option<usize>;
    fn end_of_encoding(&self) -> Option<usize>;
    /// Find the end of the internal doc type declaration, not including the ]
    fn end_of_int_subset(&self) -> Option<usize>;
}

impl<'a> XmlStr for &'a str {
    fn end_of_attribute(&self, quote: &str) -> Option<usize> {
        if self.len() == 0 ||
           self.starts_with('&') ||
           self.starts_with('<') ||
           self.starts_with(quote)
        {
            return None;
        }

        let quote_char = quote.chars().next().expect("Cant have null quote");

        self.find(&['&', '<', quote_char][..]).or(Some(self.len()))
    }

    fn end_of_char_data(&self) -> Option<usize> {
        fn find_end_of_char_data(bytes: &[u8]) -> Option<usize> {
            for (i, &b) in bytes.iter().enumerate() {
                if b == b'<' || b == b'&' { return Some(i) }

                if b == b']' && bytes[i..].starts_with(b"]]>") {
                    return Some(i)
                }
            }
            None
        }

        match find_end_of_char_data(self.as_bytes()) {
            Some(0) => None,
            Some(v) => Some(v),
            None => Some(self.len()),
        }
    }

    fn end_of_cdata(&self) -> Option<usize> {
        self.find("]]>")
    }

    fn end_of_decimal_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_decimal_char(),
                               |c| c.is_decimal_char())
    }

    fn end_of_hex_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_hex_char(),
                               |c| c.is_hex_char())
    }

    fn end_of_comment(&self) -> Option<usize> {
        // This deliberately does not include the >. -- is not allowed
        // in a comment, so we can just test the end if it matches the
        // complete close delimiter.
        self.find("--")
    }

    fn end_of_pi_value(&self) -> Option<usize> {
        self.find("?>")
    }

    fn end_of_name(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn end_of_ncname(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_ncname_start_char(), |c| c.is_ncname_char())
    }

    fn end_of_space(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }

    fn end_of_start_tag(&self) -> Option<usize> {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if '<' == c => (),
            _ => return None,
        };

        match positions.next() {
            Some((offset, c)) =>
                match c {
                    '?' | '!' | '/' => None,
                    _ => Some(offset),
                },
            None => Some(self.len()),
        }
    }

    fn end_of_encoding(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_encoding_start_char(), |c| c.is_encoding_rest_char())
    }

    fn end_of_int_subset(&self) -> Option<usize> { self.find("]") }
}

/// Predicates used when parsing an characters in an XML document.
pub trait XmlChar {
    /// Is this a [NameStartChar](http://www.w3.org/TR/xml/#NT-NameStartChar)?
    fn is_name_start_char(self) -> bool;
    /// Is this a [NameChar](http://www.w3.org/TR/xml/#NT-NameChar)?
    fn is_name_char(self) -> bool;
    /// Does this start a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_start_char(self) -> bool;
    /// Is this a component of a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_char(self) -> bool;
    /// Is this an [XML space](http://www.w3.org/TR/xml/#NT-S)?
    fn is_space_char(self) -> bool;
    fn is_decimal_char(self) -> bool;
    fn is_hex_char(self) -> bool;
    fn is_encoding_start_char(self) -> bool;
    fn is_encoding_rest_char(self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(self) -> bool {
        self == ':' || self.is_ncname_start_char()
    }

    fn is_name_char(self) -> bool {
        self.is_name_start_char() || self.is_ncname_char()
    }

    fn is_ncname_start_char(self) -> bool {
        match self {
            'A'...'Z'                   |
            '_'                         |
            'a'...'z'                   |
            '\u{0000C0}'...'\u{0000D6}' |
            '\u{0000D8}'...'\u{0000F6}' |
            '\u{0000F8}'...'\u{0002FF}' |
            '\u{000370}'...'\u{00037D}' |
            '\u{00037F}'...'\u{001FFF}' |
            '\u{00200C}'...'\u{00200D}' |
            '\u{002070}'...'\u{00218F}' |
            '\u{002C00}'...'\u{002FEF}' |
            '\u{003001}'...'\u{00D7FF}' |
            '\u{00F900}'...'\u{00FDCF}' |
            '\u{00FDF0}'...'\u{00FFFD}' |
            '\u{010000}'...'\u{0EFFFF}' => true,
            _ => false,
        }
    }

    fn is_ncname_char(self) -> bool {
        if self.is_ncname_start_char() { return true; }
        match self {
            '-'                     |
            '.'                     |
            '0'...'9'               |
            '\u{00B7}'              |
            '\u{0300}'...'\u{036F}' |
            '\u{203F}'...'\u{2040}' => true,
            _ => false
        }
    }

    fn is_space_char(self) -> bool {
        match self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }

    fn is_decimal_char(self) -> bool {
        match self {
            '0'...'9' => true,
            _ => false,
        }
    }

    fn is_hex_char(self) -> bool {
        match self {
            '0'...'9' |
            'a'...'f' |
            'A'...'F' => true,
            _ => false,
        }
    }

    fn is_encoding_start_char(self) -> bool {
        match self {
            'A'...'Z' |
            'a'...'z' => true,
            _ => false,
        }
    }

    fn is_encoding_rest_char(self) -> bool {
        match self {
            'A'...'Z' |
            'a'...'z' |
            '0'...'9' |
            '.' |
            '_' |
            '-' => true,
            _ => false,
        }
    }

}

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

    #[test]
    fn end_of_char_data_leading_ampersand() {
        assert_eq!("&".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_leading_less_than() {
        assert_eq!("<".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_leading_cdata_end() {
        assert_eq!("]]>".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_until_ampersand() {
        assert_eq!("hello&world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_until_less_than() {
        assert_eq!("hello<world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_until_cdata_end() {
        assert_eq!("hello]]>world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_includes_right_square() {
        assert_eq!("hello]world".end_of_char_data(), Some("hello]world".len()));
    }

    #[test]
    fn end_of_char_data_includes_multiple_right_squares() {
        assert_eq!("hello]]world".end_of_char_data(), Some("hello]]world".len()));
    }

    #[test]
    fn end_of_int_subset_excludes_right_square() {
        assert_eq!("hello]>world".end_of_int_subset(), Some("hello".len()))
    }
}
<commit_msg>Replace deprecated inclusive range syntax<commit_after>trait StrParseExt {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool;
}

impl<'a> StrParseExt for &'a str {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool,
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len()),
        }
    }
}

pub trait XmlStr {
    /// Find the end of the quoted attribute value, not including the quote
    fn end_of_attribute(&self, quote: &str) -> Option<usize>;
    /// Find the end of the direct character data
    fn end_of_char_data(&self) -> Option<usize>;
    /// Find the end of the CData section, not including the ]]>
    fn end_of_cdata(&self) -> Option<usize>;
    /// Find the end of a run of decimal characters
    fn end_of_decimal_chars(&self) -> Option<usize>;
    /// Find the end of a run of hexidecimal characters
    fn end_of_hex_chars(&self) -> Option<usize>;
    /// Find the end of the comment, not including the -->
    fn end_of_comment(&self) -> Option<usize>;
    /// Find the end of the processing instruction, not including the ?>
    fn end_of_pi_value(&self) -> Option<usize>;
    /// Find the end of the [Name](http://www.w3.org/TR/xml/#NT-Name)
    fn end_of_name(&self) -> Option<usize>;
    /// Find the end of the [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)
    fn end_of_ncname(&self) -> Option<usize>;
    /// Find the end of a run of space characters
    fn end_of_space(&self) -> Option<usize>;
    /// Find the end of the starting tag
    fn end_of_start_tag(&self) -> Option<usize>;
    fn end_of_encoding(&self) -> Option<usize>;
    /// Find the end of the internal doc type declaration, not including the ]
    fn end_of_int_subset(&self) -> Option<usize>;
}

impl<'a> XmlStr for &'a str {
    fn end_of_attribute(&self, quote: &str) -> Option<usize> {
        if self.len() == 0 ||
           self.starts_with('&') ||
           self.starts_with('<') ||
           self.starts_with(quote)
        {
            return None;
        }

        let quote_char = quote.chars().next().expect("Cant have null quote");

        self.find(&['&', '<', quote_char][..]).or(Some(self.len()))
    }

    fn end_of_char_data(&self) -> Option<usize> {
        fn find_end_of_char_data(bytes: &[u8]) -> Option<usize> {
            for (i, &b) in bytes.iter().enumerate() {
                if b == b'<' || b == b'&' { return Some(i) }

                if b == b']' && bytes[i..].starts_with(b"]]>") {
                    return Some(i)
                }
            }
            None
        }

        match find_end_of_char_data(self.as_bytes()) {
            Some(0) => None,
            Some(v) => Some(v),
            None => Some(self.len()),
        }
    }

    fn end_of_cdata(&self) -> Option<usize> {
        self.find("]]>")
    }

    fn end_of_decimal_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_decimal_char(),
                               |c| c.is_decimal_char())
    }

    fn end_of_hex_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_hex_char(),
                               |c| c.is_hex_char())
    }

    fn end_of_comment(&self) -> Option<usize> {
        // This deliberately does not include the >. -- is not allowed
        // in a comment, so we can just test the end if it matches the
        // complete close delimiter.
        self.find("--")
    }

    fn end_of_pi_value(&self) -> Option<usize> {
        self.find("?>")
    }

    fn end_of_name(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn end_of_ncname(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_ncname_start_char(), |c| c.is_ncname_char())
    }

    fn end_of_space(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }

    fn end_of_start_tag(&self) -> Option<usize> {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if '<' == c => (),
            _ => return None,
        };

        match positions.next() {
            Some((offset, c)) =>
                match c {
                    '?' | '!' | '/' => None,
                    _ => Some(offset),
                },
            None => Some(self.len()),
        }
    }

    fn end_of_encoding(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_encoding_start_char(), |c| c.is_encoding_rest_char())
    }

    fn end_of_int_subset(&self) -> Option<usize> { self.find("]") }
}

/// Predicates used when parsing an characters in an XML document.
pub trait XmlChar {
    /// Is this a [NameStartChar](http://www.w3.org/TR/xml/#NT-NameStartChar)?
    fn is_name_start_char(self) -> bool;
    /// Is this a [NameChar](http://www.w3.org/TR/xml/#NT-NameChar)?
    fn is_name_char(self) -> bool;
    /// Does this start a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_start_char(self) -> bool;
    /// Is this a component of a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_char(self) -> bool;
    /// Is this an [XML space](http://www.w3.org/TR/xml/#NT-S)?
    fn is_space_char(self) -> bool;
    fn is_decimal_char(self) -> bool;
    fn is_hex_char(self) -> bool;
    fn is_encoding_start_char(self) -> bool;
    fn is_encoding_rest_char(self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(self) -> bool {
        self == ':' || self.is_ncname_start_char()
    }

    fn is_name_char(self) -> bool {
        self.is_name_start_char() || self.is_ncname_char()
    }

    fn is_ncname_start_char(self) -> bool {
        match self {
            'A'..='Z'                   |
            '_'                         |
            'a'..='z'                   |
            '\u{0000C0}'..='\u{0000D6}' |
            '\u{0000D8}'..='\u{0000F6}' |
            '\u{0000F8}'..='\u{0002FF}' |
            '\u{000370}'..='\u{00037D}' |
            '\u{00037F}'..='\u{001FFF}' |
            '\u{00200C}'..='\u{00200D}' |
            '\u{002070}'..='\u{00218F}' |
            '\u{002C00}'..='\u{002FEF}' |
            '\u{003001}'..='\u{00D7FF}' |
            '\u{00F900}'..='\u{00FDCF}' |
            '\u{00FDF0}'..='\u{00FFFD}' |
            '\u{010000}'..='\u{0EFFFF}' => true,
            _ => false,
        }
    }

    fn is_ncname_char(self) -> bool {
        if self.is_ncname_start_char() { return true; }
        match self {
            '-'                     |
            '.'                     |
            '0'..='9'               |
            '\u{00B7}'              |
            '\u{0300}'..='\u{036F}' |
            '\u{203F}'..='\u{2040}' => true,
            _ => false
        }
    }

    fn is_space_char(self) -> bool {
        match self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }

    fn is_decimal_char(self) -> bool {
        match self {
            '0'..='9' => true,
            _ => false,
        }
    }

    fn is_hex_char(self) -> bool {
        match self {
            '0'..='9' |
            'a'..='f' |
            'A'..='F' => true,
            _ => false,
        }
    }

    fn is_encoding_start_char(self) -> bool {
        match self {
            'A'..='Z' |
            'a'..='z' => true,
            _ => false,
        }
    }

    fn is_encoding_rest_char(self) -> bool {
        match self {
            'A'..='Z' |
            'a'..='z' |
            '0'..='9' |
            '.' |
            '_' |
            '-' => true,
            _ => false,
        }
    }

}

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

    #[test]
    fn end_of_char_data_leading_ampersand() {
        assert_eq!("&".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_leading_less_than() {
        assert_eq!("<".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_leading_cdata_end() {
        assert_eq!("]]>".end_of_char_data(), None);
    }

    #[test]
    fn end_of_char_data_until_ampersand() {
        assert_eq!("hello&world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_until_less_than() {
        assert_eq!("hello<world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_until_cdata_end() {
        assert_eq!("hello]]>world".end_of_char_data(), Some("hello".len()));
    }

    #[test]
    fn end_of_char_data_includes_right_square() {
        assert_eq!("hello]world".end_of_char_data(), Some("hello]world".len()));
    }

    #[test]
    fn end_of_char_data_includes_multiple_right_squares() {
        assert_eq!("hello]]world".end_of_char_data(), Some("hello]]world".len()));
    }

    #[test]
    fn end_of_int_subset_excludes_right_square() {
        assert_eq!("hello]>world".end_of_int_subset(), Some("hello".len()))
    }
}
<|endoftext|>