<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A task that takes a URL and streams back the binary data.

use about_loader;
use data_loader;
use file_loader;
use http_loader;
use cookie_storage::CookieStorage;
use cookie;
use mime_classifier::MIMEClassifier;

use net_traits::{ControlMsg, LoadData, LoadResponse, LoadConsumer};
use net_traits::{Metadata, ProgressMsg, ResourceTask, AsyncResponseTarget, ResponseAction};
use net_traits::ProgressMsg::Done;
use util::opts;
use util::task::spawn_named;

use devtools_traits::{DevtoolsControlMsg};
use hyper::header::{ContentType, Header, SetCookie, UserAgent};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::borrow::ToOwned;
use std::boxed::{self, FnBox};
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::str::FromStr;
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};

static mut HOST_TABLE: Option<*mut HashMap<String, String>> = None;

pub fn global_init() {
    //TODO: handle bad file path
    let path = match env::var("HOST_FILE") {
        Ok(host_file_path) => host_file_path,
        Err(_) => return,
    };

    let mut file = match File::open(&path) {
        Ok(f) => BufReader::new(f),
        Err(_) => return,
    };

    let mut lines = String::new();
    match file.read_to_string(&mut lines) {
        Ok(_) => (),
        Err(_) => return,
    };

    unsafe {
        let host_table = boxed::into_raw(parse_hostsfile(&lines));
        HOST_TABLE = Some(host_table);
    }
}

pub enum ProgressSender {
    Channel(Sender<ProgressMsg>),
    Listener(Box<AsyncResponseTarget>),
}

impl ProgressSender {
    //XXXjdm return actual error
    pub fn send(&self, msg: ProgressMsg) -> Result<(), ()> {
        match *self {
            ProgressSender::Channel(ref c) => c.send(msg).map_err(|_| ()),
            ProgressSender::Listener(ref b) => {
                let action = match msg {
                    ProgressMsg::Payload(buf) => ResponseAction::DataAvailable(buf),
                    ProgressMsg::Done(status) => ResponseAction::ResponseComplete(status),
                };
                b.invoke_with_listener(action);
                Ok(())
            }
        }
    }
}

