<commit_before>use specs::*;
use templating::*;

use handlebars::{self, Handlebars, HelperDef};
use itertools::Itertools;
use std::collections::BTreeMap;

pub struct CodeGenerator {
    specs:      AMQProtocolDefinition,
    handlebars: Handlebars,
}

impl CodeGenerator {
    pub fn new(specs: AMQProtocolDefinition, templates: AMQPTemplates) -> CodeGenerator {
        let mut handlebars = Handlebars::new();

        handlebars.register_escape_fn(handlebars::no_escape);

        handlebars.register_helper("camel", Box::new(camel_helper));
        handlebars.register_helper("snake", Box::new(snake_helper));

        handlebars.register_template_string("main",     &templates.main).expect("Failed to register main template");
        handlebars.register_template_string("domain",   &templates.domain).expect("Failed to register domain template");
        handlebars.register_template_string("constant", &templates.constant).expect("Failed to register constant template");
        handlebars.register_template_string("class",    &templates.klass).expect("Failed to register class template");
        handlebars.register_template_string("method",   &templates.method).expect("Failed to register method template");
        handlebars.register_template_string("argument", &templates.argument).expect("Failed to register argument template");
        handlebars.register_template_string("property", &templates.property).expect("Failed to register property template");

        CodeGenerator {
            specs:      specs,
            handlebars: handlebars,
        }
    }

    pub fn with_helper(mut self, name: &str, helper: Box<HelperDef + 'static>) -> CodeGenerator {
        self.handlebars.register_helper(name, helper);
        self
    }

    pub fn with_template(mut self, name: &str, template: &str) -> CodeGenerator {
        self.handlebars.register_template_string(name, template).expect(&format!("Failed to register {} template", name));
        self
    }

    pub fn generate(&self) -> String {
        self.generate_from_template_name("main")
    }

    pub fn generate_from_template_name(&self, name: &str) -> String {
        self.handlebars.render(name, &AMQProtocolDefinitionWrapper {
            protocol:  &self.specs,
            copyright: self.specs.copyright.iter().join(""),
            domains:   self.specs.domains.iter().map(|domain| domain.codegen(&self.handlebars)).join("\n"),
            constants: self.specs.constants.iter().map(|constant| constant.codegen(&self.handlebars)).join("\n"),
            classes:   self.specs.classes.iter().map(|klass| klass.codegen(&self.handlebars)).join("\n"),
        }).expect(&format!("Failed to render {} template", name))
    }
}

#[derive(Debug, Serialize)]
struct AMQProtocolDefinitionWrapper<'a> {
    protocol : &'a AMQProtocolDefinition,
    copyright: String,
    domains:   String,
    constants: String,
    classes:   String,
}

pub trait Codegen {
    fn codegen(&self, handlebars: &Handlebars) -> String;
}

trait BasicGetData {
    fn get_data(&self, handlebars: &Handlebars) -> BTreeMap<String, &Self>;
}

trait ClassGetData {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPClassWrapper;
}

trait MethodGetData {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPMethodWrapper;
}

macro_rules! get_data {
    ($t:ty, $key:expr) => {
        impl BasicGetData for $t {
            #[allow(unused_variables)]
            fn get_data(&self, handlebars: &Handlebars) -> BTreeMap<String, &Self> {
                let mut data = BTreeMap::new();
                data.insert($key.to_string(), self);
                data
            }
        }
    }
}

#[derive(Debug, Serialize)]
struct AMQPClassWrapper<'a> {
    #[serde(rename="class")]
    klass:      &'a AMQPClass,
    methods:    String,
    properties: Option<String>,
}

impl ClassGetData for AMQPClass {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPClassWrapper {
        AMQPClassWrapper {
            klass:      self,
            methods:    self.methods.iter().map(|method| method.codegen(&handlebars)).join("\n"),
            properties: match self.properties {
                Some(ref properties) => Some(properties.iter().map(|method| method.codegen(&handlebars)).join("\n")),
                None                 => None,
            },
        }
    }
}

#[derive(Debug, Serialize)]
struct AMQPMethodWrapper<'a> {
    method:    &'a AMQPMethod,
    arguments: String,
}

impl MethodGetData for AMQPMethod {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPMethodWrapper {
        AMQPMethodWrapper {
            method:    self,
            arguments: self.arguments.iter().map(|arg| arg.codegen(&handlebars)).join("\n"),
        }
    }
}

get_data!(AMQPDomain,   "domain");
get_data!(AMQPConstant, "constant");
get_data!(AMQPArgument, "argument");
get_data!(AMQPProperty, "property");

macro_rules! codegen {
    ($t:ty, $name:expr) => {
        impl Codegen for $t {
            fn codegen (&self, handlebars: &Handlebars) -> String {
                handlebars.render($name, &self.get_data(handlebars)).expect(&format!("Failed to render {} template", $name))
            }
        }
    };
}

codegen!(AMQPDomain,   "domain");
codegen!(AMQPConstant, "constant");
codegen!(AMQPClass,    "class");
codegen!(AMQPMethod,   "method");
codegen!(AMQPArgument, "argument");
codegen!(AMQPProperty, "property");

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

    use amq_protocol_types::*;
    use serde_json::Value;
    use templating::AMQPTemplates;

    fn specs() -> AMQProtocolDefinition {
        AMQProtocolDefinition {
            name:          "AMQP".to_string(),
            major_version: 0,
            minor_version: 9,
            revision:      1,
            port:          5672,
            copyright:     vec!["Copyright 1\n".to_string(), "Copyright 2".to_string()],
            domains:       vec![
                AMQPDomain{
                    name:      "domain1".to_string(),
                    amqp_type: AMQPType::ShortInt
                }
            ],
            constants:     vec![
                AMQPConstant {
                    name:  "constant1".to_string(),
                    value: 128,
                    klass: Some("class1".to_string()),
                }
            ],
            classes:       vec![
                AMQPClass {
                    id:         42,
                    methods:    vec![
                        AMQPMethod {
                            id:          64,
                            arguments:   vec![
                                AMQPArgument {
                                    amqp_type:     Some(AMQPType::LongString),
                                    name:          "argument1".to_string(),
                                    default_value: Some(Value::String("value1".to_string())),
                                    domain:        Some("domain1".to_string()),
                                }
                            ],
                            name:        "method1".to_string(),
                            synchronous: Some(true),
                        }
                    ],
                    name:       "class1".to_string(),
                    properties: Some(vec![
                        AMQPProperty {
                            amqp_type: AMQPType::LongString,
                            name:      "property1".to_string(),
                        }
                    ]),
                }
            ],
        }
    }

    fn templates() -> AMQPTemplates {
        AMQPTemplates {
            main:     r#"
{{protocol.name}} - {{protocol.major_version}}.{{protocol.minor_version}}.{{protocol.revision}}
{{copyright}}
port {{protocol.port}}
{{domains}}
{{constants}}
{{classes}}
"#.to_string(),
            domain:   "{{domain.name}}: {{domain.type}}".to_string(),
            constant: "{{constant.name}}({{constant.class}}) = {{constant.value}}".to_string(),
            klass:    r#"
{{class.id}} - {{class.name}}
{{properties}}
{{methods}}
"#.to_string(),
            method:   r#"
{{method.id}} - {{method.name}}
synchronous: {{method.synchronous}}
{{arguments}}
"#.to_string(),
            argument: "{{argument.name}}({{argument.domain}}): {{argument.type}}".to_string(),
            property: "{{property.name}}: {{property.type}}".to_string(),
        }
    }

    #[test]
    fn main_template() {
        assert_eq!(specs().code_generator(templates()).generate(), r#"
AMQP - 0.9.1
Copyright 1
Copyright 2
port 5672
domain1: ShortInt
constant1(class1) = 128

42 - class1
property1: LongString

64 - method1
synchronous: true
argument1(domain1): LongString

"#);
    }
}
<commit_msg>codegen: add a way to get the underlying handlebars<commit_after>use specs::*;
use templating::*;

use handlebars::{self, Handlebars, HelperDef};
use itertools::Itertools;
use std::collections::BTreeMap;

pub struct CodeGenerator {
    specs:      AMQProtocolDefinition,
    handlebars: Handlebars,
}

impl CodeGenerator {
    pub fn new(specs: AMQProtocolDefinition, templates: AMQPTemplates) -> CodeGenerator {
        let mut handlebars = Handlebars::new();

        handlebars.register_escape_fn(handlebars::no_escape);

        handlebars.register_helper("camel", Box::new(camel_helper));
        handlebars.register_helper("snake", Box::new(snake_helper));

        handlebars.register_template_string("main",     &templates.main).expect("Failed to register main template");
        handlebars.register_template_string("domain",   &templates.domain).expect("Failed to register domain template");
        handlebars.register_template_string("constant", &templates.constant).expect("Failed to register constant template");
        handlebars.register_template_string("class",    &templates.klass).expect("Failed to register class template");
        handlebars.register_template_string("method",   &templates.method).expect("Failed to register method template");
        handlebars.register_template_string("argument", &templates.argument).expect("Failed to register argument template");
        handlebars.register_template_string("property", &templates.property).expect("Failed to register property template");

        CodeGenerator {
            specs:      specs,
            handlebars: handlebars,
        }
    }

    pub fn with_helper(mut self, name: &str, helper: Box<HelperDef + 'static>) -> CodeGenerator {
        self.handlebars.register_helper(name, helper);
        self
    }

    pub fn with_template(mut self, name: &str, template: &str) -> CodeGenerator {
        self.handlebars.register_template_string(name, template).expect(&format!("Failed to register {} template", name));
        self
    }

    pub fn generate(&self) -> String {
        self.generate_from_template_name("main")
    }

    pub fn generate_from_template_name(&self, name: &str) -> String {
        self.handlebars.render(name, &AMQProtocolDefinitionWrapper {
            protocol:  &self.specs,
            copyright: self.specs.copyright.iter().join(""),
            domains:   self.specs.domains.iter().map(|domain| domain.codegen(&self.handlebars)).join("\n"),
            constants: self.specs.constants.iter().map(|constant| constant.codegen(&self.handlebars)).join("\n"),
            classes:   self.specs.classes.iter().map(|klass| klass.codegen(&self.handlebars)).join("\n"),
        }).expect(&format!("Failed to render {} template", name))
    }

    pub fn get_handlebars(&mut self) -> &mut Handlebars {
        &mut self.handlebars
    }
}

#[derive(Debug, Serialize)]
struct AMQProtocolDefinitionWrapper<'a> {
    protocol : &'a AMQProtocolDefinition,
    copyright: String,
    domains:   String,
    constants: String,
    classes:   String,
}

pub trait Codegen {
    fn codegen(&self, handlebars: &Handlebars) -> String;
}

trait BasicGetData {
    fn get_data(&self, handlebars: &Handlebars) -> BTreeMap<String, &Self>;
}