/// For use by loaders in responding to a Load message.
pub fn start_sending(start_chan: LoadConsumer, metadata: Metadata) -> ProgressSender {
    start_sending_opt(start_chan, metadata).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed(start_chan: LoadConsumer, metadata: Metadata,
                             classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                             -> ProgressSender {
    start_sending_sniffed_opt(start_chan, metadata, classifier, partial_body).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed_opt(start_chan: LoadConsumer, mut metadata: Metadata,
                                 classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                                 -> Result<ProgressSender, ()> {
    if opts::get().sniff_mime_types {
        // TODO: should be calculated in the resource loader, from pull requeset #4094
        let nosniff = false;
        let check_for_apache_bug = false;

        let supplied_type = metadata.content_type.map(|ContentType(Mime(toplevel, sublevel, _))| {
            (format!("{}", toplevel), format!("{}", sublevel))
        });
        metadata.content_type = classifier.classify(nosniff, check_for_apache_bug, &supplied_type,
                                                    &partial_body).map(|(toplevel, sublevel)| {
            let mime_tp: TopLevel = FromStr::from_str(&toplevel).unwrap();
            let mime_sb: SubLevel = FromStr::from_str(&sublevel).unwrap();
            ContentType(Mime(mime_tp, mime_sb, vec!()))
        });

    }

    start_sending_opt(start_chan, metadata)
}

/// For use by loaders in responding to a Load message.
pub fn start_sending_opt(start_chan: LoadConsumer, metadata: Metadata) -> Result<ProgressSender, ()> {
    match start_chan {
        LoadConsumer::Channel(start_chan) => {
            let (progress_chan, progress_port) = channel();
            let result = start_chan.send(LoadResponse {
                metadata:      metadata,
                progress_port: progress_port,
            });
            match result {
                Ok(_) => Ok(ProgressSender::Channel(progress_chan)),
                Err(_) => Err(())
            }
        }
        LoadConsumer::Listener(target) => {
            target.invoke_with_listener(ResponseAction::HeadersAvailable(metadata));
            Ok(ProgressSender::Listener(target))
        }
    }
}

/// Create a ResourceTask
pub fn new_resource_task(user_agent: Option<String>,
                         devtools_chan: Option<Sender<DevtoolsControlMsg>>) -> ResourceTask {
    let (setup_chan, setup_port) = channel();
    let setup_chan_clone = setup_chan.clone();
    spawn_named("ResourceManager".to_owned(), move || {
        ResourceManager::new(setup_port, user_agent, setup_chan_clone, devtools_chan).start();
    });
    setup_chan
}

pub fn parse_hostsfile(hostsfile_content: &str) -> Box<HashMap<String, String>> {
    let ipv4_regex = regex!(
        r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
    let ipv6_regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");
    let mut host_table = HashMap::new();
    let lines: Vec<&str> = hostsfile_content.split('\n').collect();

    for line in lines.iter() {
        let ip_host: Vec<&str> = line.trim().split(|c: char| c == ' ' || c == '\t').collect();
        if ip_host.len() > 1 {
            if !ipv4_regex.is_match(ip_host[0]) && !ipv6_regex.is_match(ip_host[0]) { continue; }
            let address = ip_host[0].to_owned();

            for token in ip_host.iter().skip(1) {
                if token.as_bytes()[0] == b'#' {
                    break;
                }
                host_table.insert(token.to_owned().to_string(), address.clone());
            }
        }
    }
    box host_table
}

pub fn replace_hosts(mut load_data: LoadData, host_table: *mut HashMap<String, String>) -> LoadData {
    if let Some(h) = load_data.url.domain_mut() {
        unsafe {
            if let Some(ip) = (*host_table).get(h) {
                *h = ip.clone();
            }
        }
    }
    return load_data;
}

struct ResourceManager {
    from_client: Receiver<ControlMsg>,
    user_agent: Option<String>,
    cookie_storage: CookieStorage,
    resource_task: Sender<ControlMsg>,
    mime_classifier: Arc<MIMEClassifier>,
    devtools_chan: Option<Sender<DevtoolsControlMsg>>
}

impl ResourceManager {
    fn new(from_client: Receiver<ControlMsg>,
           user_agent: Option<String>,
           resource_task: Sender<ControlMsg>,
           devtools_channel: Option<Sender<DevtoolsControlMsg>>) -> ResourceManager {
        ResourceManager {
            from_client: from_client,
            user_agent: user_agent,
            cookie_storage: CookieStorage::new(),
            resource_task: resource_task,
            mime_classifier: Arc::new(MIMEClassifier::new()),
            devtools_chan: devtools_channel
        }
    }
}

impl ResourceManager {
    fn start(&mut self) {
        loop {
            match self.from_client.recv().unwrap() {
              ControlMsg::Load(load_data, consumer) => {
                self.load(load_data, consumer)
              }
              ControlMsg::SetCookiesForUrl(request, cookie_list, source) => {
                let header = Header::parse_header(&[cookie_list.into_bytes()]);
                if let Some(SetCookie(cookies)) = header {
                  for bare_cookie in cookies.into_iter() {
                    if let Some(cookie) = cookie::Cookie::new_wrapped(bare_cookie, &request, source) {
                      self.cookie_storage.push(cookie, source);
                    }
                  }
                }
              }
              ControlMsg::GetCookiesForUrl(url, consumer, source) => {
                consumer.send(self.cookie_storage.cookies_for_url(&url, source)).unwrap();
              }
              ControlMsg::Exit => {
                break
              }
            }
        }
    }

    fn load(&mut self, mut load_data: LoadData, consumer: LoadConsumer) {
        unsafe {
            if let Some(host_table) = HOST_TABLE {
                load_data = replace_hosts(load_data, host_table);
            }
        }

        self.user_agent.as_ref().map(|ua| load_data.headers.set(UserAgent(ua.clone())));

        fn from_factory(factory: fn(LoadData, LoadConsumer, Arc<MIMEClassifier>))
                        -> Box<FnBox(LoadData, LoadConsumer, Arc<MIMEClassifier>) + Send> {
            box move |load_data, senders, classifier| {
                factory(load_data, senders, classifier)
            }
        }

        let loader = match &*load_data.url.scheme {
            "file" => from_factory(file_loader::factory),
            "http" | "https" | "view-source" =>
                http_loader::factory(self.resource_task.clone(), self.devtools_chan.clone()),
            "data" => from_factory(data_loader::factory),
            "about" => from_factory(about_loader::factory),
            _ => {
                debug!("resource_task: no loader for scheme {}", load_data.url.scheme);
                start_sending(consumer, Metadata::default(load_data.url))
                    .send(ProgressMsg::Done(Err("no loader for scheme".to_string()))).unwrap();
                return
            }
        };
        debug!("resource_task: loading url: {}", load_data.url.serialize());

        loader.call_box((load_data, consumer, self.mime_classifier.clone()));
    }
}
<commit_msg>Auto merge of #6171 - eleweek:fix-4184-ok, r=Manishearth<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A task that takes a URL and streams back the binary data.

use about_loader;
use data_loader;
use file_loader;
use http_loader;
use cookie_storage::CookieStorage;
use cookie;
use mime_classifier::MIMEClassifier;

use net_traits::{ControlMsg, LoadData, LoadResponse, LoadConsumer};
use net_traits::{Metadata, ProgressMsg, ResourceTask, AsyncResponseTarget, ResponseAction};
use net_traits::ProgressMsg::Done;
use util::opts;
use util::task::spawn_named;

use devtools_traits::{DevtoolsControlMsg};
use hyper::header::{ContentType, Header, SetCookie, UserAgent};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::borrow::ToOwned;
use std::boxed::{self, FnBox};
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::str::FromStr;
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};

static mut HOST_TABLE: Option<*mut HashMap<String, String>> = None;

pub fn global_init() {
    //TODO: handle bad file path
    let path = match env::var("HOST_FILE") {
        Ok(host_file_path) => host_file_path,
        Err(_) => return,
    };

    let mut file = match File::open(&path) {
        Ok(f) => BufReader::new(f),
        Err(_) => return,
    };

    let mut lines = String::new();
    match file.read_to_string(&mut lines) {
        Ok(_) => (),
        Err(_) => return,
    };

    unsafe {
        let host_table = boxed::into_raw(parse_hostsfile(&lines));
        HOST_TABLE = Some(host_table);
    }
}

pub enum ProgressSender {
    Channel(Sender<ProgressMsg>),
    Listener(Box<AsyncResponseTarget>),
}

impl ProgressSender {
    //XXXjdm return actual error
    pub fn send(&self, msg: ProgressMsg) -> Result<(), ()> {
        match *self {
            ProgressSender::Channel(ref c) => c.send(msg).map_err(|_| ()),
            ProgressSender::Listener(ref b) => {
                let action = match msg {
                    ProgressMsg::Payload(buf) => ResponseAction::DataAvailable(buf),
                    ProgressMsg::Done(status) => ResponseAction::ResponseComplete(status),
                };
                b.invoke_with_listener(action);
                Ok(())
            }
        }
    }
}

/// For use by loaders in responding to a Load message.
pub fn start_sending(start_chan: LoadConsumer, metadata: Metadata) -> ProgressSender {
    start_sending_opt(start_chan, metadata).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed(start_chan: LoadConsumer, metadata: Metadata,
                             classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                             -> ProgressSender {
    start_sending_sniffed_opt(start_chan, metadata, classifier, partial_body).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed_opt(start_chan: LoadConsumer, mut metadata: Metadata,
                                 classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                                 -> Result<ProgressSender, ()> {
    if opts::get().sniff_mime_types {
        // TODO: should be calculated in the resource loader, from pull requeset #4094
        let mut nosniff = false;
        let mut check_for_apache_bug = false;

        if let Some(ref headers) = metadata.headers {
            if let Some(ref raw_content_type) = headers.get_raw("content-type") {
                if raw_content_type.len() > 0 {
                    let ref last_raw_content_type = raw_content_type[raw_content_type.len() - 1];
                    check_for_apache_bug = last_raw_content_type == b"text/plain"
                                        || last_raw_content_type == b"text/plain; charset=ISO-8859-1"
                                        || last_raw_content_type == b"text/plain; charset=iso-8859-1"
                                        || last_raw_content_type == b"text/plain; charset=UTF-8";
                }
            }
            if let Some(ref raw_content_type_options) = headers.get_raw("X-content-type-options") {
                nosniff = raw_content_type_options.iter().any(|ref opt| *opt == b"nosniff");
            }
        }

        let supplied_type = metadata.content_type.map(|ContentType(Mime(toplevel, sublevel, _))| {
            (format!("{}", toplevel), format!("{}", sublevel))
        });
        metadata.content_type = classifier.classify(nosniff, check_for_apache_bug, &supplied_type,
                                                    &partial_body).map(|(toplevel, sublevel)| {
            let mime_tp: TopLevel = FromStr::from_str(&toplevel).unwrap();
            let mime_sb: SubLevel = FromStr::from_str(&sublevel).unwrap();
            ContentType(Mime(mime_tp, mime_sb, vec!()))
        });

    }

    start_sending_opt(start_chan, metadata)
}

/// For use by loaders in responding to a Load message.
pub fn start_sending_opt(start_chan: LoadConsumer, metadata: Metadata) -> Result<ProgressSender, ()> {
    match start_chan {
        LoadConsumer::Channel(start_chan) => {
            let (progress_chan, progress_port) = channel();
            let result = start_chan.send(LoadResponse {
                metadata:      metadata,
                progress_port: progress_port,
            });
            match result {
                Ok(_) => Ok(ProgressSender::Channel(progress_chan)),
                Err(_) => Err(())
            }
        }
        LoadConsumer::Listener(target) => {
            target.invoke_with_listener(ResponseAction::HeadersAvailable(metadata));
            Ok(ProgressSender::Listener(target))
        }
    }
}

/// Create a ResourceTask
pub fn new_resource_task(user_agent: Option<String>,
                         devtools_chan: Option<Sender<DevtoolsControlMsg>>) -> ResourceTask {
    let (setup_chan, setup_port) = channel();
    let setup_chan_clone = setup_chan.clone();
    spawn_named("ResourceManager".to_owned(), move || {
        ResourceManager::new(setup_port, user_agent, setup_chan_clone, devtools_chan).start();
    });
    setup_chan
}

pub fn parse_hostsfile(hostsfile_content: &str) -> Box<HashMap<String, String>> {
    let ipv4_regex = regex!(
        r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
    let ipv6_regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");
    let mut host_table = HashMap::new();
    let lines: Vec<&str> = hostsfile_content.split('\n').collect();

    for line in lines.iter() {
        let ip_host: Vec<&str> = line.trim().split(|c: char| c == ' ' || c == '\t').collect();
        if ip_host.len() > 1 {
            if !ipv4_regex.is_match(ip_host[0]) && !ipv6_regex.is_match(ip_host[0]) { continue; }
            let address = ip_host[0].to_owned();

            for token in ip_host.iter().skip(1) {
                if token.as_bytes()[0] == b'#' {
                    break;
                }
                host_table.insert(token.to_owned().to_string(), address.clone());
            }
        }
    }
    box host_table
}

pub fn replace_hosts(mut load_data: LoadData, host_table: *mut HashMap<String, String>) -> LoadData {
    if let Some(h) = load_data.url.domain_mut() {
        unsafe {
            if let Some(ip) = (*host_table).get(h) {
                *h = ip.clone();
            }
        }
    }
    return load_data;
}

struct ResourceManager {
    from_client: Receiver<ControlMsg>,
    user_agent: Option<String>,
    cookie_storage: CookieStorage,
    resource_task: Sender<ControlMsg>,
    mime_classifier: Arc<MIMEClassifier>,
    devtools_chan: Option<Sender<DevtoolsControlMsg>>
}

impl ResourceManager {
    fn new(from_client: Receiver<ControlMsg>,
           user_agent: Option<String>,
           resource_task: Sender<ControlMsg>,
           devtools_channel: Option<Sender<DevtoolsControlMsg>>) -> ResourceManager {
        ResourceManager {
            from_client: from_client,
            user_agent: user_agent,
            cookie_storage: CookieStorage::new(),
            resource_task: resource_task,
            mime_classifier: Arc::new(MIMEClassifier::new()),
            devtools_chan: devtools_channel
        }
    }
}

impl ResourceManager {
    fn start(&mut self) {
        loop {
            match self.from_client.recv().unwrap() {
              ControlMsg::Load(load_data, consumer) => {
                self.load(load_data, consumer)
              }
              ControlMsg::SetCookiesForUrl(request, cookie_list, source) => {
                let header = Header::parse_header(&[cookie_list.into_bytes()]);
                if let Some(SetCookie(cookies)) = header {
                  for bare_cookie in cookies.into_iter() {
                    if let Some(cookie) = cookie::Cookie::new_wrapped(bare_cookie, &request, source) {
                      self.cookie_storage.push(cookie, source);
                    }
                  }
                }
              }
              ControlMsg::GetCookiesForUrl(url, consumer, source) => {
                consumer.send(self.cookie_storage.cookies_for_url(&url, source)).unwrap();
              }
              ControlMsg::Exit => {
                break
              }
            }
        }
    }

    fn load(&mut self, mut load_data: LoadData, consumer: LoadConsumer) {
        unsafe {
            if let Some(host_table) = HOST_TABLE {
                load_data = replace_hosts(load_data, host_table);
            }
        }

        self.user_agent.as_ref().map(|ua| load_data.headers.set(UserAgent(ua.clone())));

        fn from_factory(factory: fn(LoadData, LoadConsumer, Arc<MIMEClassifier>))
                        -> Box<FnBox(LoadData, LoadConsumer, Arc<MIMEClassifier>) + Send> {
            box move |load_data, senders, classifier| {
                factory(load_data, senders, classifier)
            }
        }

        let loader = match &*load_data.url.scheme {
            "file" => from_factory(file_loader::factory),
            "http" | "https" | "view-source" =>
                http_loader::factory(self.resource_task.clone(), self.devtools_chan.clone()),
            "data" => from_factory(data_loader::factory),
            "about" => from_factory(about_loader::factory),
            _ => {
                debug!("resource_task: no loader for scheme {}", load_data.url.scheme);
                start_sending(consumer, Metadata::default(load_data.url))
                    .send(ProgressMsg::Done(Err("no loader for scheme".to_string()))).unwrap();
                return
            }
        };
        debug!("resource_task: loading url: {}", load_data.url.serialize());

        loader.call_box((load_data, consumer, self.mime_classifier.clone()));
    }
}
<|endoftext|>


<commit_before>

use fred;
use super::utils;

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

use tokio_timer::*;

use futures::future;
use futures::{
  IntoFuture,
  Future,
  Stream
};
use futures::sync::oneshot::{
  Sender as OneshotSender,
  Receiver as OneshotReceiver,
  channel as oneshot_channel
};
use futures::sync::mpsc::{
  Sender,
  Receiver,
  channel
};

use std::time::Duration;

use fred::error::{
  RedisErrorKind,
  RedisError
};

use fred::types::*;
use fred::RedisClient;

use std::thread;
use std::sync::Arc;

use pretty_env_logger;

use super::keys as keys_tests;
use super::hashes as hashes_tests;
use super::lists as lists_tests;

#[test]
fn it_should_connect_and_disconnect() {
  let mut config = RedisConfig::default_clustered();
  let mut core = Core::new().unwrap();
  let handle = core.handle();

  let client = RedisClient::new(config);
  let connection_ft = client.connect(&handle);

  let select_ft = client.on_connect().and_then(|client| {
    client.info(None)
  })
  .and_then(|(client, _)| {
    client.quit()
  });

  let (err, client) = core.run(connection_ft.join(select_ft)).unwrap();
  assert!(err.is_none());
}

#[test]
fn it_should_connect_and_disconnect_with_policy() {
  let mut config = RedisConfig::default_clustered();
  let policy = ReconnectPolicy::Constant {
    delay: 2000,
    attempts: 0,
    max_attempts: 10
  };

  let mut core = Core::new().unwrap();
  let handle = core.handle();

  let client = RedisClient::new(config);
  let connection_ft = client.connect_with_policy(&handle, policy);

  let select_ft = client.on_connect().and_then(|client| {
    client.info(None)
  })
  .and_then(|(client, _)| {
    client.quit()
  });

  let (_, client) = core.run(connection_ft.join(select_ft)).unwrap();
}

pub mod keys {
  use super::*;

  #[test]
  fn it_should_set_and_get_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_simple_key(client)
    });
  }

  #[test]
  fn it_should_set_and_get_large_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_large_key(client)
    })
  }

  #[test]
  fn it_should_set_and_get_random_keys() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_random_keys(client)
    })
  }

  #[test]
  fn it_should_set_random_keys_with_fqdn_addresses() {
    let hosts: Vec<(String, u16)> = vec![
      ("localhost".to_owned(), 30001),
      ("localhost".to_owned(), 30002),
      ("localhost".to_owned(), 30003),
    ];
    let config = RedisConfig::new_clustered(hosts, None);

    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_random_keys(client)
    })
  }

  #[cfg(all(feature="metrics", not(feature="mock")))]
  #[test]
  fn it_should_track_latency_and_size_metrics() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_track_latency_and_size(client)
    })
  }

}

pub mod hashes {
  use super::*;

  #[test]
  fn it_should_set_and_get_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_set_and_get_simple_key(client)
    });
  }

  #[test]
  fn it_should_set_and_get_all_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_set_and_get_all_simple_key(client)
    });
  }

  #[test]
  fn it_should_check_hexists() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_check_hexists(client)
    });
  }

  #[test]
  fn it_should_read_large_hash() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_read_large_hash(client)
    });
  }

}

pub mod lists {
  use super::*;

  #[test]
  fn it_should_llen_on_empty_list() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_llen_on_empty_list(client)
    });
  }

  #[test]
  fn it_should_llen_on_list_with_elements() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_llen_on_list_with_elements(client)
    });
  }

  #[test]
  fn it_should_lpush_and_lpop_to_list() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_lpush_and_lpop_to_list(client)
    });
  }
}
<commit_msg>Add set tests to cluster.rs<commit_after>

use fred;
use super::utils;

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

use tokio_timer::*;

use futures::future;
use futures::{
  IntoFuture,
  Future,
  Stream
};
use futures::sync::oneshot::{
  Sender as OneshotSender,
  Receiver as OneshotReceiver,
  channel as oneshot_channel
};
use futures::sync::mpsc::{
  Sender,
  Receiver,
  channel
};

use std::time::Duration;

use fred::error::{
  RedisErrorKind,
  RedisError
};

use fred::types::*;
use fred::RedisClient;

use std::thread;
use std::sync::Arc;

use pretty_env_logger;

use super::keys as keys_tests;
use super::hashes as hashes_tests;
use super::lists as lists_tests;
use super::sets as sets_tests;

#[test]
fn it_should_connect_and_disconnect() {
  let mut config = RedisConfig::default_clustered();
  let mut core = Core::new().unwrap();
  let handle = core.handle();

  let client = RedisClient::new(config);
  let connection_ft = client.connect(&handle);

  let select_ft = client.on_connect().and_then(|client| {
    client.info(None)
  })
  .and_then(|(client, _)| {
    client.quit()
  });

  let (err, client) = core.run(connection_ft.join(select_ft)).unwrap();
  assert!(err.is_none());
}

#[test]
fn it_should_connect_and_disconnect_with_policy() {
  let mut config = RedisConfig::default_clustered();
  let policy = ReconnectPolicy::Constant {
    delay: 2000,
    attempts: 0,
    max_attempts: 10
  };

  let mut core = Core::new().unwrap();
  let handle = core.handle();

  let client = RedisClient::new(config);
  let connection_ft = client.connect_with_policy(&handle, policy);

  let select_ft = client.on_connect().and_then(|client| {
    client.info(None)
  })
  .and_then(|(client, _)| {
    client.quit()
  });

  let (_, client) = core.run(connection_ft.join(select_ft)).unwrap();
}

pub mod keys {
  use super::*;

  #[test]
  fn it_should_set_and_get_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_simple_key(client)
    });
  }

  #[test]
  fn it_should_set_and_get_large_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_large_key(client)
    })
  }

  #[test]
  fn it_should_set_and_get_random_keys() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_random_keys(client)
    })
  }

  #[test]
  fn it_should_set_random_keys_with_fqdn_addresses() {
    let hosts: Vec<(String, u16)> = vec![
      ("localhost".to_owned(), 30001),
      ("localhost".to_owned(), 30002),
      ("localhost".to_owned(), 30003),
    ];
    let config = RedisConfig::new_clustered(hosts, None);

    utils::setup_test_client(config, |client| {
      keys_tests::should_set_and_get_random_keys(client)
    })
  }

  #[cfg(all(feature="metrics", not(feature="mock")))]
  #[test]
  fn it_should_track_latency_and_size_metrics() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      keys_tests::should_track_latency_and_size(client)
    })
  }

}

pub mod hashes {
  use super::*;

  #[test]
  fn it_should_set_and_get_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_set_and_get_simple_key(client)
    });
  }

  #[test]
  fn it_should_set_and_get_all_simple_key() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_set_and_get_all_simple_key(client)
    });
  }

  #[test]
  fn it_should_check_hexists() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_check_hexists(client)
    });
  }

  #[test]
  fn it_should_read_large_hash() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      hashes_tests::should_read_large_hash(client)
    });
  }

}