trait ClassGetData {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPClassWrapper;
}

trait MethodGetData {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPMethodWrapper;
}

macro_rules! get_data {
    ($t:ty, $key:expr) => {
        impl BasicGetData for $t {
            #[allow(unused_variables)]
            fn get_data(&self, handlebars: &Handlebars) -> BTreeMap<String, &Self> {
                let mut data = BTreeMap::new();
                data.insert($key.to_string(), self);
                data
            }
        }
    }
}

#[derive(Debug, Serialize)]
struct AMQPClassWrapper<'a> {
    #[serde(rename="class")]
    klass:      &'a AMQPClass,
    methods:    String,
    properties: Option<String>,
}

impl ClassGetData for AMQPClass {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPClassWrapper {
        AMQPClassWrapper {
            klass:      self,
            methods:    self.methods.iter().map(|method| method.codegen(&handlebars)).join("\n"),
            properties: match self.properties {
                Some(ref properties) => Some(properties.iter().map(|method| method.codegen(&handlebars)).join("\n")),
                None                 => None,
            },
        }
    }
}

#[derive(Debug, Serialize)]
struct AMQPMethodWrapper<'a> {
    method:    &'a AMQPMethod,
    arguments: String,
}

impl MethodGetData for AMQPMethod {
    fn get_data(&self, handlebars: &Handlebars) -> AMQPMethodWrapper {
        AMQPMethodWrapper {
            method:    self,
            arguments: self.arguments.iter().map(|arg| arg.codegen(&handlebars)).join("\n"),
        }
    }
}

get_data!(AMQPDomain,   "domain");
get_data!(AMQPConstant, "constant");
get_data!(AMQPArgument, "argument");
get_data!(AMQPProperty, "property");

macro_rules! codegen {
    ($t:ty, $name:expr) => {
        impl Codegen for $t {
            fn codegen (&self, handlebars: &Handlebars) -> String {
                handlebars.render($name, &self.get_data(handlebars)).expect(&format!("Failed to render {} template", $name))
            }
        }
    };
}

codegen!(AMQPDomain,   "domain");
codegen!(AMQPConstant, "constant");
codegen!(AMQPClass,    "class");
codegen!(AMQPMethod,   "method");
codegen!(AMQPArgument, "argument");
codegen!(AMQPProperty, "property");

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

    use amq_protocol_types::*;
    use serde_json::Value;
    use templating::AMQPTemplates;

    fn specs() -> AMQProtocolDefinition {
        AMQProtocolDefinition {
            name:          "AMQP".to_string(),
            major_version: 0,
            minor_version: 9,
            revision:      1,
            port:          5672,
            copyright:     vec!["Copyright 1\n".to_string(), "Copyright 2".to_string()],
            domains:       vec![
                AMQPDomain{
                    name:      "domain1".to_string(),
                    amqp_type: AMQPType::ShortInt
                }
            ],
            constants:     vec![
                AMQPConstant {
                    name:  "constant1".to_string(),
                    value: 128,
                    klass: Some("class1".to_string()),
                }
            ],
            classes:       vec![
                AMQPClass {
                    id:         42,
                    methods:    vec![
                        AMQPMethod {
                            id:          64,
                            arguments:   vec![
                                AMQPArgument {
                                    amqp_type:     Some(AMQPType::LongString),
                                    name:          "argument1".to_string(),
                                    default_value: Some(Value::String("value1".to_string())),
                                    domain:        Some("domain1".to_string()),
                                }
                            ],
                            name:        "method1".to_string(),
                            synchronous: Some(true),
                        }
                    ],
                    name:       "class1".to_string(),
                    properties: Some(vec![
                        AMQPProperty {
                            amqp_type: AMQPType::LongString,
                            name:      "property1".to_string(),
                        }
                    ]),
                }
            ],
        }
    }

    fn templates() -> AMQPTemplates {
        AMQPTemplates {
            main:     r#"
{{protocol.name}} - {{protocol.major_version}}.{{protocol.minor_version}}.{{protocol.revision}}
{{copyright}}
port {{protocol.port}}
{{domains}}
{{constants}}
{{classes}}
"#.to_string(),
            domain:   "{{domain.name}}: {{domain.type}}".to_string(),
            constant: "{{constant.name}}({{constant.class}}) = {{constant.value}}".to_string(),
            klass:    r#"
{{class.id}} - {{class.name}}
{{properties}}
{{methods}}
"#.to_string(),
            method:   r#"
{{method.id}} - {{method.name}}
synchronous: {{method.synchronous}}
{{arguments}}
"#.to_string(),
            argument: "{{argument.name}}({{argument.domain}}): {{argument.type}}".to_string(),
            property: "{{property.name}}: {{property.type}}".to_string(),
        }
    }

    #[test]
    fn main_template() {
        assert_eq!(specs().code_generator(templates()).generate(), r#"
AMQP - 0.9.1
Copyright 1
Copyright 2
port 5672
domain1: ShortInt
constant1(class1) = 128

42 - class1
property1: LongString

64 - method1
synchronous: true
argument1(domain1): LongString

"#);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate bencher;
extern crate whatlang;
extern crate rustc_serialize;

use bencher::Bencher;
use rustc_serialize::json;
use std::collections::HashMap;
use whatlang::{detect, detect_script};

fn bench_detect(bench: &mut Bencher) {
    let example_data = include_str!("../tests/examples.json");
    let examples: HashMap<String, String> = json::decode(example_data).unwrap();

    bench.iter(|| {
        for (_, text) in &examples {
            detect(&text);
        }
    })
}

fn bench_detect_script(bench: &mut Bencher) {
    let example_data = include_str!("../tests/examples.json");
    let examples: HashMap<String, String> = json::decode(example_data).unwrap();

    bench.iter(|| {
        for (_, text) in &examples {
            detect_script(&text);
        }
    })
}

benchmark_group!(benches, bench_detect, bench_detect_script);
benchmark_main!(benches);
<commit_msg>Fix benches<commit_after>#[macro_use]
extern crate bencher;
extern crate whatlang;
extern crate serde_json;

use bencher::Bencher;
use std::collections::HashMap;
use whatlang::{detect, detect_script};

fn bench_detect(bench: &mut Bencher) {
    let example_data = include_str!("../tests/examples.json");
    let examples: HashMap<String, String> = serde_json::from_str(example_data).unwrap();

    bench.iter(|| {
        for (_, text) in &examples {
            detect(&text);
        }
    })
}

fn bench_detect_script(bench: &mut Bencher) {
    let example_data = include_str!("../tests/examples.json");
    let examples: HashMap<String, String> = serde_json::from_str(example_data).unwrap();

    bench.iter(|| {
        for (_, text) in &examples {
            detect_script(&text);
        }
    })
}

benchmark_group!(benches, bench_detect, bench_detect_script);
benchmark_main!(benches);
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate futures;
#[macro_use]
extern crate tokio_core;

use std::io;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::net::SocketAddr;

use futures::{Future, Poll};
use tokio_core::net::UdpSocket;
use tokio_core::reactor::Core;

use test::Bencher;
use std::thread;
use std::time::Duration;

use futures::stream::Stream;

/// UDP echo server
struct EchoServer {
    socket : UdpSocket,
    buf : Vec<u8>,
    to_send : Option<(usize, SocketAddr)>,
    stop : Arc<AtomicBool>,
}

impl EchoServer {
    fn new(s : UdpSocket, stop : Arc<AtomicBool>) -> Self {
        EchoServer {
            socket: s,
            to_send: None,
            buf: vec![0u8; 1600],
            stop: stop,
        }
    }
}

impl Future for EchoServer {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        loop {
            if self.stop.load(Ordering::SeqCst) {
                return Ok(futures::Async::Ready(()))
            }

            if let Some((size, peer)) = self.to_send.take() {
                match self.socket.send_to(&self.buf[..size], &peer) {
                    Err(e) => {
                        self.to_send = Some((size, peer));
                        return try_nb!(Err(e));
                    },
                    Ok(_) => { }
                }
            }

            self.to_send = Some(
                try_nb!(self.socket.recv_from(&mut self.buf))
                );
        }
    }
}

/// UDP echo server
struct ChanServer {
    rx : tokio_core::channel::Receiver<usize>,
    tx : tokio_core::channel::Sender<usize>,
    buf : Option<usize>,
    stop : Arc<AtomicBool>,
}

impl ChanServer {
    fn new(tx : tokio_core::channel::Sender<usize>, rx : tokio_core::channel::Receiver<usize>, stop : Arc<AtomicBool>) -> Self {
        ChanServer {
            rx: rx,
            tx: tx,
            buf: None,
            stop: stop,
        }
    }
}

impl Future for ChanServer {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        loop {
            if self.stop.load(Ordering::SeqCst) {
                return Ok(futures::Async::Ready(()))
            }

            if let Some(u) = self.buf.take() {
                match self.tx.send(u) {
                    Err(e) => {
                        self.buf = Some(u);
                        return try_nb!(Err(e));
                    },
                    Ok(_) => { }
                }
            }

            match self.rx.poll() {
                Ok(futures::Async::Ready(None)) => return Ok(futures::Async::Ready(())),
                Ok(futures::Async::Ready(Some(t))) => { self.buf = Some(t) },
                Ok(futures::Async::NotReady) => return Ok(futures::Async::NotReady),
                Err(e) => return try_nb!(Err(e)),
            }
        }
    }
}
#[bench]
fn udp_echo_latency(b: &mut Bencher) {
    let server_addr= "127.0.0.1:7398".to_string();
    let server_addr = server_addr.parse::<SocketAddr>().unwrap();
    let client_addr= "127.0.0.1:7399".to_string();
    let client_addr = client_addr.parse::<SocketAddr>().unwrap();

    let stop = Arc::new(AtomicBool::new(false));
    let stop2 = stop.clone();

    let child = thread::spawn(move || {
        let mut l = Core::new().unwrap();
        let handle = l.handle();

        let socket = tokio_core::net::UdpSocket::bind(&server_addr, &handle).unwrap();

        let server = EchoServer::new(socket, stop);

        l.run(server).unwrap();
    });

    // TODO: More reliable way to bind server socket and start server
    // first
    thread::sleep(Duration::from_millis(100));

    let client = std::net::UdpSocket::bind(client_addr).unwrap();

    let mut buf = [0u8; 1000];

    // warmup phase; for some reason initial couple of
    // rounds is much slower
    //
    // TODO: Describe the exact reasons; caching? branch predictor? lazy closures?
    for _ in 0..1000 {
        client.send_to(&buf, &server_addr).unwrap();
        let _ = client.recv_from(&mut buf).unwrap();
    }

    b.iter(|| {
        client.send_to(&buf, &server_addr).unwrap();
        let _ = client.recv_from(&mut buf).unwrap();
    });

    // Stop the server; TODO: Use better method
    stop2.store(true, Ordering::SeqCst);
    thread::sleep(Duration::from_millis(1));
    client.send_to(&buf, &server_addr).unwrap();

    child.join().unwrap();
}