pub mod lists {
  use super::*;

  #[test]
  fn it_should_llen_on_empty_list() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_llen_on_empty_list(client)
    });
  }

  #[test]
  fn it_should_llen_on_list_with_elements() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_llen_on_list_with_elements(client)
    });
  }

  #[test]
  fn it_should_lpush_and_lpop_to_list() {
    let config = RedisConfig::default_clustered();
    utils::setup_test_client(config, |client| {
      lists_tests::should_lpush_and_lpop_to_list(client)
    });
  }
}

pub mod sets {
  use super::*;

  #[test]
  fn it_should_sadd_members_to_set() {
    let config = RedisConfig::default();
    utils::setup_test_client(config, |client| {
      sets_tests::should_sadd_members_to_set(client)
    });
  }

  #[test]
  fn it_should_srem_members_of_set() {
    let config = RedisConfig::default();
    utils::setup_test_client(config, |client| {
      sets_tests::should_srem_members_of_set(client)
    });
  }

   #[test]
  fn it_should_smembers_of_set() {
    let config = RedisConfig::default();
    utils::setup_test_client(config, |client| {
      sets_tests::should_smembers_of_set(client)
    });
  }
}

<|endoftext|>


<commit_before>use crate::{
    assist_context::{AssistContext, Assists},
    AssistId,
};
use syntax::{
    ast::{self, NameOwner},
    AstNode, SyntaxKind, SyntaxNode, SyntaxText,
};
use test_utils::mark;

// Assist: generate_default_from_new
//
// Generates default implementation from new method.
//
// ```
// struct Example { _inner: () }
//
// impl Example {
//     pub fn n$0ew() -> Self {
//         Self { _inner: () }
//     }
// }
// ```
// ->
// ```
// struct Example { _inner: () }
//
// impl Example {
//     pub fn new() -> Self {
//         Self { _inner: () }
//     }
// }
//
// impl Default for Example {
//     fn default() -> Self {
//         Self::new()
//     }
// }
// ```
pub(crate) fn generate_default_from_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let fn_node: ast::Fn = ctx.find_node_at_offset()?;
    let fn_name = fn_node.name()?.to_string();

    if !fn_name.eq("new") {
        mark::hit!(other_function_than_new);
        return None;
    }

    if fn_node.param_list()?.params().count() != 0 {
        mark::hit!(new_function_with_parameters);
        return None;
    }

    let insert_after = scope_for_fn_insertion_node(&fn_node.syntax())?;
    let impl_obj = ast::Impl::cast(insert_after)?;
    let struct_name = impl_obj.self_ty()?.syntax().text();

    let default_fn_syntax = default_fn_node_for_new(struct_name);

    acc.add(
        AssistId("generate_default_from_new", crate::AssistKind::Generate),
        "Generate a Default impl from a new fn",
        impl_obj.syntax().text_range(),
        move |builder| {
            // FIXME: indentation logic can also go here.
            // let new_indent = IndentLevel::from_node(&insert_after);
            let insert_location = impl_obj.syntax().text_range().end();
            builder.insert(insert_location, default_fn_syntax);
        },
    )
}

fn scope_for_fn_insertion_node(node: &SyntaxNode) -> Option<SyntaxNode> {
    node.ancestors().into_iter().find(|node| node.kind() == SyntaxKind::IMPL)
}

fn default_fn_node_for_new(struct_name: SyntaxText) -> String {
    // FIXME: Update the implementation to consider the code indentation.
    format!(
        r#"

impl Default for {} {{
    fn default() -> Self {{
        Self::new()
    }}
}}"#,
        struct_name
    )
}

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

    use super::*;

    #[test]
    fn generate_default() {
        check_assist(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn ne$0w() -> Self {
        Self { _inner: () }
    }
}

fn main() {}
"#,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn new() -> Self {
        Self { _inner: () }
    }
}

impl Default for Example {
    fn default() -> Self {
        Self::new()
    }
}

fn main() {}
"#,
        );
    }

    #[test]
    fn generate_default2() {
        check_assist(
            generate_default_from_new,
            r#"
struct Test { value: u32 }

impl Test {
    pub fn ne$0w() -> Self {
        Self { value: 0 }
    }
}
"#,
            r#"
struct Test { value: u32 }

impl Test {
    pub fn new() -> Self {
        Self { value: 0 }
    }
}

impl Default for Test {
    fn default() -> Self {
        Self::new()
    }
}
"#,
        );
    }

    #[test]
    fn new_function_with_parameters() {
        mark::check!(new_function_with_parameters);
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn $0new(value: ()) -> Self {
        Self { _inner: value }
    }
}
"#,
        );
    }

    #[test]
    fn other_function_than_new() {
        mark::check!(other_function_than_new);
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Exmaple {
    pub fn a$0dd() -> Self {
        Self { _inner: () }
    }
}

"#,
        );
    }

    //     #[test]
    //     fn default_block_is_already_present() {
    //         check_assist_not_applicable(generate_default_from_new,
    //         r#"
    // struct Example { _inner: () }

    // impl Exmaple {
    //     pub fn n$0ew() -> Self {
    //         Self { _inner: () }
    //     }
    // }

    // impl Default for Example {
    //     fn default() -> Self {
    //         Self::new()
    //     }
    // }
    // "#,
    //         );
    //     }

    #[test]
    fn standalone_new_function() {
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
fn n$0ew() -> u32 {
    0
}
"#,
        );
    }
}
<commit_msg>7708: Added the work for review comments.<commit_after>use crate::{
    assist_context::{AssistContext, Assists},
    utils, AssistId,
};
use syntax::{
    ast::{self, Adt, Impl, NameOwner},
    AstNode, Direction,
};
use test_utils::mark;

// Assist: generate_default_from_new
//
// Generates default implementation from new method.
//
// ```
// struct Example { _inner: () }
//
// impl Example {
//     pub fn n$0ew() -> Self {
//         Self { _inner: () }
//     }
// }
// ```
// ->
// ```
// struct Example { _inner: () }
//
// impl Example {
//     pub fn new() -> Self {
//         Self { _inner: () }
//     }
// }
//
// impl Default for Example {
//     fn default() -> Self {
//         Self::new()
//     }
// }
// ```
pub(crate) fn generate_default_from_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let fn_node = ctx.find_node_at_offset::<ast::Fn>()?;
    let fn_name = fn_node.name()?;

    if fn_name.text() != "new" {
        mark::hit!(other_function_than_new);
        return None;
    }

    if fn_node.param_list()?.params().next().is_some() {
        mark::hit!(new_function_with_parameters);
        return None;
    }

    let impl_ = fn_node.syntax().ancestors().into_iter().find_map(ast::Impl::cast)?;

    let insert_location = impl_.syntax().text_range();

    acc.add(
        AssistId("generate_default_from_new", crate::AssistKind::Generate),
        "Generate a Default impl from a new fn",
        insert_location,
        move |builder| {
            let default_fn_syntax = default_fn_node_for_new(impl_);
            if let Some(code) = default_fn_syntax {
                builder.insert(insert_location.end(), code)
            }
        },
    )
}

fn default_fn_node_for_new(impl_: Impl) -> Option<String> {
    // the code string is this way due to formatting reason
    let code = r#"    fn default() -> Self {
        Self::new()
    }"#;
    let struct_name = impl_.self_ty()?.syntax().to_string();
    let struct_ = impl_
        .syntax()
        .siblings(Direction::Prev)
        .filter_map(ast::Struct::cast)
        .find(|struct_| struct_.name().unwrap().text() == struct_name)?;

    let adt = Adt::cast(struct_.syntax().clone())?;
    Some(utils::generate_trait_impl_text(&adt, "Default", code))
}

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

    use super::*;

    #[test]
    fn generate_default() {
        check_assist(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn ne$0w() -> Self {
        Self { _inner: () }
    }
}

fn main() {}
"#,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn new() -> Self {
        Self { _inner: () }
    }
}

impl Default for Example {
    fn default() -> Self {
        Self::new()
    }
}

fn main() {}
"#,
        );
    }

    #[test]
    fn generate_default2() {
        check_assist(
            generate_default_from_new,
            r#"
struct Test { value: u32 }

impl Test {
    pub fn ne$0w() -> Self {
        Self { value: 0 }
    }
}
"#,
            r#"
struct Test { value: u32 }

impl Test {
    pub fn new() -> Self {
        Self { value: 0 }
    }
}

impl Default for Test {
    fn default() -> Self {
        Self::new()
    }
}
"#,
        );
    }

    #[test]
    fn new_function_with_parameters() {
        mark::check!(new_function_with_parameters);
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Example {
    pub fn $0new(value: ()) -> Self {
        Self { _inner: value }
    }
}
"#,
        );
    }

    #[test]
    fn other_function_than_new() {
        mark::check!(other_function_than_new);
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
struct Example { _inner: () }

impl Exmaple {
    pub fn a$0dd() -> Self {
        Self { _inner: () }
    }
}

"#,
        );
    }

    //     #[test]
    //     fn default_block_is_already_present() {
    //         check_assist_not_applicable(generate_default_from_new,
    //         r#"
    // struct Example { _inner: () }

    // impl Exmaple {
    //     pub fn n$0ew() -> Self {
    //         Self { _inner: () }
    //     }
    // }

    // impl Default for Example {
    //     fn default() -> Self {
    //         Self::new()
    //     }
    // }
    // "#,
    //         );
    //     }

    #[test]
    fn standalone_new_function() {
        check_assist_not_applicable(
            generate_default_from_new,
            r#"
fn n$0ew() -> u32 {
    0
}
"#,
        );
    }

    #[test]
    fn multiple_struct_blocks() {
        check_assist(
            generate_default_from_new,
            r#"
struct Example { _inner: () }
struct Test { value: u32 }

impl Example {
    pub fn new$0 () -> Self {
        Self { _inner: () }
    }
}
"#,
            r#"
struct Example { _inner: () }
struct Test { value: u32 }

impl Example {
    pub fn new () -> Self {
        Self { _inner: () }
    }
}

impl Default for Example {
    fn default() -> Self {
        Self::new()
    }
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>use sys;
use bullet_vector3::BulletVector3;
use collision::collision_shapes::Shape;
use mint::{Vector3, Vector4};

pub struct RigidBody {
    rigid_body: Box<sys::btRigidBody>,
    shape: Box<Shape>,
    motion_state: Box<sys::btDefaultMotionState>,
    construction_info: Box<sys::btRigidBody_btRigidBodyConstructionInfo>,
}

impl RigidBody {
    pub fn new<T1: Into<Vector3<f64>>, T2: Into<Vector3<f64>>, T3: Into<Vector4<f64>>>(
        mass: f64,
        inertia: T1,
        shape: Shape,
        translation: T2,
        orientation: T3,
    ) -> RigidBody {
        let mut inertia: BulletVector3 = inertia.into().into();
        let shape_box = Box::new(shape);
        let translation: BulletVector3 = translation.into().into();
        let orientation: [f64; 4] = orientation.into().into();
        let transform = unsafe {
            sys::btTransform::new1(
                &orientation as *const _ as *const _,
                &translation as *const _ as *const _,
            )
        };
        let mut motion_state_box = Box::new(unsafe {
            sys::btDefaultMotionState::new(&transform as *const _, sys::btTransform::getIdentity())
        });

        let construction_info_box = Box::new(unsafe {
            sys::btRigidBody_btRigidBodyConstructionInfo::new(
                mass,
                &mut *motion_state_box as *mut _ as *mut _,
                shape_box.as_ptr(),
                inertia.0.as_mut_ptr() as *mut _,
            )
        });
        RigidBody {
            rigid_body: Box::new(unsafe {
                sys::btRigidBody::new(&*construction_info_box as *const _)
            }),
            shape: shape_box,
            motion_state: motion_state_box,
            construction_info: construction_info_box,
        }
    }

    pub(crate) unsafe fn motion_state_ptr(&self) -> *mut sys::btDefaultMotionState {
        &*self.motion_state as *const _ as *mut _
    }

    pub(crate) unsafe fn as_ptr(&self) -> *mut sys::btRigidBody {
        &*self.rigid_body as *const _ as *mut _
    }
}

#[derive(Clone)]
pub struct RigidBodyHandle {
    pub(in dynamics) ptr: *mut sys::btRigidBody,
    motion_state: *mut sys::btDefaultMotionState,
    temp_transform: sys::btTransform,
}

impl RigidBodyHandle {
    pub fn new(ptr: *mut sys::btRigidBody, motion_state: *mut sys::btDefaultMotionState) -> Self {
        let temp_transform = unsafe { sys::btTransform::new() };

        RigidBodyHandle {
            ptr,
            motion_state,
            temp_transform,
        }
    }
    pub fn set_restitution(&mut self, restitution: f64) {
        unsafe {
            sys::btCollisionObject_setRestitution(self.ptr as *mut _, restitution);
        }
    }
    pub fn set_friction(&mut self, friction: f64) {
        unsafe {
            sys::btCollisionObject_setFriction(self.ptr as *mut _, friction);
        }
    }
    pub fn set_gravity<T: Into<Vector3<f64>>>(&mut self, gravity: T) {
        let gravity: BulletVector3 = gravity.into().into();
        unsafe {
            sys::btRigidBody_setGravity(self.ptr, gravity.0.as_ptr() as *const _);
        }
    }

    pub fn set_angular_factor<T: Into<Vector3<f64>>>(&mut self, angular_factor: T) {
        let angular_factor: BulletVector3 = angular_factor.into().into();
        unsafe {
            sys::btRigidBody_setAngularFactor(self.ptr, angular_factor.0.as_ptr() as *const _);
        }
    }

    pub fn set_sleeping_thresholds(&mut self, linear: f64, angular: f64) {
        unsafe { sys::btRigidBody_setSleepingThresholds(self.ptr, linear, angular) }
    }

    pub fn apply_central_impulse<T: Into<Vector3<f64>>>(&mut self, impulse: T) {
        let impulse: BulletVector3 = impulse.into().into();
        unsafe {
            sys::btRigidBody_applyCentralImpulse(self.ptr, impulse.0.as_ptr() as *const _);
        }
    }

    pub fn get_linear_velocity(&self) -> Vector3<f64> {
        let velocity = unsafe { sys::btRigidBody_getLinearVelocity(self.ptr) };

        ::bullet_vector3::vector_from_slice(unsafe {
            ::std::slice::from_raw_parts(velocity as *const _, 4)
        })
    }

    /// Override position vector and rotation quaternion.
    pub fn reset_position_and_orientation<T, T1>(&mut self, position: T, orientation: T1)
    where
        T: Into<Vector3<f64>>,
        T1: Into<Vector4<f64>>,
    {
        let orientation: [f64; 4] = orientation.into().into();
        let position: BulletVector3 = position.into().into();
        let transform = unsafe {
            sys::btTransform::new1(
                &orientation as *const _ as *const _,
                &position as *const _ as *const _,
            )
        };

        unsafe { (*self.ptr)._base.m_worldTransform = transform };
    }

    /// Get position in world space and orientation quaternion
    pub fn get_world_position_and_orientation(&self) -> (Vector3<f64>, Vector4<f64>) {
        unsafe {
            sys::btDefaultMotionState_getWorldTransform(
                &*self.motion_state as *const _ as *mut _,
                &self.temp_transform as *const _ as *mut _,
            );
        }
        let origin = unsafe { self.temp_transform.getOrigin1().as_ref().unwrap() };

        let rotation = unsafe { self.temp_transform.getRotation() };

        (
            ::bullet_vector3::vector_from_slice(&origin.m_floats[0..3]),
            ::bullet_vector3::vector4_from_slice(&rotation._base.m_floats),
        )
    }

    /// Place data on heap as a box and set to rigid body as user pointer
    pub fn set_user_data<T: 'static>(&mut self, data: T) {
        let data_box = Box::new(data);
        unsafe {
            sys::btCollisionObject_setUserPointer(
                self.ptr as *mut _,
                Box::into_raw(data_box) as *mut _,
            )
        };
    }

    /// Get data from rigidbody's user pointer
    /// Getting data from body without previously setted data is fine.
    /// Getting wrong typed data is unsafe and will cause mem::transmute to wrong pointer type.
    pub unsafe fn get_user_data<T: 'static>(&self) -> Option<&T> {
        let pointer = sys::btCollisionObject_getUserPointer(self.ptr as *mut _);
        let pointer: *const T = ::std::mem::transmute(pointer);
        return pointer.as_ref();
    }

    /// Set user index. This will not be used in bullet and this is not related to user_data.
    pub fn set_user_index(&mut self, index: i32) {
        unsafe { sys::btCollisionObject_setUserIndex(self.ptr as *mut _, index) };
    }

    /// Get previously setted user index
    /// If index was not set - will return "-1"
    pub fn get_user_index(&self) -> i32 {
        unsafe { sys::btCollisionObject_getUserIndex(self.ptr as *mut _) }
    }

    /// Was that rigid_body removed with DynamicsWorld::remove_body()
    /// Probably may return true also for any not-added bodys
    pub fn removed(&self) -> bool {
        unsafe { (*self.ptr)._base.m_worldArrayIndex == -1 }
    }
}
<commit_msg>add RigidBodyHandle::reset_linear_velocity<commit_after>use sys;
use bullet_vector3::BulletVector3;
use collision::collision_shapes::Shape;
use mint::{Vector3, Vector4};

pub struct RigidBody {
    rigid_body: Box<sys::btRigidBody>,
    shape: Box<Shape>,
    motion_state: Box<sys::btDefaultMotionState>,
    construction_info: Box<sys::btRigidBody_btRigidBodyConstructionInfo>,
}

impl RigidBody {
    pub fn new<T1: Into<Vector3<f64>>, T2: Into<Vector3<f64>>, T3: Into<Vector4<f64>>>(
        mass: f64,
        inertia: T1,
        shape: Shape,
        translation: T2,
        orientation: T3,
    ) -> RigidBody {
        let mut inertia: BulletVector3 = inertia.into().into();
        let shape_box = Box::new(shape);
        let translation: BulletVector3 = translation.into().into();
        let orientation: [f64; 4] = orientation.into().into();
        let transform = unsafe {
            sys::btTransform::new1(
                &orientation as *const _ as *const _,
                &translation as *const _ as *const _,
            )
        };
        let mut motion_state_box = Box::new(unsafe {
            sys::btDefaultMotionState::new(&transform as *const _, sys::btTransform::getIdentity())
        });

        let construction_info_box = Box::new(unsafe {
            sys::btRigidBody_btRigidBodyConstructionInfo::new(
                mass,
                &mut *motion_state_box as *mut _ as *mut _,
                shape_box.as_ptr(),
                inertia.0.as_mut_ptr() as *mut _,
            )
        });
        RigidBody {
            rigid_body: Box::new(unsafe {
                sys::btRigidBody::new(&*construction_info_box as *const _)
            }),
            shape: shape_box,
            motion_state: motion_state_box,
            construction_info: construction_info_box,
        }
    }

    pub(crate) unsafe fn motion_state_ptr(&self) -> *mut sys::btDefaultMotionState {
        &*self.motion_state as *const _ as *mut _
    }

    pub(crate) unsafe fn as_ptr(&self) -> *mut sys::btRigidBody {
        &*self.rigid_body as *const _ as *mut _
    }
}

#[derive(Clone)]
pub struct RigidBodyHandle {
    pub(in dynamics) ptr: *mut sys::btRigidBody,
    motion_state: *mut sys::btDefaultMotionState,
    temp_transform: sys::btTransform,
}

impl RigidBodyHandle {
    pub fn new(ptr: *mut sys::btRigidBody, motion_state: *mut sys::btDefaultMotionState) -> Self {
        let temp_transform = unsafe { sys::btTransform::new() };

        RigidBodyHandle {
            ptr,
            motion_state,
            temp_transform,
        }
    }
    pub fn set_restitution(&mut self, restitution: f64) {
        unsafe {
            sys::btCollisionObject_setRestitution(self.ptr as *mut _, restitution);
        }
    }
    pub fn set_friction(&mut self, friction: f64) {
        unsafe {
            sys::btCollisionObject_setFriction(self.ptr as *mut _, friction);
        }
    }
    pub fn set_gravity<T: Into<Vector3<f64>>>(&mut self, gravity: T) {
        let gravity: BulletVector3 = gravity.into().into();
        unsafe {
            sys::btRigidBody_setGravity(self.ptr, gravity.0.as_ptr() as *const _);
        }
    }

    pub fn set_angular_factor<T: Into<Vector3<f64>>>(&mut self, angular_factor: T) {
        let angular_factor: BulletVector3 = angular_factor.into().into();
        unsafe {
            sys::btRigidBody_setAngularFactor(self.ptr, angular_factor.0.as_ptr() as *const _);
        }
    }

    pub fn set_sleeping_thresholds(&mut self, linear: f64, angular: f64) {
        unsafe { sys::btRigidBody_setSleepingThresholds(self.ptr, linear, angular) }
    }

    pub fn apply_central_impulse<T: Into<Vector3<f64>>>(&mut self, impulse: T) {
        let impulse: BulletVector3 = impulse.into().into();
        unsafe {
            sys::btRigidBody_applyCentralImpulse(self.ptr, impulse.0.as_ptr() as *const _);
        }
    }

    pub fn get_linear_velocity(&self) -> Vector3<f64> {
        let velocity = unsafe { sys::btRigidBody_getLinearVelocity(self.ptr) };

        ::bullet_vector3::vector_from_slice(unsafe {
            ::std::slice::from_raw_parts(velocity as *const _, 4)
        })
    }

    /// Override velocity vector.
    pub fn reset_linear_velocity<T>(&mut self, vel: T)
    where
        T: Into<Vector3<f64>>,
    {
        let v = vel.into();

        unsafe { (*self.ptr).m_linearVelocity = sys::btVector3::new1(&v.x,&v.y,&v.z); };
    }

    /// Override position vector and rotation quaternion.
    pub fn reset_position_and_orientation<T, T1>(&mut self, position: T, orientation: T1)
    where
        T: Into<Vector3<f64>>,
        T1: Into<Vector4<f64>>,
    {
        let orientation: [f64; 4] = orientation.into().into();
        let position: BulletVector3 = position.into().into();
        let transform = unsafe {
            sys::btTransform::new1(
                &orientation as *const _ as *const _,
                &position as *const _ as *const _,
            )
        };

        unsafe { (*self.ptr)._base.m_worldTransform = transform };
    }

    /// Get position in world space and orientation quaternion
    pub fn get_world_position_and_orientation(&self) -> (Vector3<f64>, Vector4<f64>) {
        unsafe {
            sys::btDefaultMotionState_getWorldTransform(
                &*self.motion_state as *const _ as *mut _,
                &self.temp_transform as *const _ as *mut _,
            );
        }
        let origin = unsafe { self.temp_transform.getOrigin1().as_ref().unwrap() };

        let rotation = unsafe { self.temp_transform.getRotation() };

        (
            ::bullet_vector3::vector_from_slice(&origin.m_floats[0..3]),
            ::bullet_vector3::vector4_from_slice(&rotation._base.m_floats),
        )
    }

    /// Place data on heap as a box and set to rigid body as user pointer
    pub fn set_user_data<T: 'static>(&mut self, data: T) {
        let data_box = Box::new(data);
        unsafe {
            sys::btCollisionObject_setUserPointer(
                self.ptr as *mut _,
                Box::into_raw(data_box) as *mut _,
            )
        };
    }

    /// Get data from rigidbody's user pointer
    /// Getting data from body without previously setted data is fine.
    /// Getting wrong typed data is unsafe and will cause mem::transmute to wrong pointer type.
    pub unsafe fn get_user_data<T: 'static>(&self) -> Option<&T> {
        let pointer = sys::btCollisionObject_getUserPointer(self.ptr as *mut _);
        let pointer: *const T = ::std::mem::transmute(pointer);
        return pointer.as_ref();
    }

    /// Set user index. This will not be used in bullet and this is not related to user_data.
    pub fn set_user_index(&mut self, index: i32) {
        unsafe { sys::btCollisionObject_setUserIndex(self.ptr as *mut _, index) };
    }

    /// Get previously setted user index
    /// If index was not set - will return "-1"
    pub fn get_user_index(&self) -> i32 {
        unsafe { sys::btCollisionObject_getUserIndex(self.ptr as *mut _) }
    }

    /// Was that rigid_body removed with DynamicsWorld::remove_body()
    /// Probably may return true also for any not-added bodys
    pub fn removed(&self) -> bool {
        unsafe { (*self.ptr)._base.m_worldArrayIndex == -1 }
    }
}
<|endoftext|>