#[bench]
fn channel_latency(b: &mut Bencher) {
    let stop = Arc::new(AtomicBool::new(false));
    let stop2 = stop.clone();

    // TODO: Any way to start Loop on a separate thread and yet get
    // a tokio_core::channel to it?
    let (tx, rx) = std::sync::mpsc::channel();

    let child = thread::spawn(move || {
        let mut l = Core::new().unwrap();
        let handle = l.handle();

        let (in_tx, in_rx) = tokio_core::channel::channel(&handle).unwrap();
        let (out_tx, out_rx) = tokio_core::channel::channel(&handle).unwrap();

        let server = ChanServer::new(out_tx, in_rx, stop);

        tx.send((in_tx, out_rx)).unwrap();
        l.run(server).unwrap();
    });

    let (in_tx, out_rx) = rx.recv().unwrap();

    // TODO: More reliable way to bind server socket and start server
    // first
    thread::sleep(Duration::from_millis(100));

    let mut rx_iter = out_rx.wait();

    // warmup phase; for some reason initial couple of
    // rounds is much slower
    //
    // TODO: Describe the exact reasons; caching? branch predictor? lazy closures?
    for _ in 0..1000 {
        in_tx.send(1usize).unwrap();
        let _ = rx_iter.next();
    }

    b.iter(|| {
        in_tx.send(1usize).unwrap();
        let _ = rx_iter.next();
    });

    // Stop the server; TODO: Use better method
    stop2.store(true, Ordering::SeqCst);
    thread::sleep(Duration::from_millis(1));
    in_tx.send(1usize).unwrap();

    child.join().unwrap();
}
<commit_msg>Improve latency benchmarks.<commit_after>#![feature(test)]

extern crate test;
extern crate futures;
#[macro_use]
extern crate tokio_core;

use std::io;
use std::net::SocketAddr;

use futures::{Future, Poll};
use tokio_core::net::UdpSocket;
use tokio_core::reactor::Core;

use test::Bencher;
use std::thread;

use futures::stream::Stream;
use futures::{oneshot, Oneshot};

/// UDP echo server
struct EchoServer {
    socket : UdpSocket,
    buf : Vec<u8>,
    to_send : Option<(usize, SocketAddr)>,
    stop : Oneshot<()>,
}

impl EchoServer {
    fn new(s : UdpSocket, stop : Oneshot<()>) -> Self {
        EchoServer {
            socket: s,
            to_send: None,
            buf: vec![0u8; 1600],
            stop: stop,
        }
    }
}

impl Future for EchoServer {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        loop {
            if self.stop.poll() != Ok(futures::Async::NotReady) {
                return Ok(futures::Async::Ready(()))
            }

            if let Some(&(size, peer)) = self.to_send.as_ref() {
                let _ = try_nb!(self.socket.send_to(&self.buf[..size], &peer));
            }
            self.to_send = None;

            self.to_send = Some(
                try_nb!(self.socket.recv_from(&mut self.buf))
                );
        }
    }
}

/// UDP echo server
///
/// TODO: This may be replace-able with the newly-minted Sink::send_all function in the futures crate
struct ChanServer {
    rx : tokio_core::channel::Receiver<usize>,
    tx : tokio_core::channel::Sender<usize>,
    buf : Option<usize>,
}

impl ChanServer {
    fn new(tx : tokio_core::channel::Sender<usize>, rx : tokio_core::channel::Receiver<usize>) -> Self {
        ChanServer {
            rx: rx,
            tx: tx,
            buf: None,
        }
    }
}

impl Future for ChanServer {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        loop {
            if let Some(u) = self.buf.take() {
                match self.tx.send(u) {
                    Err(e) => {
                        self.buf = Some(u);
                        return try_nb!(Err(e));
                    },
                    Ok(_) => { }
                }
            }

            match self.rx.poll() {
                Ok(futures::Async::Ready(None)) => return Ok(futures::Async::Ready(())),
                Ok(futures::Async::Ready(Some(t))) => { self.buf = Some(t) },
                Ok(futures::Async::NotReady) => return Ok(futures::Async::NotReady),
                Err(e) => return try_nb!(Err(e)),
            }
        }
    }
}

#[bench]
fn udp_echo_latency(b: &mut Bencher) {
    let server_addr= "127.0.0.1:0".to_string();
    let server_addr = server_addr.parse::<SocketAddr>().unwrap();
    let client_addr= "127.0.0.1:0".to_string();
    let client_addr = client_addr.parse::<SocketAddr>().unwrap();

    let (stop_c, stop_p) = oneshot::<()>();

    let (tx, rx) = std::sync::mpsc::channel();

    let child = thread::spawn(move || {
        let mut l = Core::new().unwrap();
        let handle = l.handle();

        let socket = tokio_core::net::UdpSocket::bind(&server_addr, &handle).unwrap();

        tx.send(socket.local_addr().unwrap()).unwrap();
        let server = EchoServer::new(socket, stop_p);

        l.run(server).unwrap();
    });

    let client = std::net::UdpSocket::bind(client_addr).unwrap();

    let server_addr = rx.recv().unwrap();
    let mut buf = [0u8; 1000];

    // warmup phase; for some reason initial couple of
    // runs are much slower
    //
    // TODO: Describe the exact reasons; caching? branch predictor? lazy closures?
    for _ in 0..8 {
        client.send_to(&buf, &server_addr).unwrap();
        let _ = client.recv_from(&mut buf).unwrap();
    }

    b.iter(|| {
        client.send_to(&buf, &server_addr).unwrap();
        let _ = client.recv_from(&mut buf).unwrap();
    });

    stop_c.complete(());

    child.join().unwrap();
}

#[bench]
fn channel_latency(b: &mut Bencher) {

    let (tx, rx) = std::sync::mpsc::channel();

    let child = thread::spawn(move || {
        let mut l = Core::new().unwrap();
        let handle = l.handle();

        let (in_tx, in_rx) = tokio_core::channel::channel(&handle).unwrap();
        let (out_tx, out_rx) = tokio_core::channel::channel(&handle).unwrap();

        let server = ChanServer::new(out_tx, in_rx);

        tx.send((in_tx, out_rx)).unwrap();
        l.run(server).unwrap();
    });

    let (in_tx, out_rx) = rx.recv().unwrap();

    let mut rx_iter = out_rx.wait();

    // warmup phase; for some reason initial couple of runs are much slower
    //
    // TODO: Describe the exact reasons; caching? branch predictor? lazy closures?
    for _ in 0..8 {
        in_tx.send(1usize).unwrap();
        let _ = rx_iter.next();
    }

    b.iter(|| {
        in_tx.send(1usize).unwrap();
        let _ = rx_iter.next();
    });

    drop(in_tx);

    child.join().unwrap();
}
<|endoftext|>


<commit_before>#[derive(Debug, Clone, PartialEq)]
pub enum SimilarityModel {
    TfIdf,
    Bm25{k1: f64, b: f64},
}

/// tf(term_frequency) = log(term_frequency) + 1.0
#[inline]
fn tf(term_frequency: u32) -> f64 {
    (term_frequency as f64).ln() + 1.0
}

/// idf(term_docs, total_docs) = log((total_docs + 1.0) / (term_docs + 1.0)) + 1.0
#[inline]
fn idf(term_docs: u64, total_docs: u64) -> f64 {
    ((total_docs as f64 + 1.0) / (term_docs as f64 + 1.0)).ln() + 1.0
}

impl SimilarityModel {
    pub fn score(&self, term_frequency: u32, length: u32, total_tokens: u64, total_docs: u64, total_docs_with_term: u64) -> f64 {
        match *self {
            SimilarityModel::TfIdf => {
                let tf = tf(term_frequency);
                let idf = idf(total_docs_with_term, total_docs);

                tf * idf
            }
            SimilarityModel::Bm25{k1, b} => {
                let tf = tf(term_frequency);
                let idf = idf(total_docs_with_term, total_docs);
                let average_length = (total_tokens as f64) / (total_docs as f64);

                idf * (k1 + 1.0) * (tf / (tf + (k1 * ((1.0 - b) + b * (length as f64).sqrt() / average_length.sqrt()))))
            }
        }
    }
}

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

    #[test]
    fn test_tf_idf_higher_term_freq_increases_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(2, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 5));
    }

    #[test]
    fn test_tf_idf_lower_term_docs_increases_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 10));
    }

    #[test]
    fn test_tf_idf_field_length_doesnt_affect_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 100, 100, 20, 5) == similarity.score(1, 40, 100, 20, 5));
    }

    #[test]
    fn test_tf_idf_total_tokens_doesnt_affect_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 40, 1000, 20, 5) == similarity.score(1, 40, 100, 20, 5));
    }

    #[test]
    fn test_bm25_higher_term_freq_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(2, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 5));
    }

    #[test]
    fn test_bm25_lower_term_docs_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 10));
    }

    #[test]
    fn test_bm25_lower_field_length_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 100, 20, 5) > similarity.score(1, 100, 100, 20, 5));
    }

    #[test]
    fn test_bm25_higher_total_tokens_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 1000, 20, 5) > similarity.score(1, 40, 100, 20, 5));
    }
}
<commit_msg>Prevent similarity algorithms from producing NaNs when one of the inputs are 0<commit_after>#[derive(Debug, Clone, PartialEq)]
pub enum SimilarityModel {
    TfIdf,
    Bm25{k1: f64, b: f64},
}

/// tf(term_frequency) = log(term_frequency + 1.0) + 1.0
#[inline]
fn tf(term_frequency: u32) -> f64 {
    (term_frequency as f64 + 1.0f64).ln() + 1.0
}

/// idf(term_docs, total_docs) = log((total_docs + 1.0) / (term_docs + 1.0)) + 1.0
#[inline]
fn idf(term_docs: u64, total_docs: u64) -> f64 {
    ((total_docs as f64 + 1.0) / (term_docs as f64 + 1.0)).ln() + 1.0
}