<commit_before>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::Context;
use crate::{AnyTensor, DataType, Result, Status, Tensor, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts).unwrap();
///
/// let t = Tensor::from(&[3i32]);
/// let h = TensorHandle::new(&ctx, &t).unwrap();
/// let v: Tensor<i32> = h.resolve().unwrap();
/// assert_eq!(&v[..], &[3i32]);
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts).unwrap();
///
///     let t = Tensor::from(&[3i32]);
///     let h = TensorHandle::new(&ctx, &t).unwrap();
///     h
/// };
/// ```
#[derive(Debug)]
pub struct TensorHandle<'a> {
    inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Create a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(_ctx: &'a Context, t: &Tensor<T>) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<Tensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(Tensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(&t[..], &t2[..]);
    }

    /// Following tests are disabled by default because it requires a GPU and some setup.
    ///
    /// To run this test, you need to pass the `-- --ignored` arguumnet to cargo test.
    /// ```sh
    /// cargo test --features "eager tensorflow_gpu" -- --ignored
    /// ```
    #[cfg(feature = "tensorflow_gpu")]
    mod gpu {
        use super::*;
        use crate::eager::ContextOptions;

        #[test]
        #[ignore]
        fn test_copy_to_device() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
            let h = TensorHandle::new(&ctx, &t).unwrap();
            let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&t[..], &t2[..]);
        }

        #[test]
        #[ignore]
        fn test_copy_to_device_lifetime() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let h_gpu = {
                // Create a temporal Context
                let opts = ContextOptions::new();
                let ctx2 = Context::new(opts).unwrap();
                let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

                // Create a TensorHandle managed by the context `ctx2`.
                let h = TensorHandle::new(&ctx2, &t).unwrap();

                // Copy to GPU. This creates a new handle managed by the context `ctx`.
                h.copy_to_device(&ctx, target_device).unwrap()
            };
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&values[..], &t2[..]);
        }
    }
}
<commit_msg>fix typo<commit_after>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::Context;
use crate::{AnyTensor, DataType, Result, Status, Tensor, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts).unwrap();
///
/// let t = Tensor::from(&[3i32]);
/// let h = TensorHandle::new(&ctx, &t).unwrap();
/// let v: Tensor<i32> = h.resolve().unwrap();
/// assert_eq!(&v[..], &[3i32]);
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts).unwrap();
///
///     let t = Tensor::from(&[3i32]);
///     let h = TensorHandle::new(&ctx, &t).unwrap();
///     h
/// };
/// ```
#[derive(Debug)]
pub struct TensorHandle<'a> {
    inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Create a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(_ctx: &'a Context, t: &Tensor<T>) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<Tensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(Tensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(&t[..], &t2[..]);
    }

    /// Following tests are disabled by default because it requires a GPU and some setup.
    ///
    /// To run this test, you need to pass the `-- --ignored` argument to cargo test.
    /// ```sh
    /// cargo test --features "eager tensorflow_gpu" -- --ignored
    /// ```
    #[cfg(feature = "tensorflow_gpu")]
    mod gpu {
        use super::*;
        use crate::eager::ContextOptions;

        #[test]
        #[ignore]
        fn test_copy_to_device() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
            let h = TensorHandle::new(&ctx, &t).unwrap();
            let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&t[..], &t2[..]);
        }

        #[test]
        #[ignore]
        fn test_copy_to_device_lifetime() {
            let values = [0_i32, 1, 2, 3];

            let opts = ContextOptions::new();
            let ctx = Context::new(opts).unwrap();
            let devices = ctx.device_list().unwrap();
            let gpu_device = devices
                .iter()
                .find(|d| d.device_type == "GPU")
                .expect("No GPU device was found.");
            let target_device = &gpu_device.name;

            let h_gpu = {
                // Create a temporal Context
                let opts = ContextOptions::new();
                let ctx2 = Context::new(opts).unwrap();
                let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

                // Create a TensorHandle managed by the context `ctx2`.
                let h = TensorHandle::new(&ctx2, &t).unwrap();

                // Copy to GPU. This creates a new handle managed by the context `ctx`.
                h.copy_to_device(&ctx, target_device).unwrap()
            };
            assert_eq!(&h_gpu.device_name().unwrap(), target_device);
            let t2 = h_gpu.resolve::<i32>().unwrap();

            assert_eq!(&values[..], &t2[..]);
        }
    }
}
<|endoftext|>