impl SimilarityModel {
    pub fn score(&self, term_frequency: u32, length: u32, total_tokens: u64, total_docs: u64, total_docs_with_term: u64) -> f64 {
        match *self {
            SimilarityModel::TfIdf => {
                let tf = tf(term_frequency);
                let idf = idf(total_docs_with_term, total_docs);

                tf * idf
            }
            SimilarityModel::Bm25{k1, b} => {
                let tf = tf(term_frequency);
                let idf = idf(total_docs_with_term, total_docs);
                let average_length = (total_tokens as f64 + 1.0f64) / (total_docs as f64 + 1.0f64);

                idf * (k1 + 1.0) * (tf / (tf + (k1 * ((1.0 - b) + b * (length as f64).sqrt() / average_length.sqrt())) + 1.0f64))
            }
        }
    }
}

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

    #[test]
    fn test_tf_idf_higher_term_freq_increases_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(2, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 5));
    }

    #[test]
    fn test_tf_idf_lower_term_docs_increases_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 10));
    }

    #[test]
    fn test_tf_idf_field_length_doesnt_affect_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 100, 100, 20, 5) == similarity.score(1, 40, 100, 20, 5));
    }

    #[test]
    fn test_tf_idf_total_tokens_doesnt_affect_score() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(1, 40, 1000, 20, 5) == similarity.score(1, 40, 100, 20, 5));
    }

    #[test]
    fn test_tf_idf_handles_zeros() {
        let similarity = SimilarityModel::TfIdf;

        assert!(similarity.score(0, 0, 0, 0, 0).is_finite());
    }

    #[test]
    fn test_bm25_higher_term_freq_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(2, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 5));
    }

    #[test]
    fn test_bm25_lower_term_docs_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 100, 10, 5) > similarity.score(1, 40, 100, 10, 10));
    }

    #[test]
    fn test_bm25_lower_field_length_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 100, 20, 5) > similarity.score(1, 100, 100, 20, 5));
    }

    #[test]
    fn test_bm25_higher_total_tokens_increases_score() {
        let similarity = SimilarityModel::Bm25 {
            k1: 1.2,
            b: 0.75,
        };

        assert!(similarity.score(1, 40, 1000, 20, 5) > similarity.score(1, 40, 100, 20, 5));
    }

    #[test]
    fn test_bm25_handles_zeros() {
        let similarity = SimilarityModel::Bm25 {
            k1: 0.0,
            b: 0.0,
        };

        assert!(similarity.score(0, 0, 0, 0, 0).is_finite());
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::{env, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use toml;

/// A Crate Dependency
pub type Dependency = (String, toml::Value);

/// Enumeration of errors which can occur when working with a rust manifest.
quick_error! {
    #[derive(Debug)]
    pub enum ManifestError {
        /// Cargo.toml could not be found.
        MissingManifest {
            description("missing manifest")
            display("Your Cargo.toml is missing.")
        }
        /// The TOML table could not be found.
        NonExistentTable(table: String) {
            description("non existent table")
            display("The table `{}` could not be found.", table)
        }
        /// The dependency could not be found.
        NonExistentDependency(name: String, table: String) {
            description("non existent dependency")
            display("The dependency `{}` could not be found in `{}`.", name, table)
        }
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&str>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError::MissingManifest).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file(&self, file: &mut File) -> Result<(), Box<Error>> {
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError::MissingManifest));

        let new_contents = format!(
          "[{}]\n{}{}",
          proj_header,
          proj_data,
          toml::Value::Table(toml)
        );
        let new_contents_bytes = new_contents.as_bytes();

        // We need to truncate the file, otherwise the new contents
        // will be mixed up with the old ones.
        file.set_len(new_contents_bytes.len() as u64).unwrap();
        file.write_all(new_contents_bytes)
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn insert_into_table(&mut self,
                             table: &str,
                             &(ref name, ref data): &Dependency)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match *entry {
            toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError::NonExistentTable(table.into())),
        }
    }

    /// Remove entry from a Cargo.toml.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate cargo_edit;
    /// # extern crate toml;
    /// # fn main() {
    ///     use cargo_edit::Manifest;
    ///     use toml;
    ///
    ///     let mut manifest = Manifest { data: toml::Table::new() };
    ///     let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
    ///     let _ = manifest.insert_into_table("dependencies", &dep);
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    /// # }
    /// ```
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn remove_from_table(&mut self, table: &str, name: &str) -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table));

        match entry {
            Entry::Vacant(_) => Err(ManifestError::NonExistentTable(table.into())),
            Entry::Occupied(entry) => {
                match *entry.into_mut() {
                    toml::Value::Table(ref mut deps) => {
                        deps.remove(name)
                            .map(|_| ())
                            .ok_or(ManifestError::NonExistentDependency(name.into(), table.into()))
                    }
                    _ => Err(ManifestError::NonExistentTable(table.into())),
                }
            }
        }
    }

    /// Add multiple dependencies to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;
    use toml;

    #[test]
    fn add_remove_dependency() {
        let mut manifest = Manifest { data: toml::Table::new() };
        // Create a copy containing empty "dependencies" table because removing
        //   the last entry in a table does not remove the section.
        let mut copy = Manifest { data: toml::Table::new() };
        copy.data.insert("dependencies".to_owned(),
                         toml::Value::Table(BTreeMap::new()));
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
        assert_eq!(manifest, copy);
    }

    #[test]
    fn remove_dependency_no_section() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    }

    #[test]
    fn remove_dependency_non_existent() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        let other_dep = ("other-dep".to_owned(),
                         toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &other_dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());

    }
}
<commit_msg>wrap .set_len() in a try and remove unused imports<commit_after>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::{env, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use toml;

/// A Crate Dependency
pub type Dependency = (String, toml::Value);

/// Enumeration of errors which can occur when working with a rust manifest.
quick_error! {
    #[derive(Debug)]
    pub enum ManifestError {
        /// Cargo.toml could not be found.
        MissingManifest {
            description("missing manifest")
            display("Your Cargo.toml is missing.")
        }
        /// The TOML table could not be found.
        NonExistentTable(table: String) {
            description("non existent table")
            display("The table `{}` could not be found.", table)
        }
        /// The dependency could not be found.
        NonExistentDependency(name: String, table: String) {
            description("non existent dependency")
            display("The dependency `{}` could not be found in `{}`.", name, table)
        }
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&str>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError::MissingManifest).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file(&self, file: &mut File) -> Result<(), Box<Error>> {
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError::MissingManifest));

        let new_contents = format!(
          "[{}]\n{}{}",
          proj_header,
          proj_data,
          toml::Value::Table(toml)
        );
        let new_contents_bytes = new_contents.as_bytes();

        // We need to truncate the file, otherwise the new contents
        // will be mixed up with the old ones.
        try!(file.set_len(new_contents_bytes.len() as u64));
        file.write_all(new_contents_bytes)
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn insert_into_table(&mut self,
                             table: &str,
                             &(ref name, ref data): &Dependency)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match *entry {
            toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError::NonExistentTable(table.into())),
        }
    }

    /// Remove entry from a Cargo.toml.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate cargo_edit;
    /// # extern crate toml;
    /// # fn main() {
    ///     use cargo_edit::Manifest;
    ///     use toml;
    ///
    ///     let mut manifest = Manifest { data: toml::Table::new() };
    ///     let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
    ///     let _ = manifest.insert_into_table("dependencies", &dep);
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    /// # }
    /// ```
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn remove_from_table(&mut self, table: &str, name: &str) -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table));

        match entry {
            Entry::Vacant(_) => Err(ManifestError::NonExistentTable(table.into())),
            Entry::Occupied(entry) => {
                match *entry.into_mut() {
                    toml::Value::Table(ref mut deps) => {
                        deps.remove(name)
                            .map(|_| ())
                            .ok_or(ManifestError::NonExistentDependency(name.into(), table.into()))
                    }
                    _ => Err(ManifestError::NonExistentTable(table.into())),
                }
            }
        }
    }

    /// Add multiple dependencies to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;
    use toml;

    #[test]
    fn add_remove_dependency() {
        let mut manifest = Manifest { data: toml::Table::new() };
        // Create a copy containing empty "dependencies" table because removing
        //   the last entry in a table does not remove the section.
        let mut copy = Manifest { data: toml::Table::new() };
        copy.data.insert("dependencies".to_owned(),
                         toml::Value::Table(BTreeMap::new()));
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
        assert_eq!(manifest, copy);
    }

    #[test]
    fn remove_dependency_no_section() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    }

    #[test]
    fn remove_dependency_non_existent() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(),
                   toml::Value::String("0.1.0".to_owned()));
        let other_dep = ("other-dep".to_owned(),
                         toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &other_dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());

    }
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;
use std::cmp::min;

pub fn compute_matches<'a>(choices: &[&'a str], query: &str) -> Vec<&'a str> {
    struct ScoredChoice<'a> {
        score: f64,
        choice: &'a str,
    };
    let mut ret = Vec::new();
    for choice in choices.iter() {
        let score = score(&choice, query);
        if score > 0_f64 {
            ret.push(ScoredChoice{ score: score, choice: &choice });
        }
    }
    ret.sort_by(|x, y| y.score.partial_cmp(&x.score).unwrap());
    ret.iter().map(|s| s.choice).collect()
}

#[allow(dead_code)]
fn get_slice_indices(length: usize, workers: usize, idx: usize) -> (usize, usize) {
    let lb = (length as f64 / workers as f64) * idx as f64;
    let ub = (length as f64 / workers as f64) * (idx + 1) as f64;
    (lb as usize, ub as usize)
}

#[test]
fn get_slice_indices_test() {
    assert_eq!(get_slice_indices(100, 1, 0), (0, 100));

    assert_eq!(get_slice_indices(100, 2, 1), (50, 100));
    assert_eq!(get_slice_indices(100, 2, 0), (0, 50));

    assert_eq!(get_slice_indices(100, 3, 0), (0, 33));
    assert_eq!(get_slice_indices(100, 3, 1), (33, 66));
    assert_eq!(get_slice_indices(100, 3, 2), (66, 100));

    assert_eq!(get_slice_indices(100, 4, 0), (0, 25));
    assert_eq!(get_slice_indices(100, 4, 1), (25, 50));
    assert_eq!(get_slice_indices(100, 4, 2), (50, 75));
    assert_eq!(get_slice_indices(100, 4, 3), (75, 100));

    assert_eq!(get_slice_indices(12, 12, 11), (11, 12));
    assert_eq!(get_slice_indices(12, 12, 0), (0, 1));

    assert_eq!(get_slice_indices(1, 2, 0), (0, 0));
    assert_eq!(get_slice_indices(1, 2, 1), (0, 1));
}

fn score(choice: &str, query: &str) -> f64 {
    if query.len() == 0 {
        return 1.0;
    }
    if choice.len() == 0 {
        return 0.0;
    }

    let query = query.to_ascii_lowercase();
    let choice = choice.to_ascii_lowercase();

    match compute_match_length(&choice, &query) {
        None => return 0.0,
        Some(match_length) => {
            let score = query.len() as f64 / match_length as f64;
            return score as f64 / choice.len() as f64;
        }
    }
}

fn compute_match_length(string: &str, chars: &str) -> Option<usize> {
    let (first_char, rest) = match chars.slice_shift_char() {
        Some((f, r)) => (f, r),
        None => return None,
    };
    let indices = find_char_in_string(string, first_char);

    let mut current_min = None;
    for i in indices.iter() {
        let last_index = find_end_of_match(string, rest, *i);
        if last_index.is_some() {
            let idx = last_index.unwrap() - *i + 1;
            if current_min.is_some() {
                let cm = min(current_min.unwrap(), idx);
                current_min = Some(cm);
            } else {
                current_min = Some(idx)
            }
        }
    }
    return current_min;
}