<commit_before>//! Heap sort algorithm.
//!
//! The `efficient` module contains the efficient sorting algorithm "Heap sort".
//!
//! Source: https://en.wikipedia.org/wiki/Heapsort

extern crate libc;

use std::slice;

/// Heap sort
///
/// # Details
///
///
/// # Scaling
///
///
///
/// # Optimizations
///
///
///
/// # Notes
///
/// The type T of the vector elements to sort _must_ implement the `PartialOrd` trait so the
/// compiler knows how to compare the elements and sort them.
///
///
///
/// # Examples
///
/// ```
/// let mut data: Vec<i32> = vec![4, 2, 3, 1, 5];
/// sorting::efficientsorts::heap::sort(&mut data);
/// assert_eq!(vec![1, 2, 3, 4, 5], data);
/// ```
///
pub fn sort<T: PartialOrd>(input: &mut [T]) {
    let n = input.len();

    heapify(input);

    let mut end = n - 1;
    while end > 0 {
        input.swap(end, 0);
        end -= 1;
        sift_down(input, 0, end);
    }
}

fn node_parent_id(i: usize) -> usize {
    if i == 0 { 0 }
    else      { (i - 1) / 2 }
}

fn node_child_left(i: usize) -> usize {
    2 * i + 1
}

// fn node_child_right(i: usize) -> usize {
//     2 * i + 2
// }

fn heapify<T: PartialOrd>(input: &mut [T]) {
    let n = input.len();
    let end = n - 1;

    // Last element is at "n-1". Find its parent:
    let mut start: usize = node_parent_id(n-1);

    loop {
        sift_down(input, start, end);
        if start == 0 { break; }
        start -= 1;
    }
}