fn find_char_in_string(string: &str, char: char) -> Vec<usize> {
    let mut indices = Vec::new();
    for (i, c) in string.chars().enumerate() {
        if char == c {
            indices.push(i);
        }
    }
    return indices;
}

fn find_end_of_match(string: &str, rest_of_query: &str, first_index: usize) -> Option<usize> {
    let mut last_index = first_index + 1;
    for c in rest_of_query.chars() {
        match string[last_index..].find(c) {
            None => return None,
            Some(ref i) => {
                last_index += *i + 1;
            }
        }
    }
    return Some(last_index - 1);
}

#[test]
fn find_end_of_match_test() {
    assert_eq!(find_end_of_match("a", "a", 0), None);
    assert_eq!(find_end_of_match("ba", "a", 1), None);
    assert_eq!(find_end_of_match("aaa", "aa", 0), Some(2));
    assert_eq!(find_end_of_match("aaa", "b", 0), None);
    assert_eq!(find_end_of_match("this is a long match", "this is a match", 0), None);
    assert_eq!(find_end_of_match("this is a long match", "his is a match", 0), Some(19));
    assert_eq!(find_end_of_match("./rust/x86_64-apple-darwin/test/run-pass/process-spawn-with-unicode-params-πЯ音æ∞/child.stage2-x86_64-apple-darwin", "ust", 2), Some(5));
}

#[test]
fn unicode_boundary_handling_test() {
    score("./rust/x86_64-apple-darwin/test/run-pass/process-spawn-with-unicode-params-πЯ音æ∞/child.stage2-x86_64-apple-darwin", "he");
}

#[test]
fn basic_scoring() {
    assert_eq!(score("", "a"), 0.0);
    assert_eq!(score("a", ""), 1.0);
    assert_eq!(score("short", "longer"), 0.0);
    assert_eq!(score("a", "b"), 0.0);
    assert_eq!(score("ab", "ac"), 0.0);

    assert!(score("a", "a") > 0.0);
    assert!(score("ab", "a") > 0.0);
    assert!(score("ba", "a") > 0.0);
    assert!(score("bab", "a") > 0.0);
    assert!(score("babababab", "aaaa") > 0.0);

    assert_eq!(score("a", "a"), 1_f64 / "a".len() as f64);
    assert_eq!(score("ab", "ab"), 0.5);
    assert_eq!(score("a long string", "a long string"), 1_f64 / "a long string".len() as f64);
    assert_eq!(score("spec/search_spec.rb", "sear"), 1_f64 / "spec/search_spec.rb".len() as f64);
}

#[test]
fn character_matching() {
    assert!(score("/! symbols $^", "/!$^") > 0.0);

    assert_eq!(score("a", "A"), 1.0);
    assert_eq!(score("A", "a"), 1.0);

    assert_eq!(score("a", "aa"), 0.0);
}

#[test]
fn match_equality() {
    assert!(score("selecta.gemspec", "asp") > score("algorithm4_spec.rb", "asp"));
    assert!(score("README.md", "em") > score("benchmark.rb", "em"));
    assert!(score("search.rb", "sear") > score("spec/search_spec.rb", "sear"));

    assert!(score("fbb", "fbb") > score("foo bar baz", "fbb"));
    assert!(score("foo", "foo") > score("longer foo", "foo"));
    assert!(score("foo", "foo") > score("foo longer", "foo"));
    assert!(score("1/2/3/4", "1/2/3") > score("1/9/2/3/4", "1/2/3"));

    assert!(score("long 12 long", "12") > score("1 long 2", "12"));

    assert_eq!(score("121padding2", "12"), 1.0 / "121padding2".len() as f64);
    assert_eq!(score("1padding212", "12"), 1.0 / "1padding212".len() as f64);
}
<commit_msg>Use &[char] instead of &str when matching<commit_after>use std::ascii::AsciiExt;
use std::cmp::min;

macro_rules! chars {
    ($str:expr) => (
        &$str.chars().collect::<Vec<char>>()
    );
}

pub fn compute_matches<'a>(choices: &[&'a str], query: &str) -> Vec<&'a str> {
    struct ScoredChoice<'a> {
        score: f64,
        choice: &'a str,
    };
    let mut ret = Vec::new();
    for choice in choices.iter() {
        let score = score(&choice, query);
        if score > 0_f64 {
            ret.push(ScoredChoice{ score: score, choice: &choice });
        }
    }
    ret.sort_by(|x, y| y.score.partial_cmp(&x.score).unwrap());
    ret.iter().map(|s| s.choice).collect()
}

#[allow(dead_code)]
fn get_slice_indices(length: usize, workers: usize, idx: usize) -> (usize, usize) {
    let lb = (length as f64 / workers as f64) * idx as f64;
    let ub = (length as f64 / workers as f64) * (idx + 1) as f64;
    (lb as usize, ub as usize)
}

#[test]
fn get_slice_indices_test() {
    assert_eq!(get_slice_indices(100, 1, 0), (0, 100));

    assert_eq!(get_slice_indices(100, 2, 1), (50, 100));
    assert_eq!(get_slice_indices(100, 2, 0), (0, 50));

    assert_eq!(get_slice_indices(100, 3, 0), (0, 33));
    assert_eq!(get_slice_indices(100, 3, 1), (33, 66));
    assert_eq!(get_slice_indices(100, 3, 2), (66, 100));

    assert_eq!(get_slice_indices(100, 4, 0), (0, 25));
    assert_eq!(get_slice_indices(100, 4, 1), (25, 50));
    assert_eq!(get_slice_indices(100, 4, 2), (50, 75));
    assert_eq!(get_slice_indices(100, 4, 3), (75, 100));

    assert_eq!(get_slice_indices(12, 12, 11), (11, 12));
    assert_eq!(get_slice_indices(12, 12, 0), (0, 1));

    assert_eq!(get_slice_indices(1, 2, 0), (0, 0));
    assert_eq!(get_slice_indices(1, 2, 1), (0, 1));
}

fn score(choice: &str, query: &str) -> f64 {
    if query.len() == 0 {
        return 1.0;
    }
    if choice.len() == 0 {
        return 0.0;
    }

    let query = query.to_ascii_lowercase();
    let choice = choice.to_ascii_lowercase();

    match compute_match_length(chars!(choice), chars!(query)) {
        None => return 0.0,
        Some(match_length) => {
            let score = query.len() as f64 / match_length as f64;
            return score as f64 / choice.len() as f64;
        }
    }
}

fn compute_match_length(string: &[char], chars: &[char]) -> Option<usize> {
    let first_char = chars[0];
    let rest = &chars[1..];
    let indices = find_char_in_string(string, first_char);

    let mut current_min = None;
    for i in indices.iter() {
        let last_index = find_end_of_match(string, rest, *i);
        if last_index.is_some() {
            let idx = last_index.unwrap() - *i + 1;
            if current_min.is_some() {
                let cm = min(current_min.unwrap(), idx);
                current_min = Some(cm);
            } else {
                current_min = Some(idx)
            }
        }
    }
    return current_min;
}

fn find_char_in_string(string: &[char], char: char) -> Vec<usize> {
    let mut indices = Vec::new();
    for (i, c) in string.iter().enumerate() {
        if char == *c {
            indices.push(i);
        }
    }
    return indices;
}

fn find_end_of_match(string: &[char], rest_of_query: &[char], first_index: usize) -> Option<usize> {
    let mut last_index = first_index + 1;
    for c in rest_of_query.iter() {
        let current_substring = &string[last_index..];
        let mut index = None;
        for (i, x) in current_substring.iter().enumerate() {
            if x == c {
                index = Some(i);
                break;
            }
        }
        if index.is_some() {
            last_index += index.unwrap() + 1;
        } else {
            return None;
        }
    }
    return Some(last_index - 1);
}

#[test]
fn find_end_of_match_test() {
    assert_eq!(find_end_of_match(chars!("a"), chars!("a"), 0), None);
    assert_eq!(find_end_of_match(chars!("ba"), chars!("a"), 1), None);
    assert_eq!(find_end_of_match(chars!("aaa"), chars!("aa"), 0), Some(2));
    assert_eq!(find_end_of_match(chars!("aaa"), chars!("b"), 0), None);
    assert_eq!(find_end_of_match(chars!("this is a long match"), chars!("this is a match"), 0), None);
    assert_eq!(find_end_of_match(chars!("this is a long match"), chars!("his is a match"), 0), Some(19));
    assert_eq!(find_end_of_match(chars!("./rust/x86_64-apple-darwin/test/run-pass/process-spawn-with-unicode-params-πЯ音æ∞/child.stage2-x86_64-apple-darwin"), chars!("ust"), 2), Some(5));
}

#[test]
fn unicode_boundary_handling_test() {
    score("./rust/x86_64-apple-darwin/test/run-pass/process-spawn-with-unicode-params-πЯ音æ∞/child.stage2-x86_64-apple-darwin", "he");
}

#[test]
fn basic_scoring() {
    assert_eq!(score("", "a"), 0.0);
    assert_eq!(score("a", ""), 1.0);
    assert_eq!(score("short", "longer"), 0.0);
    assert_eq!(score("a", "b"), 0.0);
    assert_eq!(score("ab", "ac"), 0.0);

    assert!(score("a", "a") > 0.0);
    assert!(score("ab", "a") > 0.0);
    assert!(score("ba", "a") > 0.0);
    assert!(score("bab", "a") > 0.0);
    assert!(score("babababab", "aaaa") > 0.0);

    assert_eq!(score("a", "a"), 1_f64 / "a".len() as f64);
    assert_eq!(score("ab", "ab"), 0.5);
    assert_eq!(score("a long string", "a long string"), 1_f64 / "a long string".len() as f64);
    assert_eq!(score("spec/search_spec.rb", "sear"), 1_f64 / "spec/search_spec.rb".len() as f64);
}

#[test]
fn character_matching() {
    assert!(score("/! symbols $^", "/!$^") > 0.0);

    assert_eq!(score("a", "A"), 1.0);
    assert_eq!(score("A", "a"), 1.0);

    assert_eq!(score("a", "aa"), 0.0);
}