fn sift_down<T: PartialOrd>(input: &mut [T], start: usize, end: usize) {
    let mut root = start;

    while node_child_left(root) <= end {
        let child = node_child_left(root);
        let mut swap = root;

        if input[swap] < input[child] {
            swap = child;
        }
        if child+1 <= end && input[swap] < input[child+1] {
            swap = child + 1;
        }

        if swap == root {
            return;
        } else {
            input.swap(root, swap);
            root = swap;
        }
    }
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_i8(array_pointer: *const libc::int8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i16(array_pointer: *const libc::int16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i32(array_pointer: *const libc::int32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i64(array_pointer: *const libc::int64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_u8(array_pointer: *const libc::uint8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u16(array_pointer: *const libc::uint16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u32(array_pointer: *const libc::uint32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u64(array_pointer: *const libc::uint64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_f32(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_f64(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f64, n as usize)
    };
    sort(&mut to_sort);
}
<commit_msg>Rename "root" to "i"<commit_after>//! Heap sort algorithm.
//!
//! The `efficient` module contains the efficient sorting algorithm "Heap sort".
//!
//! Source: https://en.wikipedia.org/wiki/Heapsort

extern crate libc;

use std::slice;

/// Heap sort
///
/// # Details
///
///
/// # Scaling
///
///
///
/// # Optimizations
///
///
///
/// # Notes
///
/// The type T of the vector elements to sort _must_ implement the `PartialOrd` trait so the
/// compiler knows how to compare the elements and sort them.
///
///
///
/// # Examples
///
/// ```
/// let mut data: Vec<i32> = vec![4, 2, 3, 1, 5];
/// sorting::efficientsorts::heap::sort(&mut data);
/// assert_eq!(vec![1, 2, 3, 4, 5], data);
/// ```
///
pub fn sort<T: PartialOrd>(input: &mut [T]) {
    let n = input.len();

    heapify(input);

    let mut end = n - 1;
    while end > 0 {
        input.swap(end, 0);
        end -= 1;
        sift_down(input, 0, end);
    }
}

fn node_parent_id(i: usize) -> usize {
    if i == 0 { 0 }
    else      { (i - 1) / 2 }
}

fn node_child_left(i: usize) -> usize {
    2 * i + 1
}

// fn node_child_right(i: usize) -> usize {
//     2 * i + 2
// }

fn heapify<T: PartialOrd>(input: &mut [T]) {
    let n = input.len();
    let end = n - 1;

    // Last element is at "n-1". Find its parent:
    let mut start: usize = node_parent_id(n-1);

    loop {
        sift_down(input, start, end);
        if start == 0 { break; }
        start -= 1;
    }
}

fn sift_down<T: PartialOrd>(input: &mut [T], start: usize, end: usize) {
    let mut i = start;

    while node_child_left(i) <= end {
        let child = node_child_left(i);
        let mut swap = i;

        if input[swap] < input[child] {
            swap = child;
        }
        if child+1 <= end && input[swap] < input[child+1] {
            swap = child + 1;
        }

        if swap == i {
            return;
        } else {
            input.swap(i, swap);
            i = swap;
        }
    }
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_i8(array_pointer: *const libc::int8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i16(array_pointer: *const libc::int16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i32(array_pointer: *const libc::int32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_i64(array_pointer: *const libc::int64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_u8(array_pointer: *const libc::uint8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u16(array_pointer: *const libc::uint16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u32(array_pointer: *const libc::uint32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_u64(array_pointer: *const libc::uint64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_heapsort_f32(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_heapsort_f64(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f64, n as usize)
    };
    sort(&mut to_sort);
}
<|endoftext|>


<commit_before>//! Provides an animated color value, so you can have colors change over time

use std::cmp::{Eq, Ord, PartialOrd, PartialEq, Ordering};

use linalg;
use film::Colorf;

/// ColorKeyframe is a color associated with a specific time
#[derive(Debug, Copy, Clone)]
pub struct ColorKeyframe {
    pub color: Colorf,
    pub time: f32,
}

impl ColorKeyframe {
    pub fn new(color: &Colorf, time: f32) -> ColorKeyframe {
        ColorKeyframe { color: *color, time: time }
    }
}
impl Ord for ColorKeyframe {
    fn cmp(&self, other: &ColorKeyframe) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}
impl PartialOrd for ColorKeyframe {
    fn partial_cmp(&self, other: &ColorKeyframe) -> Option<Ordering> {
        self.time.partial_cmp(&other.time)
    }
}
impl Eq for ColorKeyframe {}
impl PartialEq for ColorKeyframe {
    fn eq(&self, other: &ColorKeyframe) -> bool {
        self.time == other.time
    }
}

/// AnimatedColor is a list of colors associated with time points in the scene
/// that will compute the color at the desired time by blending the two nearest ones
#[derive(Debug, Clone)]
pub struct AnimatedColor {
    /// List of color keyframes in time order
    keyframes: Vec<ColorKeyframe>,
}

impl AnimatedColor {
    /// Create an animated transform that will blend between the passed keyframes
    pub fn with_keyframes(mut keyframes: Vec<ColorKeyframe>) -> AnimatedColor {
        keyframes.sort();
        AnimatedColor { keyframes: keyframes }
    }
    /// Compute the color at the desired time
    pub fn color(&self, time: f32) -> Colorf {
        if self.keyframes.is_empty() {
            Colorf::black()
        } else if self.keyframes.len() == 1 {
            self.keyframes[0].color
        } else {
            // TODO: Binary search here somehow? Or does the BTreeSet have some faster impl
            // of take/skip while?
            let first = self.keyframes.iter().take_while(|k| k.time < time).last();
            let second = self.keyframes.iter().skip_while(|k| k.time < time).next();
            if first.is_none() {
                self.keyframes.first().unwrap().color
            } else if second.is_none() {
                self.keyframes.last().unwrap().color
            } else {
                let mut color = Colorf::black();
                let f = first.unwrap().color;
                let s = second.unwrap().color;
                color.r = linalg::lerp(time, &f.r, &s.r);
                color.g = linalg::lerp(time, &f.g, &s.g);
                color.b = linalg::lerp(time, &f.b, &s.b);
                color.a = linalg::lerp(time, &f.a, &s.a);
                color
            }
        }
    }
}

<commit_msg>Keep color we blend normalized<commit_after>//! Provides an animated color value, so you can have colors change over time

use std::cmp::{Eq, Ord, PartialOrd, PartialEq, Ordering};

use linalg;
use film::Colorf;

/// ColorKeyframe is a color associated with a specific time
#[derive(Debug, Copy, Clone)]
pub struct ColorKeyframe {
    pub color: Colorf,
    pub time: f32,
}

impl ColorKeyframe {
    pub fn new(color: &Colorf, time: f32) -> ColorKeyframe {
        ColorKeyframe { color: *color, time: time }
    }
}
impl Ord for ColorKeyframe {
    fn cmp(&self, other: &ColorKeyframe) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}
impl PartialOrd for ColorKeyframe {
    fn partial_cmp(&self, other: &ColorKeyframe) -> Option<Ordering> {
        self.time.partial_cmp(&other.time)
    }
}
impl Eq for ColorKeyframe {}
impl PartialEq for ColorKeyframe {
    fn eq(&self, other: &ColorKeyframe) -> bool {
        self.time == other.time
    }
}

/// AnimatedColor is a list of colors associated with time points in the scene
/// that will compute the color at the desired time by blending the two nearest ones
#[derive(Debug, Clone)]
pub struct AnimatedColor {
    /// List of color keyframes in time order
    keyframes: Vec<ColorKeyframe>,
}

impl AnimatedColor {
    /// Create an animated transform that will blend between the passed keyframes
    pub fn with_keyframes(mut keyframes: Vec<ColorKeyframe>) -> AnimatedColor {
        keyframes.sort();
        AnimatedColor { keyframes: keyframes }
    }
    /// Compute the color at the desired time
    pub fn color(&self, time: f32) -> Colorf {
        if self.keyframes.is_empty() {
            Colorf::black()
        } else if self.keyframes.len() == 1 {
            self.keyframes[0].color
        } else {
            // TODO: Binary search here somehow? Or does the BTreeSet have some faster impl
            // of take/skip while?
            let first = self.keyframes.iter().take_while(|k| k.time < time).last();
            let second = self.keyframes.iter().skip_while(|k| k.time < time).next();
            if first.is_none() {
                self.keyframes.first().unwrap().color
            } else if second.is_none() {
                self.keyframes.last().unwrap().color
            } else {
                let mut color = Colorf::black();
                let f = first.unwrap().color;
                let s = second.unwrap().color;
                color.r = linalg::lerp(time, &f.r, &s.r);
                color.g = linalg::lerp(time, &f.g, &s.g);
                color.b = linalg::lerp(time, &f.b, &s.b);
                color.a = linalg::lerp(time, &f.a, &s.a);
                color.clamp()
            }
        }
    }
}

<|endoftext|>


<commit_before>use super::integration::IntegrationId;
use palette::Lch;
use std::time::Instant;

/// simple on/off devices such as relays, lights
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct OnOffDevice {
    power: bool,
}

/// lights with adjustable brightness and/or color
#[derive(Clone, Debug, PartialEq)]
pub struct Light {
    pub power: bool,
    pub brightness: Option<f64>,
    pub color: Option<Lch>,
}

/// button sensors, motion sensors
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SensorKind {
    OnOffSensor {
        value: bool,
    },
    DimmerSwitch {
        on: bool,
        up: bool,
        down: bool,
        off: bool,
    },
    Unknown,
}

#[derive(Clone, Debug, PartialEq)]
pub enum DeviceKind {
    OnOffDevice(OnOffDevice),
    Light(Light),
    Sensor(SensorKind),
}

/// active scene that's controlling the device state, if any
#[derive(Clone, Debug, PartialEq)]
pub struct DeviceSceneState {
    scene_name: String,
    activation_time: Instant,
}

/// unique identifier for the Device
type DeviceId = String;

#[derive(Clone, Debug, PartialEq)]
pub struct Device<T = DeviceKind> {
    pub id: DeviceId,
    pub name: String,
    pub integration_id: IntegrationId,
    pub scene: Option<DeviceSceneState>,
    pub kind: T,
}
<commit_msg>MultiSourceLight device<commit_after>use super::integration::IntegrationId;
use palette::Lch;
use std::time::Instant;

/// simple on/off devices such as relays, lights
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct OnOffDevice {
    power: bool,
}

/// lights with adjustable brightness and/or color
#[derive(Clone, Debug, PartialEq)]
pub struct Light {
    pub power: bool,

    /// Current brightness, if supported
    pub brightness: Option<f64>,

    /// Current color, if supported
    pub color: Option<Lch>,
}

/// lights with multiple individually adjustable light sources
#[derive(Clone, Debug, PartialEq)]
pub struct MultiSourceLight {
    pub power: bool,

    /// Global brightness control for all lights in this MultiSourceLight
    pub brightness: f64,

    /// List of colors, one for each light in this MultiSourceLight
    pub lights: Vec<Lch>,
}

/// button sensors, motion sensors
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SensorKind {
    OnOffSensor {
        value: bool,
    },
    DimmerSwitch {
        on: bool,
        up: bool,
        down: bool,
        off: bool,
    },
    Unknown,
}

#[derive(Clone, Debug, PartialEq)]
pub enum DeviceKind {
    OnOffDevice(OnOffDevice),
    Light(Light),
    MultiSourceLight(MultiSourceLight),
    Sensor(SensorKind),
}

/// active scene that's controlling the device state, if any
#[derive(Clone, Debug, PartialEq)]
pub struct DeviceSceneState {
    scene_name: String,
    activation_time: Instant,
}

/// unique identifier for the Device
type DeviceId = String;

#[derive(Clone, Debug, PartialEq)]
pub struct Device<T = DeviceKind> {
    pub id: DeviceId,
    pub name: String,
    pub integration_id: IntegrationId,
    pub scene: Option<DeviceSceneState>,
    pub kind: T,
}
<|endoftext|>