#[test]
fn match_equality() {
    assert!(score("selecta.gemspec", "asp") > score("algorithm4_spec.rb", "asp"));
    assert!(score("README.md", "em") > score("benchmark.rb", "em"));
    assert!(score("search.rb", "sear") > score("spec/search_spec.rb", "sear"));

    assert!(score("fbb", "fbb") > score("foo bar baz", "fbb"));
    assert!(score("foo", "foo") > score("longer foo", "foo"));
    assert!(score("foo", "foo") > score("foo longer", "foo"));
    assert!(score("1/2/3/4", "1/2/3") > score("1/9/2/3/4", "1/2/3"));

    assert!(score("long 12 long", "12") > score("1 long 2", "12"));

    assert_eq!(score("121padding2", "12"), 1.0 / "121padding2".len() as f64);
    assert_eq!(score("1padding212", "12"), 1.0 / "1padding212".len() as f64);
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::MAX_TX_SIZE;
use crate::cpu_features::CpuFeatureLevel;
use crate::predict::{native, PredictionMode, PredictionVariant};
use crate::tiling::PlaneRegionMut;
use crate::transform::TxSize;
use crate::util::AlignedArray;
use crate::Pixel;
use libc;
use std::mem::size_of;

macro_rules! decl_angular_ipred_fn {
  ($($f:ident),+) => {
    extern {
      $(
        fn $f(
          dst: *mut u8, stride: libc::ptrdiff_t, topleft: *const u8,
          width: libc::c_int, height: libc::c_int, angle: libc::c_int,
        );
      )*
    }
  };
}

decl_angular_ipred_fn! {
  rav1e_ipred_dc_avx2,
  rav1e_ipred_dc_ssse3,
  rav1e_ipred_dc_128_avx2,
  rav1e_ipred_dc_128_ssse3,
  rav1e_ipred_dc_left_avx2,
  rav1e_ipred_dc_left_ssse3,
  rav1e_ipred_dc_top_avx2,
  rav1e_ipred_dc_top_ssse3,
  rav1e_ipred_v_avx2,
  rav1e_ipred_v_ssse3,
  rav1e_ipred_h_avx2,
  rav1e_ipred_h_ssse3,
  rav1e_ipred_z1_avx2,
  rav1e_ipred_z2_avx2,
  rav1e_ipred_z3_avx2,
  rav1e_ipred_smooth_avx2,
  rav1e_ipred_smooth_ssse3,
  rav1e_ipred_smooth_v_avx2,
  rav1e_ipred_smooth_v_ssse3,
  rav1e_ipred_smooth_h_avx2,
  rav1e_ipred_smooth_h_ssse3,
  rav1e_ipred_paeth_avx2,
  rav1e_ipred_paeth_ssse3
}

macro_rules! decl_cfl_pred_fn {
  ($($f:ident),+) => {
    extern {
      $(
        fn $f(
          dst: *mut u8, stride: libc::ptrdiff_t, topleft: *const u8,
          width: libc::c_int, height: libc::c_int, ac: *const u8,
          alpha: libc::c_int,
        );
      )*
    }
  };
}

decl_cfl_pred_fn! {
  rav1e_ipred_cfl_avx2,
  rav1e_ipred_cfl_ssse3,
  rav1e_ipred_cfl_128_avx2,
  rav1e_ipred_cfl_128_ssse3,
  rav1e_ipred_cfl_left_avx2,
  rav1e_ipred_cfl_left_ssse3,
  rav1e_ipred_cfl_top_avx2,
  rav1e_ipred_cfl_top_ssse3
}

#[inline(always)]
pub fn dispatch_predict_intra<T: Pixel>(
  mode: PredictionMode, variant: PredictionVariant,
  dst: &mut PlaneRegionMut<'_, T>, tx_size: TxSize, bit_depth: usize,
  ac: &[i16], angle: isize, edge_buf: &AlignedArray<[T; 4 * MAX_TX_SIZE + 1]>,
  cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::dispatch_predict_intra(
      mode, variant, dst, tx_size, bit_depth, ac, angle, edge_buf, cpu,
    );
  };

  if size_of::<T>() != 1 {
    return call_native(dst);
  }

  unsafe {
    let dst_ptr = dst.data_ptr_mut() as *mut _;
    let stride = dst.plane_cfg.stride as libc::ptrdiff_t;
    let edge_ptr =
      edge_buf.array.as_ptr().offset(2 * MAX_TX_SIZE as isize) as *const _;
    let w = tx_size.width() as libc::c_int;
    let h = tx_size.height() as libc::c_int;
    let angle = angle as libc::c_int;

    if cpu >= CpuFeatureLevel::AVX2 {
      match mode {
        PredictionMode::DC_PRED => {
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_dc_128_avx2,
            PredictionVariant::LEFT => rav1e_ipred_dc_left_avx2,
            PredictionVariant::TOP => rav1e_ipred_dc_top_avx2,
            PredictionVariant::BOTH => rav1e_ipred_dc_avx2,
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED if angle == 90 => {
          rav1e_ipred_v_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::H_PRED if angle == 180 => {
          rav1e_ipred_h_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED
        | PredictionMode::H_PRED
        | PredictionMode::D45_PRED
        | PredictionMode::D135_PRED
        | PredictionMode::D113_PRED
        | PredictionMode::D157_PRED
        | PredictionMode::D203_PRED
        | PredictionMode::D67_PRED => {
          (if angle <= 90 {
            rav1e_ipred_z1_avx2
          } else if angle < 180 {
            rav1e_ipred_z2_avx2
          } else {
            rav1e_ipred_z3_avx2
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_PRED => {
          rav1e_ipred_smooth_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_V_PRED => {
          rav1e_ipred_smooth_v_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_H_PRED => {
          rav1e_ipred_smooth_h_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::PAETH_PRED => {
          rav1e_ipred_paeth_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::UV_CFL_PRED => {
          let ac_ptr = ac.as_ptr() as *const _;
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_cfl_128_avx2,
            PredictionVariant::LEFT => rav1e_ipred_cfl_left_avx2,
            PredictionVariant::TOP => rav1e_ipred_cfl_top_avx2,
            PredictionVariant::BOTH => rav1e_ipred_cfl_avx2,
          })(dst_ptr, stride, edge_ptr, w, h, ac_ptr, angle);
        }
        _ => call_native(dst),
      }
    } else if cpu >= CpuFeatureLevel::SSSE3 {
      match mode {
        PredictionMode::DC_PRED => {
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_dc_128_ssse3,
            PredictionVariant::LEFT => rav1e_ipred_dc_left_ssse3,
            PredictionVariant::TOP => rav1e_ipred_dc_top_ssse3,
            PredictionVariant::BOTH => rav1e_ipred_dc_ssse3,
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED if angle == 90 => {
          rav1e_ipred_v_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::H_PRED if angle == 180 => {
          rav1e_ipred_h_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_PRED => {
          rav1e_ipred_smooth_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_V_PRED => {
          rav1e_ipred_smooth_v_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_H_PRED => {
          rav1e_ipred_smooth_h_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::PAETH_PRED => {
          rav1e_ipred_paeth_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::UV_CFL_PRED => {
          let ac_ptr = ac.as_ptr() as *const _;
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_cfl_128_ssse3,
            PredictionVariant::LEFT => rav1e_ipred_cfl_left_ssse3,
            PredictionVariant::TOP => rav1e_ipred_cfl_top_ssse3,
            PredictionVariant::BOTH => rav1e_ipred_cfl_ssse3,
          })(dst_ptr, stride, edge_ptr, w, h, ac_ptr, angle);
        }
        _ => call_native(dst),
      }
    } else {
      call_native(dst);
    }
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use crate::frame::{AsRegion, Plane};
  use crate::predict::native;
  use num_traits::*;

  #[test]
  fn pred_matches_u8() {
    let tx_size = TxSize::TX_4X4;
    let bit_depth = 8;
    let cpu = CpuFeatureLevel::default();
    let ac = [0i16; 32 * 32];

    let mut edge_buf: AlignedArray<[u8; 4 * MAX_TX_SIZE + 1]> =
      AlignedArray::uninitialized();
    for i in 0..edge_buf.array.len() {
      edge_buf.array[i] = (i + 32).saturating_sub(2 * MAX_TX_SIZE).as_();
    }

    for (mode, variant) in [
      (PredictionMode::DC_PRED, PredictionVariant::BOTH),
      (PredictionMode::DC_PRED, PredictionVariant::TOP),
      (PredictionMode::DC_PRED, PredictionVariant::LEFT),
      (PredictionMode::DC_PRED, PredictionVariant::NONE),
      (PredictionMode::V_PRED, PredictionVariant::BOTH),
      (PredictionMode::H_PRED, PredictionVariant::BOTH),
      (PredictionMode::D45_PRED, PredictionVariant::BOTH),
      (PredictionMode::D135_PRED, PredictionVariant::BOTH),
      (PredictionMode::D203_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_V_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_H_PRED, PredictionVariant::BOTH),
      (PredictionMode::PAETH_PRED, PredictionVariant::BOTH),
    ]
    .iter()
    {
      let angles = match mode {
        PredictionMode::V_PRED => [81, 84, 87, 90, 93, 96, 99].iter(),
        PredictionMode::H_PRED => [171, 174, 177, 180, 183, 186, 189].iter(),
        PredictionMode::D45_PRED => [
          3, 6, 9, 14, 17, 20, 23, 26, 29, 32, 36, 39, 42, 45, 48, 51, 54, 58,
          61, 64, 67, 70, 73, 76,
        ]
        .iter(),
        PredictionMode::D135_PRED => [
          104, 107, 110, 113, 116, 119, 122, 126, 129, 132, 135, 138, 141,
          144, 148, 151, 154, 157, 160, 163, 166,
        ]
        .iter(),
        PredictionMode::D203_PRED => [
          194, 197, 200, 203, 206, 209, 212, 216, 219, 222, 225, 228, 231,
          234, 238, 241, 244, 247, 250, 253, 256, 261, 264, 267,
        ]
        .iter(),
        _ => [0].iter(),
      };
      for angle in angles {
        let expected = {
          let mut plane = Plane::wrap(vec![0u8; 4 * 4], 4);
          native::dispatch_predict_intra(
            *mode,
            *variant,
            &mut plane.as_region_mut(),
            tx_size,
            bit_depth,
            &ac,
            *angle,
            &edge_buf,
            cpu,
          );
          let mut data = [0u8; 4 * 4];
          for (v, d) in data.iter_mut().zip(plane.data[..].iter()) {
            *v = *d;
          }
          data
        };

        let mut output = Plane::wrap(vec![0u8; 4 * 4], 4);
        dispatch_predict_intra(
          *mode,
          *variant,
          &mut output.as_region_mut(),
          tx_size,
          bit_depth,
          &ac,
          *angle,
          &edge_buf,
          cpu,
        );
        assert_eq!(expected, &output.data[..]);
      }
    }
  }
}
<commit_msg>Group angles by mode in intra prediction ASM tests<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::MAX_TX_SIZE;
use crate::cpu_features::CpuFeatureLevel;
use crate::predict::{native, PredictionMode, PredictionVariant};
use crate::tiling::PlaneRegionMut;
use crate::transform::TxSize;
use crate::util::AlignedArray;
use crate::Pixel;
use libc;
use std::mem::size_of;

macro_rules! decl_angular_ipred_fn {
  ($($f:ident),+) => {
    extern {
      $(
        fn $f(
          dst: *mut u8, stride: libc::ptrdiff_t, topleft: *const u8,
          width: libc::c_int, height: libc::c_int, angle: libc::c_int,
        );
      )*
    }
  };
}

decl_angular_ipred_fn! {
  rav1e_ipred_dc_avx2,
  rav1e_ipred_dc_ssse3,
  rav1e_ipred_dc_128_avx2,
  rav1e_ipred_dc_128_ssse3,
  rav1e_ipred_dc_left_avx2,
  rav1e_ipred_dc_left_ssse3,
  rav1e_ipred_dc_top_avx2,
  rav1e_ipred_dc_top_ssse3,
  rav1e_ipred_v_avx2,
  rav1e_ipred_v_ssse3,
  rav1e_ipred_h_avx2,
  rav1e_ipred_h_ssse3,
  rav1e_ipred_z1_avx2,
  rav1e_ipred_z2_avx2,
  rav1e_ipred_z3_avx2,
  rav1e_ipred_smooth_avx2,
  rav1e_ipred_smooth_ssse3,
  rav1e_ipred_smooth_v_avx2,
  rav1e_ipred_smooth_v_ssse3,
  rav1e_ipred_smooth_h_avx2,
  rav1e_ipred_smooth_h_ssse3,
  rav1e_ipred_paeth_avx2,
  rav1e_ipred_paeth_ssse3
}

macro_rules! decl_cfl_pred_fn {
  ($($f:ident),+) => {
    extern {
      $(
        fn $f(
          dst: *mut u8, stride: libc::ptrdiff_t, topleft: *const u8,
          width: libc::c_int, height: libc::c_int, ac: *const u8,
          alpha: libc::c_int,
        );
      )*
    }
  };
}

decl_cfl_pred_fn! {
  rav1e_ipred_cfl_avx2,
  rav1e_ipred_cfl_ssse3,
  rav1e_ipred_cfl_128_avx2,
  rav1e_ipred_cfl_128_ssse3,
  rav1e_ipred_cfl_left_avx2,
  rav1e_ipred_cfl_left_ssse3,
  rav1e_ipred_cfl_top_avx2,
  rav1e_ipred_cfl_top_ssse3
}

#[inline(always)]
pub fn dispatch_predict_intra<T: Pixel>(
  mode: PredictionMode, variant: PredictionVariant,
  dst: &mut PlaneRegionMut<'_, T>, tx_size: TxSize, bit_depth: usize,
  ac: &[i16], angle: isize, edge_buf: &AlignedArray<[T; 4 * MAX_TX_SIZE + 1]>,
  cpu: CpuFeatureLevel,
) {
  let call_native = |dst: &mut PlaneRegionMut<'_, T>| {
    native::dispatch_predict_intra(
      mode, variant, dst, tx_size, bit_depth, ac, angle, edge_buf, cpu,
    );
  };

  if size_of::<T>() != 1 {
    return call_native(dst);
  }

  unsafe {
    let dst_ptr = dst.data_ptr_mut() as *mut _;
    let stride = dst.plane_cfg.stride as libc::ptrdiff_t;
    let edge_ptr =
      edge_buf.array.as_ptr().offset(2 * MAX_TX_SIZE as isize) as *const _;
    let w = tx_size.width() as libc::c_int;
    let h = tx_size.height() as libc::c_int;
    let angle = angle as libc::c_int;

    if cpu >= CpuFeatureLevel::AVX2 {
      match mode {
        PredictionMode::DC_PRED => {
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_dc_128_avx2,
            PredictionVariant::LEFT => rav1e_ipred_dc_left_avx2,
            PredictionVariant::TOP => rav1e_ipred_dc_top_avx2,
            PredictionVariant::BOTH => rav1e_ipred_dc_avx2,
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED if angle == 90 => {
          rav1e_ipred_v_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::H_PRED if angle == 180 => {
          rav1e_ipred_h_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED
        | PredictionMode::H_PRED
        | PredictionMode::D45_PRED
        | PredictionMode::D135_PRED
        | PredictionMode::D113_PRED
        | PredictionMode::D157_PRED
        | PredictionMode::D203_PRED
        | PredictionMode::D67_PRED => {
          (if angle <= 90 {
            rav1e_ipred_z1_avx2
          } else if angle < 180 {
            rav1e_ipred_z2_avx2
          } else {
            rav1e_ipred_z3_avx2
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_PRED => {
          rav1e_ipred_smooth_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_V_PRED => {
          rav1e_ipred_smooth_v_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_H_PRED => {
          rav1e_ipred_smooth_h_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::PAETH_PRED => {
          rav1e_ipred_paeth_avx2(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::UV_CFL_PRED => {
          let ac_ptr = ac.as_ptr() as *const _;
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_cfl_128_avx2,
            PredictionVariant::LEFT => rav1e_ipred_cfl_left_avx2,
            PredictionVariant::TOP => rav1e_ipred_cfl_top_avx2,
            PredictionVariant::BOTH => rav1e_ipred_cfl_avx2,
          })(dst_ptr, stride, edge_ptr, w, h, ac_ptr, angle);
        }
        _ => call_native(dst),
      }
    } else if cpu >= CpuFeatureLevel::SSSE3 {
      match mode {
        PredictionMode::DC_PRED => {
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_dc_128_ssse3,
            PredictionVariant::LEFT => rav1e_ipred_dc_left_ssse3,
            PredictionVariant::TOP => rav1e_ipred_dc_top_ssse3,
            PredictionVariant::BOTH => rav1e_ipred_dc_ssse3,
          })(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::V_PRED if angle == 90 => {
          rav1e_ipred_v_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::H_PRED if angle == 180 => {
          rav1e_ipred_h_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_PRED => {
          rav1e_ipred_smooth_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_V_PRED => {
          rav1e_ipred_smooth_v_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::SMOOTH_H_PRED => {
          rav1e_ipred_smooth_h_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::PAETH_PRED => {
          rav1e_ipred_paeth_ssse3(dst_ptr, stride, edge_ptr, w, h, angle);
        }
        PredictionMode::UV_CFL_PRED => {
          let ac_ptr = ac.as_ptr() as *const _;
          (match variant {
            PredictionVariant::NONE => rav1e_ipred_cfl_128_ssse3,
            PredictionVariant::LEFT => rav1e_ipred_cfl_left_ssse3,
            PredictionVariant::TOP => rav1e_ipred_cfl_top_ssse3,
            PredictionVariant::BOTH => rav1e_ipred_cfl_ssse3,
          })(dst_ptr, stride, edge_ptr, w, h, ac_ptr, angle);
        }
        _ => call_native(dst),
      }
    } else {
      call_native(dst);
    }
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use crate::frame::{AsRegion, Plane};
  use crate::predict::native;
  use num_traits::*;

  #[test]
  fn pred_matches_u8() {
    let tx_size = TxSize::TX_4X4;
    let bit_depth = 8;
    let cpu = CpuFeatureLevel::default();
    let ac = [0i16; 32 * 32];

    let mut edge_buf: AlignedArray<[u8; 4 * MAX_TX_SIZE + 1]> =
      AlignedArray::uninitialized();
    for i in 0..edge_buf.array.len() {
      edge_buf.array[i] = (i + 32).saturating_sub(2 * MAX_TX_SIZE).as_();
    }

    for (mode, variant) in [
      (PredictionMode::DC_PRED, PredictionVariant::BOTH),
      (PredictionMode::DC_PRED, PredictionVariant::TOP),
      (PredictionMode::DC_PRED, PredictionVariant::LEFT),
      (PredictionMode::DC_PRED, PredictionVariant::NONE),
      (PredictionMode::V_PRED, PredictionVariant::BOTH),
      (PredictionMode::H_PRED, PredictionVariant::BOTH),
      (PredictionMode::D45_PRED, PredictionVariant::BOTH),
      (PredictionMode::D135_PRED, PredictionVariant::BOTH),
      (PredictionMode::D113_PRED, PredictionVariant::BOTH),
      (PredictionMode::D157_PRED, PredictionVariant::BOTH),
      (PredictionMode::D203_PRED, PredictionVariant::BOTH),
      (PredictionMode::D67_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_V_PRED, PredictionVariant::BOTH),
      (PredictionMode::SMOOTH_H_PRED, PredictionVariant::BOTH),
      (PredictionMode::PAETH_PRED, PredictionVariant::BOTH),
    ]
    .iter()
    {
      let angles = match mode {
        PredictionMode::V_PRED => [81, 84, 87, 90, 93, 96, 99].iter(),
        PredictionMode::H_PRED => [171, 174, 177, 180, 183, 186, 189].iter(),
        PredictionMode::D45_PRED => [36, 39, 42, 45, 48, 51, 54].iter(),
        PredictionMode::D135_PRED => {
          [126, 129, 132, 135, 138, 141, 144].iter()
        }
        PredictionMode::D113_PRED => {
          [104, 107, 110, 113, 116, 119, 122].iter()
        }
        PredictionMode::D157_PRED => {
          [148, 151, 154, 157, 160, 163, 166].iter()
        }
        PredictionMode::D203_PRED => {
          [194, 197, 200, 203, 206, 209, 212].iter()
        }
        PredictionMode::D67_PRED => [58, 61, 64, 67, 70, 73, 76].iter(),
        _ => [0].iter(),
      };
      for angle in angles {
        let expected = {
          let mut plane = Plane::wrap(vec![0u8; 4 * 4], 4);
          native::dispatch_predict_intra(
            *mode,
            *variant,
            &mut plane.as_region_mut(),
            tx_size,
            bit_depth,
            &ac,
            *angle,
            &edge_buf,
            cpu,
          );
          let mut data = [0u8; 4 * 4];
          for (v, d) in data.iter_mut().zip(plane.data[..].iter()) {
            *v = *d;
          }
          data
        };

        let mut output = Plane::wrap(vec![0u8; 4 * 4], 4);
        dispatch_predict_intra(
          *mode,
          *variant,
          &mut output.as_region_mut(),
          tx_size,
          bit_depth,
          &ac,
          *angle,
          &edge_buf,
          cpu,
        );
        assert_eq!(
          expected,
          &output.data[..],
          "mode={:?} variant={:?} angle={}",
          *mode,
          *variant,
          *angle
        );
      }
    }
  }
}
<|endoftext|>


<commit_before>use super::{custom_float, CustomFloat, Time};
use super::vector::Vector;

#[derive(Clone, Debug)]
pub struct Particle {
  pub x: Vector,
  pub v: Vector,
  pub r: CustomFloat,
  pub m: CustomFloat,
}

/// Computes solutions to the quadratic formula:
/// ( -b +/- sqrt(b^2 - 4ac) ) / 2a
/// The smaller solution is always on the left.
/// Returns None for imaginary results.
///
/// # Examples
/// ```
/// use particles::quadratic_formula;
/// let (s1, s2) = quadratic_formula(1., 0., -1.).unwrap();
/// assert!((s1 + 1.).abs() < 1e-10);
/// assert!((s2 - 1.).abs() < 1e-10);
/// ```
// TODO: make private, write unit tests for it
pub fn quadratic_formula(a: CustomFloat, b: CustomFloat, c: CustomFloat) -> Option<(CustomFloat, CustomFloat)> {
  // account for floating error: we can't have degenerate solutions come up imaginary
  let b2 = b * b + (b / 2.0_f32.powi((custom_float::MANTISSA_DIGITS - 1) as i32));
  println!("b2: {:?}", b2);
  let ac = 4. * a * c;
  println!("ac: {:?}", ac);

  if b2 < ac { None } // imaginary result
  else {
    let a2 = 2. * a;
    let fst = - b / a2;
    let snd = ((b2 - ac).sqrt() / a2).abs();
    Some((fst - snd, fst + snd))
  }
}

impl Particle {
  /// Computes the next time the two given particles will impact each other.
  /// Returns None if no such impact will occur.
  ///
  /// # Panics
  /// - if the two particles given overlap (i.e. they have fused together)
  ///
  /// # Examples
  /// ```
  /// use particles::{Particle, Time, Vector};
  /// let p1 = Particle {
  ///   x: Vector((-2., 0.)),
  ///   v: Vector((1., 0.)),
  ///   r: 1.,
  ///   m: 1.
  /// };
  /// let p2 = Particle {
  ///   x: Vector((2., 0.)),
  ///   v: Vector((-1., 0.)),
  ///   r: 1.,
  ///   m: 1.
  /// };
  /// let Time(t) = p1.impact_time(&p2).unwrap();
  /// assert!((t - 1.).abs() < 1e-10);
  /// ```
  pub fn impact_time(&self, other: &Particle) -> Option<Time> {
    // solves for t:
    // | self.x - self.x + (self.v - other.v) * t | = self.r + other.r
    let dv = &(&self.v - &other.v);
    let dx = &(&self.x - &other.x);
    let sr = self.r + other.r;

    // quadratic formula for t:
    // |dv|^2 t^2 + 2 (dx * dv) t + |dx|^2 - sr^2 = 0
    let a = dv.norm2();
    let b = 2. * ( dx * dv );
    let c = dx.norm2() - sr.powi(2);

    let s = quadratic_formula(a, b, c);
    match s {
      Some((less, _)) if less >= 0. => Some(Time(less)),
      Some((_, more)) if more <= 0. => None,
      None => None, // particles will never impact
      Some(solns) => // less < 0 && more > 0 <=> particles overlap
        unreachable!(
          "impact_time found overlapping particles:\n{:?}\n{:?}\n{:?}",
          solns, self, other
        ),
    }
  }

  /// Returns new particles after a collision.
  /// Assumes that the particles are tangent to each other.
  /// The first particle returned corresponds to self.
  ///
  /// # Panics
  /// - if the two particles are not tangent (or within 10e-5 units)
  ///
  /// # Examples
  /// ```
  /// use particles::{Particle, Time, Vector};
  /// let p1 = Particle {
  ///   x: Vector((-1., 0.)),
  ///   v: Vector((1., 0.)),
  ///   r: 1.,
  ///   m: 1.
  /// };
  /// let p2 = Particle {
  ///   x: Vector((1., 0.)),
  ///   v: Vector((-1., 0.)),
  ///   r: 1.,
  ///   m: 1.
  /// };
  /// let (p1_, p2_) = p1.bounce(&p2);
  /// assert!((&p1_.v - &Vector((-1., 0.))).norm() < 1e-10);
  /// assert!((&p2_.v - &Vector((1., 0.))).norm() < 1e-10);
  /// ```
  pub fn bounce(&self, other: &Particle) -> (Particle, Particle) {
    let r_t = self.r + other.r;
    let dx = &self.x - &other.x;
    // only works for particles in contact
    assert!((dx.norm() - r_t).abs() < 10e-5);

    let dv = &self.v - &other.v;
    let m_r = self.m * other.m / (self.m + other.m);

    // dp = 2 m1 m2 / (m1 + m2) (dv . \hat{dx}) \hat{dx}
    let dp = dx.scale(2. * m_r * (&dv * &dx) / dx.norm2());
    let v1 = &self.v - &dp.scale(1. / self.m);
    let v2 = &other.v + &dp.scale(1. / other.m);

    let p1 = Particle { v: v1, .. self.clone() };
    let p2 = Particle { v: v2, .. other.clone() };
    (p1, p2)
  }

  pub fn evolve(&self, t: Time) -> Particle {
    let Time(t_) = t;
    Particle { x: &self.x + &self.v.scale(t_), .. self.clone() }
  }
}

impl PartialEq for Particle {
  fn eq(&self, other: &Particle) -> bool {
    self.r.eq(&other.r) &&
    self.m.eq(&other.m) &&
    self.x.eq(&other.x) &&
    self.v.eq(&other.v)
  }

  fn ne(&self, other: &Particle) -> bool {
    self.r.ne(&other.r) ||
    self.m.ne(&other.m) ||
    self.x.ne(&other.x) ||
    self.v.ne(&other.v)
  }
}

impl Eq for Particle {}

<commit_msg>[t] move particles examples into unit test<commit_after>use super::{custom_float, CustomFloat, Time};
use super::vector::Vector;

#[derive(Clone, Debug)]
pub struct Particle {
  pub x: Vector,
  pub v: Vector,
  pub r: CustomFloat,
  pub m: CustomFloat,
}

/// Computes solutions to the quadratic formula:
/// ( -b +/- sqrt(b^2 - 4ac) ) / 2a
/// The smaller solution is always on the left.
/// Returns None for imaginary results.
pub fn quadratic_formula(a: CustomFloat, b: CustomFloat, c: CustomFloat) -> Option<(CustomFloat, CustomFloat)> {
  // account for rounding error:
  // in the case of b = 4ac, rounding error may cause b < 4ac.
  // So, we increment the last bit in the mantissa by one
  let b2 = b * b + (b / 2.0_f32.powi((custom_float::MANTISSA_DIGITS - 1) as i32));
  let ac = 4. * a * c;

  if b2 < ac { None } // imaginary result
  else {
    let a2 = 2. * a;
    let fst = - b / a2;
    let snd = ((b2 - ac).sqrt() / a2).abs();
    Some((fst - snd, fst + snd))
  }
}

impl Particle {
  /// Computes the next time the two given particles will impact each other.
  /// Returns None if no such impact will occur.
  ///
  /// # Panics
  /// - if the two particles given overlap (i.e. they have fused together)
  pub fn impact_time(&self, other: &Particle) -> Option<Time> {
    // solves for t:
    // | self.x - self.x + (self.v - other.v) * t | = self.r + other.r
    let dv = &(&self.v - &other.v);
    let dx = &(&self.x - &other.x);
    let sr = self.r + other.r;

    // quadratic formula for t:
    // |dv|^2 t^2 + 2 (dx * dv) t + |dx|^2 - sr^2 = 0
    let a = dv.norm2();
    let b = 2. * ( dx * dv );
    let c = dx.norm2() - sr.powi(2);

    let s = quadratic_formula(a, b, c);
    match s {
      Some((less, _)) if less >= 0. => Some(Time(less)),
      Some((_, more)) if more <= 0. => None,
      None => None, // particles will never impact
      Some(solns) => // less < 0 && more > 0 <=> particles overlap
        unreachable!(
          "impact_time found overlapping particles:\n{:?}\n{:?}\n{:?}",
          solns, self, other
        ),
    }
  }

  /// Returns new particles after a collision.
  /// Assumes that the particles are tangent to each other.
  /// The first particle returned corresponds to self.
  ///
  /// # Panics
  /// - if the two particles are not tangent (or within 10e-5 units)
  pub fn bounce(&self, other: &Particle) -> (Particle, Particle) {
    let r_t = self.r + other.r;
    let dx = &self.x - &other.x;
    // only works for particles in contact
    assert!((dx.norm() - r_t).abs() < 10e-5);

    let dv = &self.v - &other.v;
    let m_r = self.m * other.m / (self.m + other.m);

    // dp = 2 m1 m2 / (m1 + m2) (dv . \hat{dx}) \hat{dx}
    let dp = dx.scale(2. * m_r * (&dv * &dx) / dx.norm2());
    let v1 = &self.v - &dp.scale(1. / self.m);
    let v2 = &other.v + &dp.scale(1. / other.m);

    let p1 = Particle { v: v1, .. self.clone() };
    let p2 = Particle { v: v2, .. other.clone() };
    (p1, p2)
  }

  pub fn evolve(&self, t: Time) -> Particle {
    let Time(t_) = t;
    Particle { x: &self.x + &self.v.scale(t_), .. self.clone() }
  }
}

impl PartialEq for Particle {
  fn eq(&self, other: &Particle) -> bool {
    self.r.eq(&other.r) &&
    self.m.eq(&other.m) &&
    self.x.eq(&other.x) &&
    self.v.eq(&other.v)
  }

  fn ne(&self, other: &Particle) -> bool {
    self.r.ne(&other.r) ||
    self.m.ne(&other.m) ||
    self.x.ne(&other.x) ||
    self.v.ne(&other.v)
  }
}

impl Eq for Particle {}

#[cfg(test)]
mod tests {
  use super::super::{FloatOps, Particle, Time, Vector};

  #[test]
  fn quadratic_formula_simple() {
    use super::quadratic_formula;
    let (s1, s2) = quadratic_formula(1., 0., -1.).unwrap();
    assert!(FloatOps(s1).close(&FloatOps(-1.)));
    assert!(FloatOps(s2).close(&FloatOps(1.)));
  }

  #[test]
  fn bounce_symmetrical_particles() {
    let p1 = Particle {
      x: Vector((-1., 0.)),
      v: Vector((1., 0.)),
      r: 1.,
      m: 1.
    };
    let p2 = Particle {
      x: Vector((1., 0.)),
      v: Vector((-1., 0.)),
      r: 1.,
      m: 1.
    };
    let (p1_, p2_) = p1.bounce(&p2);
    assert!((&p1_.v - &Vector((-1., 0.))).norm() < 1e-10);
    assert!((&p2_.v - &Vector((1., 0.))).norm() < 1e-10);
  }

  #[test]
  fn impact_time_symmetrical_partices() {
    let p1 = Particle {
      x: Vector((-2., 0.)),
      v: Vector((1., 0.)),
      r: 1.,
      m: 1.
    };
    let p2 = Particle {
      x: Vector((2., 0.)),
      v: Vector((-1., 0.)),
      r: 1.,
      m: 1.
    };
    let Time(t) = p1.impact_time(&p2).unwrap();
    assert!((t - 1.).abs() < 1e-10);
  }
}
<|endoftext|>