<commit_before>use common::util::*;
use std::fs::{metadata, OpenOptions, set_permissions};
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};

extern crate libc;
use self::libc::umask;

static TEST_FILE: &'static str = "file";
static REFERENCE_FILE: &'static str = "reference";
static REFERENCE_PERMS: u32 = 0o247;

struct TestCase {
    args: Vec<&'static str>,
    before: u32,
    after: u32
}

fn mkfile(file: &str, mode: u32) {
    OpenOptions::new().mode(mode).create(true).write(true).open(file).unwrap();
    let mut perms = metadata(file).unwrap().permissions();
    perms.set_mode(mode);
    set_permissions(file, perms).unwrap();
}

fn run_single_test(test: &TestCase, at: AtPath, mut ucmd: UCommand) {
        mkfile(&at.plus_as_string(TEST_FILE), test.before);
        let perms = at.metadata(TEST_FILE).permissions().mode();
        if perms != test.before {
            panic!(format!("{}: expected: {:o} got: {:o}", "setting permissions failed", test.after, perms));
        }

        for arg in &test.args {
            ucmd.arg(arg);
        }
        let r = ucmd.run();
        if !r.success {
            println!("{}", r.stderr);
            panic!(format!("{:?}: failed", ucmd.raw));
        }

        let perms = at.metadata(TEST_FILE).permissions().mode();
        if perms != test.after {
            panic!(format!("{:?}: expected: {:o} got: {:o}", ucmd.raw, test.after, perms));
        }
}

fn run_tests(tests: Vec<TestCase>) {
    for test in tests {
        let (at, ucmd) = at_and_ucmd!();
        run_single_test(&test, at, ucmd);
    }
}

#[test]
fn test_chmod_octal() {
    let tests = vec!{
        TestCase{args: vec!{"0700",  TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"0070",  TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"0007",  TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"-0700", TEST_FILE}, before: 0o700, after: 0o000},
        TestCase{args: vec!{"-0070", TEST_FILE}, before: 0o060, after: 0o000},
        TestCase{args: vec!{"-0007", TEST_FILE}, before: 0o001, after: 0o000},
        TestCase{args: vec!{"+0100", TEST_FILE}, before: 0o600, after: 0o700},
        TestCase{args: vec!{"+0020", TEST_FILE}, before: 0o050, after: 0o070},
        TestCase{args: vec!{"+0004", TEST_FILE}, before: 0o003, after: 0o007},
    };
    run_tests(tests);
}

#[test]
fn test_chmod_ugoa() {
    let last = unsafe {
        umask(0)
    };
    let tests = vec!{
        TestCase{args: vec!{"u=rwx", TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"g=rwx", TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"o=rwx", TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"a=rwx", TEST_FILE}, before: 0o000, after: 0o777},
        TestCase{args: vec!{"-r", TEST_FILE}, before: 0o777, after: 0o333},
        TestCase{args: vec!{"-w", TEST_FILE}, before: 0o777, after: 0o555},
        TestCase{args: vec!{"-x", TEST_FILE}, before: 0o777, after: 0o666},
    };
    run_tests(tests);

    unsafe {
        umask(0o022);
    }
    let tests = vec!{
        TestCase{args: vec!{"u=rwx", TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"g=rwx", TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"o=rwx", TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"a=rwx", TEST_FILE}, before: 0o000, after: 0o777},
        TestCase{args: vec!{"+rw", TEST_FILE}, before: 0o000, after: 0o644},
        TestCase{args: vec!{"=rwx", TEST_FILE}, before: 0o000, after: 0o755},
        TestCase{args: vec!{"-w", TEST_FILE}, before: 0o777, after: 0o577},
        TestCase{args: vec!{"-x", TEST_FILE}, before: 0o777, after: 0o666},
    };
    run_tests(tests);
    unsafe {
        umask(last);
    }
}

#[test]
fn test_chmod_ugo_copy() {
    let tests = vec!{
        TestCase{args: vec!{"u=g", TEST_FILE}, before: 0o070, after: 0o770},
        TestCase{args: vec!{"g=o", TEST_FILE}, before: 0o005, after: 0o055},
        TestCase{args: vec!{"o=u", TEST_FILE}, before: 0o200, after: 0o202},
        TestCase{args: vec!{"u-g", TEST_FILE}, before: 0o710, after: 0o610},
        TestCase{args: vec!{"u+g", TEST_FILE}, before: 0o250, after: 0o750},
    };
    run_tests(tests);
}

#[test]
fn test_chmod_many_options() {
    let tests = vec!{
        TestCase{args: vec!{"-r,a+w", TEST_FILE}, before: 0o444, after: 0o222},
    };
    run_tests(tests);
}

#[test]
fn test_chmod_reference_file() {
    let tests = vec!{
        TestCase{args: vec!{"--reference", REFERENCE_FILE, TEST_FILE}, before: 0o070, after: 0o247},
        TestCase{args: vec!{"a-w", "--reference", REFERENCE_FILE, TEST_FILE}, before: 0o070, after: 0o247},
    };
    let (at, ucmd) = at_and_ucmd!();
    mkfile(&at.plus_as_string(REFERENCE_FILE), REFERENCE_PERMS);
    run_single_test(&tests[0], at, ucmd);
}
<commit_msg>tests/chmod: don't make assumptions about umask<commit_after>use common::util::*;
use std::fs::{metadata, OpenOptions, set_permissions};
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};

extern crate libc;
use self::libc::umask;

static TEST_FILE: &'static str = "file";
static REFERENCE_FILE: &'static str = "reference";
static REFERENCE_PERMS: u32 = 0o247;

struct TestCase {
    args: Vec<&'static str>,
    before: u32,
    after: u32
}

fn mkfile(file: &str, mode: u32) {
    OpenOptions::new().mode(mode).create(true).write(true).open(file).unwrap();
    let mut perms = metadata(file).unwrap().permissions();
    perms.set_mode(mode);
    set_permissions(file, perms).unwrap();
}

fn run_single_test(test: &TestCase, at: AtPath, mut ucmd: UCommand) {
        mkfile(&at.plus_as_string(TEST_FILE), test.before);
        let perms = at.metadata(TEST_FILE).permissions().mode();
        if perms != test.before {
            panic!(format!("{}: expected: {:o} got: {:o}", "setting permissions failed", test.after, perms));
        }

        for arg in &test.args {
            ucmd.arg(arg);
        }
        let r = ucmd.run();
        if !r.success {
            println!("{}", r.stderr);
            panic!(format!("{:?}: failed", ucmd.raw));
        }

        let perms = at.metadata(TEST_FILE).permissions().mode();
        if perms != test.after {
            panic!(format!("{:?}: expected: {:o} got: {:o}", ucmd.raw, test.after, perms));
        }
}

fn run_tests(tests: Vec<TestCase>) {
    for test in tests {
        let (at, ucmd) = at_and_ucmd!();
        run_single_test(&test, at, ucmd);
    }
}

#[test]
fn test_chmod_octal() {
    let tests = vec!{
        TestCase{args: vec!{"0700",  TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"0070",  TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"0007",  TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"-0700", TEST_FILE}, before: 0o700, after: 0o000},
        TestCase{args: vec!{"-0070", TEST_FILE}, before: 0o060, after: 0o000},
        TestCase{args: vec!{"-0007", TEST_FILE}, before: 0o001, after: 0o000},
        TestCase{args: vec!{"+0100", TEST_FILE}, before: 0o600, after: 0o700},
        TestCase{args: vec!{"+0020", TEST_FILE}, before: 0o050, after: 0o070},
        TestCase{args: vec!{"+0004", TEST_FILE}, before: 0o003, after: 0o007},
    };
    run_tests(tests);
}

#[test]
fn test_chmod_ugoa() {
    let last = unsafe {
        umask(0)
    };
    let tests = vec!{
        TestCase{args: vec!{"u=rwx", TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"g=rwx", TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"o=rwx", TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"a=rwx", TEST_FILE}, before: 0o000, after: 0o777},
        TestCase{args: vec!{"-r", TEST_FILE}, before: 0o777, after: 0o333},
        TestCase{args: vec!{"-w", TEST_FILE}, before: 0o777, after: 0o555},
        TestCase{args: vec!{"-x", TEST_FILE}, before: 0o777, after: 0o666},
    };
    run_tests(tests);

    unsafe {
        umask(0o022);
    }
    let tests = vec!{
        TestCase{args: vec!{"u=rwx", TEST_FILE}, before: 0o000, after: 0o700},
        TestCase{args: vec!{"g=rwx", TEST_FILE}, before: 0o000, after: 0o070},
        TestCase{args: vec!{"o=rwx", TEST_FILE}, before: 0o000, after: 0o007},
        TestCase{args: vec!{"a=rwx", TEST_FILE}, before: 0o000, after: 0o777},
        TestCase{args: vec!{"+rw", TEST_FILE}, before: 0o000, after: 0o644},
        TestCase{args: vec!{"=rwx", TEST_FILE}, before: 0o000, after: 0o755},
        TestCase{args: vec!{"-w", TEST_FILE}, before: 0o777, after: 0o577},
        TestCase{args: vec!{"-x", TEST_FILE}, before: 0o777, after: 0o666},
    };
    run_tests(tests);
    unsafe {
        umask(last);
    }
}

#[test]
fn test_chmod_ugo_copy() {
    let tests = vec!{
        TestCase{args: vec!{"u=g", TEST_FILE}, before: 0o070, after: 0o770},
        TestCase{args: vec!{"g=o", TEST_FILE}, before: 0o005, after: 0o055},
        TestCase{args: vec!{"o=u", TEST_FILE}, before: 0o200, after: 0o202},
        TestCase{args: vec!{"u-g", TEST_FILE}, before: 0o710, after: 0o610},
        TestCase{args: vec!{"u+g", TEST_FILE}, before: 0o250, after: 0o750},
    };
    run_tests(tests);
}

#[test]
fn test_chmod_many_options() {
    let original_umask = unsafe {
        umask(0)
    };
    let tests = vec!{
        TestCase{args: vec!{"-r,a+w", TEST_FILE}, before: 0o444, after: 0o222},
    };
    run_tests(tests);
    unsafe {
        umask(original_umask);
    }
}

#[test]
fn test_chmod_reference_file() {
    let tests = vec!{
        TestCase{args: vec!{"--reference", REFERENCE_FILE, TEST_FILE}, before: 0o070, after: 0o247},
        TestCase{args: vec!{"a-w", "--reference", REFERENCE_FILE, TEST_FILE}, before: 0o070, after: 0o247},
    };
    let (at, ucmd) = at_and_ucmd!();
    mkfile(&at.plus_as_string(REFERENCE_FILE), REFERENCE_PERMS);
    run_single_test(&tests[0], at, ucmd);
}
<|endoftext|>


<commit_before>extern crate pombase;

use std::collections::hash_set::HashSet;
use std::iter::Iterator;

use self::pombase::api::query::*;
use self::pombase::api::result::*;
use self::pombase::api::server_data::*;
use self::pombase::api::query_exec::*;

fn get_server_data() -> ServerData {
    use std::path::PathBuf;
    let mut search_maps_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    search_maps_path.push("tests/test_search_data.json");
    let mut gene_subsets_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    gene_subsets_path.push("tests/test_gene_subsets.json");
    let mut config_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    config_path.push("tests/test_config.json");
    ServerData::new(config_path.to_str().unwrap(), search_maps_path.to_str().unwrap(),
                    gene_subsets_path.to_str().unwrap())
}

fn check_gene_result(query: &Query, genes: Vec<&str>) {
    let server_data = get_server_data();
    let query_exec = QueryExec::new(server_data);
    let result = query_exec.exec(&query);

    let result_genes_iter =
        result.rows.into_iter()
        .map(|row: ResultRow| row.gene_uniquename)
        .collect::<HashSet<_>>();
    let expect_genes_iter =
        genes.iter().cloned()
        .map(|s: &str| String::from(s))
        .collect::<HashSet<_>>();
    assert_eq!(result_genes_iter,
               expect_genes_iter);
}

#[test]
fn test_and_or_not() {
    let qp1 = QueryNode::GeneList { ids: vec!["id_one".into(), "id_two".into(), "id_three".into()] };
    let qp2 = QueryNode::GeneList { ids: vec!["id_two".into(), "id_three".into(), "id_four".into()] };
    let qp3 = QueryNode::GeneList { ids: vec!["id_one".into(), "id_two".into(),
                                              "id_three".into(), "id_four".into()] };
    let qp4 = QueryNode::GeneList { ids: vec!["id_two".into(), "id_three".into(),
                                              "id_four".into(), "id_five".into()] };
    let qp5 = QueryNode::GeneList { ids: vec!["id_three".into(), "id_four".into(),
                                              "id_five".into()] };

    let and_query_node_1 = QueryNode::And(vec![qp1.clone(), qp2.clone()]);
    let and_query_node_2 = QueryNode::And(vec![qp3.clone(), qp4.clone()]);

    let and_query =
        Query::from_node(QueryNode::And(vec![and_query_node_1.clone(), and_query_node_2.clone()]));

    check_gene_result(&and_query, vec!["id_two", "id_three"]);

    let or_query =
        Query::from_node(QueryNode::Or(vec![and_query_node_1, and_query_node_2]));

    check_gene_result(&or_query, vec!["id_two", "id_three", "id_four"]);

    let not_query_node =
        QueryNode::Not { node_a: Box::new(qp1.clone()), node_b: Box::new(qp2.clone()) };
    let not_query = Query::from_node(not_query_node);
    check_gene_result(&not_query, vec!["id_one"]);

    let not_query_2_node = QueryNode::Not { node_a: Box::new(qp1), node_b: Box::new(qp5) };
    let not_query_2 = Query::from_node(not_query_2_node);
    check_gene_result(&not_query_2, vec!["id_one", "id_two"]);
}

#[test]
fn test_termid() {
    let qp1 = QueryNode::TermId { termid: "GO:0044237".into(), definition: None,
                                  single_or_multi_allele: None, expression: None };
    let q1 = Query::from_node(qp1);

    check_gene_result(&q1, vec!["SPAC24B11.06c", "SPAC19G12.03", "SPAC2F3.09", "SPAC27E2.05"]);
}

#[test]
fn test_gene_subset() {
    let qp1 = QueryNode::Subset { subset_name: "PTHR17490:SF9".into() };
    let q1 = Query::from_node(qp1);

    check_gene_result(&q1, vec!["SPCC895.03c"]);
}
<commit_msg>Fix broken test<commit_after>extern crate pombase;

use std::collections::hash_set::HashSet;
use std::iter::Iterator;

use self::pombase::api::query::*;
use self::pombase::api::result::*;
use self::pombase::api::server_data::*;
use self::pombase::api::query_exec::*;

fn get_server_data() -> ServerData {
    use std::path::PathBuf;
    let mut search_maps_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    search_maps_path.push("tests/test_search_data.json");
    let mut gene_subsets_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    gene_subsets_path.push("tests/test_gene_subsets.json");
    let mut config_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    config_path.push("tests/test_config.json");
    ServerData::new(config_path.to_str().unwrap(), search_maps_path.to_str().unwrap(),
                    gene_subsets_path.to_str().unwrap())
}

fn check_gene_result(query: &Query, genes: Vec<&str>) {
    let server_data = get_server_data();
    let query_exec = QueryExec::new(server_data);
    let result = query_exec.exec(&query);

    let result_genes_iter =
        result.rows.into_iter()
        .map(|row: ResultRow| row.gene_uniquename)
        .collect::<HashSet<_>>();
    let expect_genes_iter =
        genes.iter().cloned()
        .map(|s: &str| String::from(s))
        .collect::<HashSet<_>>();
    assert_eq!(result_genes_iter,
               expect_genes_iter);
}

#[test]
fn test_and_or_not() {
    let qp1 = QueryNode::GeneList { ids: vec!["id_one".into(), "id_two".into(), "id_three".into()] };
    let qp2 = QueryNode::GeneList { ids: vec!["id_two".into(), "id_three".into(), "id_four".into()] };
    let qp3 = QueryNode::GeneList { ids: vec!["id_one".into(), "id_two".into(),
                                              "id_three".into(), "id_four".into()] };
    let qp4 = QueryNode::GeneList { ids: vec!["id_two".into(), "id_three".into(),
                                              "id_four".into(), "id_five".into()] };
    let qp5 = QueryNode::GeneList { ids: vec!["id_three".into(), "id_four".into(),
                                              "id_five".into()] };

    let and_query_node_1 = QueryNode::And(vec![qp1.clone(), qp2.clone()]);
    let and_query_node_2 = QueryNode::And(vec![qp3.clone(), qp4.clone()]);

    let and_query =
        Query::from_node(QueryNode::And(vec![and_query_node_1.clone(), and_query_node_2.clone()]));

    check_gene_result(&and_query, vec!["id_two", "id_three"]);

    let or_query =
        Query::from_node(QueryNode::Or(vec![and_query_node_1, and_query_node_2]));

    check_gene_result(&or_query, vec!["id_two", "id_three", "id_four"]);

    let not_query_node =
        QueryNode::Not { node_a: Box::new(qp1.clone()), node_b: Box::new(qp2.clone()) };
    let not_query = Query::from_node(not_query_node);
    check_gene_result(&not_query, vec!["id_one"]);

    let not_query_2_node = QueryNode::Not { node_a: Box::new(qp1), node_b: Box::new(qp5) };
    let not_query_2 = Query::from_node(not_query_2_node);
    check_gene_result(&not_query_2, vec!["id_one", "id_two"]);
}

#[test]
fn test_termid() {
    let qp1 = QueryNode::Term { termid: "GO:0044237".into(), name: None,
                                single_or_multi_allele: None, expression: None };
    let q1 = Query::from_node(qp1);

    check_gene_result(&q1, vec!["SPAC24B11.06c", "SPAC19G12.03", "SPAC2F3.09", "SPAC27E2.05"]);
}

#[test]
fn test_gene_subset() {
    let qp1 = QueryNode::Subset { subset_name: "PTHR17490:SF9".into() };
    let q1 = Query::from_node(qp1);

    check_gene_result(&q1, vec!["SPCC895.03c"]);
}
<|endoftext|>


<commit_before>extern crate discord;
extern crate dotenv;
#[macro_use]
extern crate serde_derive;
extern crate envy;

use discord::{Discord, State};
use discord::model::{Event};
use discord::voice::AudioSource;

use std::collections::HashMap;
use std::sync::mpsc::{channel, Sender, Receiver};

#[derive(Default)]
struct Dj {
  controls: HashMap<u64, Sender<AudioSourceControl>>
}

pub enum MusicItem {
  YouTube(String)
}

fn main() {
  println!("Loading .env");

  dotenv::dotenv().ok();

  println!("Reading environment variables");

  let environment: Environment = envy::from_env().expect("Invalid or missing environment variables");

  println!("Creating bot");

  let discord = Discord::from_bot_token(&environment.discord_bot_token).expect("could not start discord from bot token");
  let (mut connection, ready) = discord.connect().expect("could not connect");
  let mut state = State::new(ready);

  let mut dj = Dj::default();

  loop {
    let event = match connection.recv_event() {
      Ok(e) => e,
      Err(e) => {
        println!("could not receive event: {}", e);
        continue;
      }
    };

    state.update(&event);

    let valid_commands = &["!play", "!stop", "!clear", "!pause", "!resume"];

    match event {
      Event::MessageCreate(message) => {
        let params: Vec<&str> = message.content.split_whitespace().collect();
        if params.is_empty() || params[0].len() < 2 {
          continue;
        }

        let command: String = params[0].to_lowercase().chars().skip(1).collect();
        if !params[0].starts_with('!') || !valid_commands.contains(&command.as_ref()) {
          continue;
        }

        let (server_id, voice_channel) = match state.find_voice_user(message.author.id) {
          Some((Some(server_id), voice_channel)) => (server_id, voice_channel),
          _ => {
            if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Must be in a public voice channel to DJ.")) {
              println!("could not send message: {}", e);
            }
            continue;
          }
        };

        let voice = connection.voice(Some(server_id));

        match command.as_ref() {
          "play" => {
            if params.len() < 2 {
              if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Please provide a YouTube URL to play.")) {
                println!("could not send message: {}", e);
              }
              continue;
            }
            let url = params[1];
            if let Some((_, controls)) = dj.controls.iter().find(|&(id, _)| *id == voice_channel.0) {
              if let Err(e) = controls.send(AudioSourceControl::AddQueue(url.to_string())) {
                println!("could not send control: {}", e);
              }
              continue;
            }
            let source = match discord::voice::open_ytdl_stream(url) {
              Ok(s) => s,
              Err(e) => {
                println!("could not open source: {}", e);
                if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Could not open video for streaming.")) {
                  println!("could not send message: {}", e);
                }
                continue;
              }
            };

            voice.connect(voice_channel);

            let (tx, rx) = channel();
            dj.controls.insert(voice_channel.0, tx);
            let buf_source = BufferedQueuedAudioSource::new(source, rx);
            voice.play(Box::new(buf_source));
          },
          "stop" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Stop) {
              println!("couldn't send control: {}", e);
            }
          },
          "pause" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Pause) {
              println!("couldn't send control: {}", e);
            }
          },
          "resume" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Resume) {
              println!("couldn't send control: {}", e);
            }
          },
          "clear" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::ClearQueue) {
              println!("couldn't send control: {}", e);
            }
          },
          _ => continue
        }
      },
      Event::VoiceStateUpdate(server_id, _) => {
        let chan = match connection.voice(server_id).current_channel() {
          Some(c) => c,
          None => continue
        };
        if let Some(server_id) = server_id {
          if let Some(srv) = state.servers().iter().find(|srv| srv.id == server_id) {
            if srv.voice_states.iter().filter(|vs| vs.channel_id == Some(chan)).count() <= 1 {
              connection.voice(Some(server_id)).disconnect();
            }
          }
        }
      },
      _ => {}
    }
  }
}

#[derive(Debug, Deserialize)]
pub struct Environment {
  #[serde(rename = "lb_discord_bot_token")]
  pub discord_bot_token: String,
}

pub enum AudioSourceControl {
  Pause,
  Resume,
  Stop,
  AddQueue(String),
  ClearQueue
}

struct BufferedQueuedAudioSource {
  controls: Receiver<AudioSourceControl>,
  source: Box<AudioSource>,
  buffer: Box<[i16]>,
  queue: Vec<String>,
  position: usize,
  capacity: usize,
  paused: bool,
  finished: bool
}

impl BufferedQueuedAudioSource {
  fn new(source: Box<AudioSource>, controls: Receiver<AudioSourceControl>) -> Self {
    BufferedQueuedAudioSource {
      controls: controls,
      source: source,
      buffer: vec![0; 61440].into_boxed_slice(),
      queue: Vec::default(),
      position: 0,
      capacity: 0,
      paused: false,
      finished: false
    }
  }
}

impl AudioSource for BufferedQueuedAudioSource {
  fn is_stereo(&mut self) -> bool {
    self.source.is_stereo()
  }

  fn read_frame(&mut self, buffer: &mut [i16]) -> Option<usize> {
    self.read_controls();
    self.handle_queue();
    if self.paused {
      for byte in buffer.iter_mut() {
        *byte = 0;
      }
      return Some(buffer.len());
    }
    self.read(buffer)
  }
}

impl BufferedQueuedAudioSource {
  fn read_controls(&mut self) {
    while let Ok(control) = self.controls.try_recv() {
      match control {
        AudioSourceControl::Resume => self.paused = false,
        AudioSourceControl::Pause => self.paused = true,
        AudioSourceControl::AddQueue(s) => self.queue.push(s),
        AudioSourceControl::ClearQueue => self.queue.clear(),
        AudioSourceControl::Stop => {
          self.queue.clear();
          self.finished = true;
        }
      }
    }
  }

  fn handle_queue(&mut self) {
    if !self.finished || self.queue.is_empty() {
      return;
    }
    while !self.queue.is_empty() {
      let first = self.queue.remove(0);
      match discord::voice::open_ytdl_stream(&first) {
        Ok(s) => {
          self.source = s;
          self.finished = false;
          break;
        },
        Err(e) => {
          println!("could not open ytdl stream in queue: {}", e);
          continue;
        }
      }
    }
  }

  fn read(&mut self, buffer: &mut [i16]) -> Option<usize> {
    if self.position == self.capacity && buffer.len() >= self.buffer.len() {
      return self.source.read_frame(buffer);
    }
    let nread = {
      let rem = match self.fill_buff() {
        Some(s) => s,
        None => return None
      };
      let read = std::cmp::min(buffer.len(), rem.len());
      for i in 0..read {
        buffer[i] = rem[i];
      }
      read
    };
    self.consume(nread);
    Some(nread)
  }

  fn fill_buff(&mut self) -> Option<&[i16]> {
    if self.position >= self.capacity {
      self.capacity = match self.source.read_frame(&mut self.buffer) {
        Some(s) => s,
        None => return None
      };
      // FIXME: This should happen when read_frame turns None, but discord-rs has a bug
      // https://github.com/SpaceManiac/discord-rs/issues/107
      if self.capacity == 0 && !self.finished {
        self.finished = true;
      }
      self.position = 0;
    }
    Some(&self.buffer[self.position..self.capacity])
  }

  fn consume(&mut self, amt: usize) {
    self.position = std::cmp::min(self.position + amt, self.capacity);
  }
}
<commit_msg>fix: remove redundant check<commit_after>extern crate discord;
extern crate dotenv;
#[macro_use]
extern crate serde_derive;
extern crate envy;

use discord::{Discord, State};
use discord::model::{Event};
use discord::voice::AudioSource;

use std::collections::HashMap;
use std::sync::mpsc::{channel, Sender, Receiver};

#[derive(Default)]
struct Dj {
  controls: HashMap<u64, Sender<AudioSourceControl>>
}

pub enum MusicItem {
  YouTube(String)
}

fn main() {
  println!("Loading .env");

  dotenv::dotenv().ok();

  println!("Reading environment variables");

  let environment: Environment = envy::from_env().expect("Invalid or missing environment variables");

  println!("Creating bot");

  let discord = Discord::from_bot_token(&environment.discord_bot_token).expect("could not start discord from bot token");
  let (mut connection, ready) = discord.connect().expect("could not connect");
  let mut state = State::new(ready);

  let mut dj = Dj::default();

  loop {
    let event = match connection.recv_event() {
      Ok(e) => e,
      Err(e) => {
        println!("could not receive event: {}", e);
        continue;
      }
    };

    state.update(&event);

    let valid_commands = &["!play", "!stop", "!clear", "!pause", "!resume"];

    match event {
      Event::MessageCreate(message) => {
        let params: Vec<&str> = message.content.split_whitespace().collect();
        if params.is_empty() || params[0].len() < 2 {
          continue;
        }

        let command: String = params[0].to_lowercase().chars().skip(1).collect();
        if !valid_commands.contains(&command.as_ref()) {
          continue;
        }

        let (server_id, voice_channel) = match state.find_voice_user(message.author.id) {
          Some((Some(server_id), voice_channel)) => (server_id, voice_channel),
          _ => {
            if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Must be in a public voice channel to DJ.")) {
              println!("could not send message: {}", e);
            }
            continue;
          }
        };

        let voice = connection.voice(Some(server_id));

        match command.as_ref() {
          "play" => {
            if params.len() < 2 {
              if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Please provide a YouTube URL to play.")) {
                println!("could not send message: {}", e);
              }
              continue;
            }
            let url = params[1];
            if let Some((_, controls)) = dj.controls.iter().find(|&(id, _)| *id == voice_channel.0) {
              if let Err(e) = controls.send(AudioSourceControl::AddQueue(url.to_string())) {
                println!("could not send control: {}", e);
              }
              continue;
            }
            let source = match discord::voice::open_ytdl_stream(url) {
              Ok(s) => s,
              Err(e) => {
                println!("could not open source: {}", e);
                if let Err(e) = discord.send_embed(message.channel_id, "", |e| e.description("Could not open video for streaming.")) {
                  println!("could not send message: {}", e);
                }
                continue;
              }
            };

            voice.connect(voice_channel);

            let (tx, rx) = channel();
            dj.controls.insert(voice_channel.0, tx);
            let buf_source = BufferedQueuedAudioSource::new(source, rx);
            voice.play(Box::new(buf_source));
          },
          "stop" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Stop) {
              println!("couldn't send control: {}", e);
            }
          },
          "pause" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Pause) {
              println!("couldn't send control: {}", e);
            }
          },
          "resume" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::Resume) {
              println!("couldn't send control: {}", e);
            }
          },
          "clear" => {
            let control = match dj.controls.get(&voice_channel.0) {
              Some(c) => c,
              None => continue
            };
            if let Err(e) = control.send(AudioSourceControl::ClearQueue) {
              println!("couldn't send control: {}", e);
            }
          },
          _ => continue
        }
      },
      Event::VoiceStateUpdate(server_id, _) => {
        let chan = match connection.voice(server_id).current_channel() {
          Some(c) => c,
          None => continue
        };
        if let Some(server_id) = server_id {
          if let Some(srv) = state.servers().iter().find(|srv| srv.id == server_id) {
            if srv.voice_states.iter().filter(|vs| vs.channel_id == Some(chan)).count() <= 1 {
              connection.voice(Some(server_id)).disconnect();
            }
          }
        }
      },
      _ => {}
    }
  }
}

#[derive(Debug, Deserialize)]
pub struct Environment {
  #[serde(rename = "lb_discord_bot_token")]
  pub discord_bot_token: String,
}

pub enum AudioSourceControl {
  Pause,
  Resume,
  Stop,
  AddQueue(String),
  ClearQueue
}

struct BufferedQueuedAudioSource {
  controls: Receiver<AudioSourceControl>,
  source: Box<AudioSource>,
  buffer: Box<[i16]>,
  queue: Vec<String>,
  position: usize,
  capacity: usize,
  paused: bool,
  finished: bool
}

impl BufferedQueuedAudioSource {
  fn new(source: Box<AudioSource>, controls: Receiver<AudioSourceControl>) -> Self {
    BufferedQueuedAudioSource {
      controls: controls,
      source: source,
      buffer: vec![0; 61440].into_boxed_slice(),
      queue: Vec::default(),
      position: 0,
      capacity: 0,
      paused: false,
      finished: false
    }
  }
}

impl AudioSource for BufferedQueuedAudioSource {
  fn is_stereo(&mut self) -> bool {
    self.source.is_stereo()
  }

  fn read_frame(&mut self, buffer: &mut [i16]) -> Option<usize> {
    self.read_controls();
    self.handle_queue();
    if self.paused {
      for byte in buffer.iter_mut() {
        *byte = 0;
      }
      return Some(buffer.len());
    }
    self.read(buffer)
  }
}

impl BufferedQueuedAudioSource {
  fn read_controls(&mut self) {
    while let Ok(control) = self.controls.try_recv() {
      match control {
        AudioSourceControl::Resume => self.paused = false,
        AudioSourceControl::Pause => self.paused = true,
        AudioSourceControl::AddQueue(s) => self.queue.push(s),
        AudioSourceControl::ClearQueue => self.queue.clear(),
        AudioSourceControl::Stop => {
          self.queue.clear();
          self.finished = true;
        }
      }
    }
  }

  fn handle_queue(&mut self) {
    if !self.finished || self.queue.is_empty() {
      return;
    }
    while !self.queue.is_empty() {
      let first = self.queue.remove(0);
      match discord::voice::open_ytdl_stream(&first) {
        Ok(s) => {
          self.source = s;
          self.finished = false;
          break;
        },
        Err(e) => {
          println!("could not open ytdl stream in queue: {}", e);
          continue;
        }
      }
    }
  }

  fn read(&mut self, buffer: &mut [i16]) -> Option<usize> {
    if self.position == self.capacity && buffer.len() >= self.buffer.len() {
      return self.source.read_frame(buffer);
    }
    let nread = {
      let rem = match self.fill_buff() {
        Some(s) => s,
        None => return None
      };
      let read = std::cmp::min(buffer.len(), rem.len());
      for i in 0..read {
        buffer[i] = rem[i];
      }
      read
    };
    self.consume(nread);
    Some(nread)
  }

  fn fill_buff(&mut self) -> Option<&[i16]> {
    if self.position >= self.capacity {
      self.capacity = match self.source.read_frame(&mut self.buffer) {
        Some(s) => s,
        None => return None
      };
      // FIXME: This should happen when read_frame turns None, but discord-rs has a bug
      // https://github.com/SpaceManiac/discord-rs/issues/107
      if self.capacity == 0 && !self.finished {
        self.finished = true;
      }
      self.position = 0;
    }
    Some(&self.buffer[self.position..self.capacity])
  }

  fn consume(&mut self, amt: usize) {
    self.position = std::cmp::min(self.position + amt, self.capacity);
  }
}
<|endoftext|>


<commit_before>use irc;
use matrix;
use irc::protocol::{Command,Message};
use mio;
use mio::{EventLoop,Handler,Token,EventSet,PollOpt,Sender};
use std::thread;
use std::collections::HashMap;

const CLIENT: Token = Token(0);
const MATRIX: Token = Token(1);

#[derive(Debug)]
pub enum Event {
    EndPoll,
    Matrix(matrix::events::Event)
}

pub struct Bridge {
    client: irc::streams::Client,
    matrix: matrix::client::Client,
    rooms: HashMap<matrix::events::RoomID, Room>,
    seenEvents: Vec<matrix::events::EventID>,
    server_name: String,
}

impl Handler for Bridge {
    type Timeout = ();
    type Message = Event;

    fn ready(&mut self, event_loop: &mut EventLoop<Bridge>, token: Token, _: EventSet) {
        match token {
            CLIENT =>
                self.handle_client(event_loop),
            _ => unreachable!()
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Bridge>, msg: Self::Message) {
        println!("Got message from matrix! {:?}", msg);
        match msg {
            Event::EndPoll => {
                self.poll_matrix(event_loop.channel());
            },
            Event::Matrix(e) =>
                self.handle_matrix(e)
        };
    }
}

unsafe impl Sync for Bridge{}

struct Room {
    id: matrix::events::RoomID,
    canonical_alias: Option<String>,
    join_rules: Option<String>,
    members: Vec<matrix::events::UserID>,
    pending_events: Vec<matrix::events::RoomEvent>
}

impl Room {
    fn new(id: matrix::events::RoomID) -> Self {
        Room {
            id: id,
            canonical_alias: None,
            join_rules: None,
            members: vec![],
            pending_events: vec![]
        }
    }

    fn run_pending<F>(&mut self, mut callback: &mut F)
            where F: FnMut(irc::protocol::Message) {
        while let Some(evt) = self.pending_events.pop() {
            self.handle_with_alias(evt, callback);
        }
    }

    fn handle_with_alias<F>(&mut self, evt: matrix::events::RoomEvent, mut callback: &mut F)
            where F: FnMut(irc::protocol::Message) {
        if self.canonical_alias != None {
            match evt {
                matrix::events::RoomEvent::Membership(user, matrix::events::MembershipAction::Join) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Join,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: None
                    });
                    self.members.push(user);
                },
                matrix::events::RoomEvent::Membership(user, matrix::events::MembershipAction::Leave) => {
                    if self.canonical_alias != None {
                        callback(irc::protocol::Message {
                            prefix: Some(format!("{}@anony.oob", user.nickname)),
                            command: irc::protocol::Command::Part,
                            args: vec![self.canonical_alias.clone().unwrap()],
                            suffix: None
                        });
                    }
                    self.members.push(user);
                },
                matrix::events::RoomEvent::Membership(user, _) => (),
                matrix::events::RoomEvent::Message(user, text) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Privmsg,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: Some(text)
                    });
                },
                matrix::events::RoomEvent::Topic(user, topic) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Topic,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: Some(topic.clone())
                    });
                },
                matrix::events::RoomEvent::CanonicalAlias(_) => unreachable!(),
                _ => {
                    println!("Unhandled event {:?}", evt)
                }
            }
        } else {
            self.pending_events.push(evt);
        }
    }

    fn handle_event<F>(&mut self, evt: matrix::events::RoomEvent, mut callback: F)
            where F: FnMut(irc::protocol::Message) {
        match evt {
            matrix::events::RoomEvent::CanonicalAlias(name) => {
                let was_empty = self.canonical_alias == None;
                self.canonical_alias = Some(name.clone());
                for ref user in &self.members {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Join,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: None
                    });
                };
                if was_empty {
                    self.run_pending(&mut callback);
                }
            },
            matrix::events::RoomEvent::JoinRules(rules) =>
                self.join_rules = Some(rules.clone()),
            matrix::events::RoomEvent::Create => (),
            matrix::events::RoomEvent::Aliases => (),
            matrix::events::RoomEvent::PowerLevels => (),
            matrix::events::RoomEvent::HistoryVisibility(_) => (),
            matrix::events::RoomEvent::Name(_, _) => (),
            matrix::events::RoomEvent::Avatar(_, _) => (),
            matrix::events::RoomEvent::Unknown(json) => {
                println!("Unknown room event: {:?}", json);
            }
            _ => self.handle_with_alias(evt, &mut callback)
        };
    }
}

impl Bridge {
    pub fn room(&mut self, id: &matrix::events::RoomID) -> &mut Room {
        if !self.rooms.contains_key(id) {
            self.rooms.insert(id.clone(), Room::new(id.clone()));
        }
        match self.rooms.get_mut(id) {
            Some(room) => room,
            None => unreachable!()
        }
    }

    pub fn new(client: irc::streams::Client, url: &str) -> Self {
        Bridge {
            client: client,
            matrix: matrix::client::Client::new(url),
            rooms: HashMap::new(),
            seenEvents: vec![],
            server_name: "oob.systems".to_string()
        }
    }

    pub fn run(&mut self) {
        let mut events = EventLoop::new().unwrap();
        events.register(self.client.stream(), CLIENT, EventSet::all(), PollOpt::edge()).unwrap();
        events.run(self).unwrap();
    }

    fn handle_matrix(&mut self, evt: matrix::events::Event) {
        let duplicate = match evt.id {
            Some(ref id) =>
                self.seenEvents.contains(id),
            _ => false
        };
        if !duplicate {
            let mut messages: Vec<irc::protocol::Message> = vec![];
            {
                let appendMsg = |msg: irc::protocol::Message| {
                    messages.push(msg);
                };
                match evt.data {
                    matrix::events::EventData::Room(room_id, room_event) => {
                        self.room(&room_id).handle_event(room_event, appendMsg);
                    },
                    _ => println!("Unhandled {:?}", evt)
                }
            }
            match evt.id {
                Some(id) =>
                    self.seenEvents.push(id),
                None => ()
            };
            for ref msg in messages {
                self.client.send(msg);
            }
        }
    }

    fn poll_matrix(&mut self, channel: mio::Sender<Event>) ->
        thread::JoinHandle<matrix::client::Result> {
        let poll = self.matrix.pollAsync();
        thread::spawn(move|| {
            poll.send().and_then(|evts| {
                for evt in evts {
                    channel.send(Event::Matrix(evt));
                };
                channel.send(Event::EndPoll);
                Ok(())
            })
        })
    }

    fn start_matrix(&mut self, channel: mio::Sender<Event>) ->
        matrix::client::Result {
        self.matrix.sync(|evt: matrix::events::Event| {
            channel.send(Event::Matrix(evt));
        }).and_then(|r| {
            self.poll_matrix(channel);
            Ok(r)
        })
    }

    fn handle_client(&mut self, events: &mut EventLoop<Bridge>) {
        loop {
            match self.client.read_message() {
                None => return,
                Some(message) => {
                    println!("Got a message! {:?}", message);
                    match message.command {
                        Command::Pass => {
                            self.client.auth.set_password(message.args[0].clone())
                        }
                        Command::Nick => {
                            self.client.set_nickname(message.args[0].clone());
                        },
                        Command::User => {
                            println!("User logged in: {}", message.args[0]);
                            self.client.auth.set_username(message.args[0].clone());
                            let auth = self.client.auth.consume();
                            match (auth.username, auth.password) {
                                (Some(username), Some(password)) => {
                                    self.matrix.login(username.trim(), password.trim())
                                        .and(self.start_matrix(events.channel()))
                                        .and_then(|_| {
                                            self.client.welcome(username.trim());
                                            println!("Logged in {:?}", username);
                                            Ok(())
                                        }).expect("Could not login!");
                                },
                                _ => panic!("Username and/or password missing")
                            };
                        },
                        Command::Join => {
                            self.client.join(&message.args[0]);
                        },
                        Command::Ping => {
                            self.client.pong();
                        },
                        Command::Quit => {
                            // FIXME: Logout of matrix and exit thread
                            return;
                        },
                        Command::Privmsg => {
                            let uid = matrix::events::UserID::from_str("@tdfischer:localhost");
                            let roomid = matrix::events::RoomID::from_str("!SNCDinqFeGteFrlCoN:localhost");
                            let evt = matrix::events::EventData::Room(
                                roomid,
                                matrix::events::RoomEvent::Message(
                                    uid, message.suffix.unwrap()));
                            self.seenEvents.push(self.matrix.send(evt));
                        },
                        _ =>
                            println!("unhandled {:?}", message)
                    }
                }
            }
        }
    }
}

<commit_msg>bridge: map irc channels to matrix channels, and handle PRIVMSG<commit_after>use irc;
use matrix;
use irc::protocol::{Command,Message};
use mio;
use mio::{EventLoop,Handler,Token,EventSet,PollOpt,Sender};
use std::thread;
use std::collections::HashMap;

const CLIENT: Token = Token(0);
const MATRIX: Token = Token(1);

#[derive(Debug)]
pub enum Event {
    EndPoll,
    Matrix(matrix::events::Event)
}

pub struct Bridge {
    client: irc::streams::Client,
    matrix: matrix::client::Client,
    rooms: HashMap<matrix::events::RoomID, Room>,
    seenEvents: Vec<matrix::events::EventID>,
    server_name: String,
}

impl Handler for Bridge {
    type Timeout = ();
    type Message = Event;

    fn ready(&mut self, event_loop: &mut EventLoop<Bridge>, token: Token, _: EventSet) {
        match token {
            CLIENT =>
                self.handle_client(event_loop),
            _ => unreachable!()
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Bridge>, msg: Self::Message) {
        println!("Got message from matrix! {:?}", msg);
        match msg {
            Event::EndPoll => {
                self.poll_matrix(event_loop.channel());
            },
            Event::Matrix(e) =>
                self.handle_matrix(e)
        };
    }
}

unsafe impl Sync for Bridge{}

struct Room {
    id: matrix::events::RoomID,
    canonical_alias: Option<String>,
    join_rules: Option<String>,
    members: Vec<matrix::events::UserID>,
    pending_events: Vec<matrix::events::RoomEvent>
}

impl Room {
    fn new(id: matrix::events::RoomID) -> Self {
        Room {
            id: id,
            canonical_alias: None,
            join_rules: None,
            members: vec![],
            pending_events: vec![]
        }
    }

    fn run_pending<F>(&mut self, mut callback: &mut F)
            where F: FnMut(irc::protocol::Message) {
        while let Some(evt) = self.pending_events.pop() {
            self.handle_with_alias(evt, callback);
        }
    }

    fn handle_with_alias<F>(&mut self, evt: matrix::events::RoomEvent, mut callback: &mut F)
            where F: FnMut(irc::protocol::Message) {
        if self.canonical_alias != None {
            match evt {
                matrix::events::RoomEvent::Membership(user, matrix::events::MembershipAction::Join) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Join,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: None
                    });
                    self.members.push(user);
                },
                matrix::events::RoomEvent::Membership(user, matrix::events::MembershipAction::Leave) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Part,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: None
                    });
                    self.members.push(user);
                },
                matrix::events::RoomEvent::Membership(user, _) => (),
                matrix::events::RoomEvent::Message(user, text) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Privmsg,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: Some(text)
                    });
                },
                matrix::events::RoomEvent::Topic(user, topic) => {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Topic,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: Some(topic.clone())
                    });
                },
                matrix::events::RoomEvent::CanonicalAlias(_) => unreachable!(),
                _ => {
                    println!("Unhandled event {:?}", evt)
                }
            }
        } else {
            self.pending_events.push(evt);
        }
    }

    fn handle_event<F>(&mut self, evt: matrix::events::RoomEvent, mut callback: F)
            where F: FnMut(irc::protocol::Message) {
        match evt {
            matrix::events::RoomEvent::CanonicalAlias(name) => {
                let was_empty = self.canonical_alias == None;
                self.canonical_alias = Some(name.clone());
                for ref user in &self.members {
                    callback(irc::protocol::Message {
                        prefix: Some(format!("{}@anony.oob", user.nickname)),
                        command: irc::protocol::Command::Join,
                        args: vec![self.canonical_alias.clone().unwrap()],
                        suffix: None
                    });
                };
                if was_empty {
                    self.run_pending(&mut callback);
                }
            },
            matrix::events::RoomEvent::JoinRules(rules) =>
                self.join_rules = Some(rules.clone()),
            matrix::events::RoomEvent::Create => (),
            matrix::events::RoomEvent::Aliases => (),
            matrix::events::RoomEvent::PowerLevels => (),
            matrix::events::RoomEvent::HistoryVisibility(_) => (),
            matrix::events::RoomEvent::Name(_, _) => (),
            matrix::events::RoomEvent::Avatar(_, _) => (),
            matrix::events::RoomEvent::Unknown(json) => {
                println!("Unknown room event: {:?}", json);
            }
            _ => self.handle_with_alias(evt, &mut callback)
        };
    }
}

impl Bridge {
    pub fn room_from_matrix(&mut self, id: &matrix::events::RoomID) -> &mut Room {
        if !self.rooms.contains_key(id) {
            self.rooms.insert(id.clone(), Room::new(id.clone()));
        }
        match self.rooms.get_mut(id) {
            Some(room) => room,
            None => unreachable!()
        }
    }

    pub fn room_from_irc(&mut self, id: &String) -> &mut Room {
        let mut room_id: Option<matrix::events::RoomID> = None;
        for (_, r) in self.rooms.iter_mut() {
            if let Some(ref alias) = r.canonical_alias {
                if alias == id {
                    room_id = Some(r.id.clone())
                }
            }
        }
        self.room_from_matrix(&room_id.unwrap())
    }

    pub fn new(client: irc::streams::Client, url: &str) -> Self {
        Bridge {
            client: client,
            matrix: matrix::client::Client::new(url),
            rooms: HashMap::new(),
            seenEvents: vec![],
            server_name: "oob.systems".to_string()
        }
    }

    pub fn run(&mut self) {
        let mut events = EventLoop::new().unwrap();
        events.register(self.client.stream(), CLIENT, EventSet::all(), PollOpt::edge()).unwrap();
        events.run(self).unwrap();
    }

    fn handle_matrix(&mut self, evt: matrix::events::Event) {
        let duplicate = match evt.id {
            Some(ref id) =>
                self.seenEvents.contains(id),
            _ => false
        };
        if !duplicate {
            let mut messages: Vec<irc::protocol::Message> = vec![];
            {
                let appendMsg = |msg: irc::protocol::Message| {
                    messages.push(msg);
                };
                match evt.data {
                    matrix::events::EventData::Room(room_id, room_event) => {
                        self.room_from_matrix(&room_id).handle_event(room_event, appendMsg);
                    },
                    _ => println!("Unhandled {:?}", evt)
                }
            }
            match evt.id {
                Some(id) =>
                    self.seenEvents.push(id),
                None => ()
            };
            for ref msg in messages {
                self.client.send(msg);
            }
        }
    }

    fn poll_matrix(&mut self, channel: mio::Sender<Event>) ->
        thread::JoinHandle<matrix::client::Result> {
        let poll = self.matrix.pollAsync();
        thread::spawn(move|| {
            poll.send().and_then(|evts| {
                for evt in evts {
                    channel.send(Event::Matrix(evt));
                };
                channel.send(Event::EndPoll);
                Ok(())
            })
        })
    }

    fn start_matrix(&mut self, channel: mio::Sender<Event>) ->
        matrix::client::Result {
        self.matrix.sync(|evt: matrix::events::Event| {
            channel.send(Event::Matrix(evt));
        }).and_then(|r| {
            self.poll_matrix(channel);
            Ok(r)
        })
    }

    fn handle_client(&mut self, events: &mut EventLoop<Bridge>) {
        loop {
            match self.client.read_message() {
                None => return,
                Some(message) => {
                    println!("Got a message! {:?}", message);
                    match message.command {
                        Command::Pass => {
                            self.client.auth.set_password(message.args[0].clone())
                        }
                        Command::Nick => {
                            self.client.set_nickname(message.args[0].clone());
                        },
                        Command::User => {
                            println!("User logged in: {}", message.args[0]);
                            self.client.auth.set_username(message.args[0].clone());
                            let auth = self.client.auth.consume();
                            match (auth.username, auth.password) {
                                (Some(username), Some(password)) => {
                                    self.matrix.login(username.trim(), password.trim())
                                        .and(self.start_matrix(events.channel()))
                                        .and_then(|_| {
                                            self.client.welcome(username.trim());
                                            println!("Logged in {:?}", username);
                                            Ok(())
                                        }).expect("Could not login!");
                                },
                                _ => panic!("Username and/or password missing")
                            };
                        },
                        Command::Join => {
                            self.client.join(&message.args[0]);
                        },
                        Command::Ping => {
                            self.client.pong();
                        },
                        Command::Quit => {
                            // FIXME: Logout of matrix and exit thread
                            return;
                        },
                        Command::Privmsg => {
                            let evt = {
                                let id = self.matrix.uid.clone().unwrap();
                                let room = self.room_from_irc(&message.args[0]);
                                matrix::events::EventData::Room(
                                    room.id.clone(),
                                    matrix::events::RoomEvent::Message(
                                        id, message.suffix.unwrap()))
                            };
                            self.seenEvents.push(self.matrix.send(evt).unwrap());
                        },
                        _ =>
                            println!("unhandled {:?}", message)
                    }
                }
            }
        }
    }
}

<|endoftext|>


<commit_before>use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, quote_spanned, ToTokens};

#[derive(Clone, Copy, PartialEq)]
enum RuntimeFlavor {
    CurrentThread,
    Threaded,
}

impl RuntimeFlavor {
    fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
        match s {
            "current_thread" => Ok(RuntimeFlavor::CurrentThread),
            "multi_thread" => Ok(RuntimeFlavor::Threaded),
            "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
            "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
            "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
            _ => Err(format!("No such runtime flavor `{}`. The runtime flavors are `current_thread` and `multi_thread`.", s)),
        }
    }
}

struct FinalConfig {
    flavor: RuntimeFlavor,
    worker_threads: Option<usize>,
    start_paused: Option<bool>,
}

struct Configuration {
    rt_multi_thread_available: bool,
    default_flavor: RuntimeFlavor,
    flavor: Option<RuntimeFlavor>,
    worker_threads: Option<(usize, Span)>,
    start_paused: Option<(bool, Span)>,
    is_test: bool,
}

impl Configuration {
    fn new(is_test: bool, rt_multi_thread: bool) -> Self {
        Configuration {
            rt_multi_thread_available: rt_multi_thread,
            default_flavor: match is_test {
                true => RuntimeFlavor::CurrentThread,
                false => RuntimeFlavor::Threaded,
            },
            flavor: None,
            worker_threads: None,
            start_paused: None,
            is_test,
        }
    }

    fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.flavor.is_some() {
            return Err(syn::Error::new(span, "`flavor` set multiple times."));
        }

        let runtime_str = parse_string(runtime, span, "flavor")?;
        let runtime =
            RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
        self.flavor = Some(runtime);
        Ok(())
    }

    fn set_worker_threads(
        &mut self,
        worker_threads: syn::Lit,
        span: Span,
    ) -> Result<(), syn::Error> {
        if self.worker_threads.is_some() {
            return Err(syn::Error::new(
                span,
                "`worker_threads` set multiple times.",
            ));
        }

        let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
        if worker_threads == 0 {
            return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
        }
        self.worker_threads = Some((worker_threads, span));
        Ok(())
    }

    fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.start_paused.is_some() {
            return Err(syn::Error::new(span, "`start_paused` set multiple times."));
        }

        let start_paused = parse_bool(start_paused, span, "start_paused")?;
        self.start_paused = Some((start_paused, span));
        Ok(())
    }

    fn macro_name(&self) -> &'static str {
        if self.is_test {
            "tokio::test"
        } else {
            "tokio::main"
        }
    }

    fn build(&self) -> Result<FinalConfig, syn::Error> {
        let flavor = self.flavor.unwrap_or(self.default_flavor);
        use RuntimeFlavor::*;

        let worker_threads = match (flavor, self.worker_threads) {
            (CurrentThread, Some((_, worker_threads_span))) => {
                let msg = format!(
                    "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(worker_threads_span, msg));
            }
            (CurrentThread, None) => None,
            (Threaded, worker_threads) if self.rt_multi_thread_available => {
                worker_threads.map(|(val, _span)| val)
            }
            (Threaded, _) => {
                let msg = if self.flavor.is_none() {
                    "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
                } else {
                    "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
                };
                return Err(syn::Error::new(Span::call_site(), msg));
            }
        };

        let start_paused = match (flavor, self.start_paused) {
            (Threaded, Some((_, start_paused_span))) => {
                let msg = format!(
                    "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(start_paused_span, msg));
            }
            (CurrentThread, Some((start_paused, _))) => Some(start_paused),
            (_, None) => None,
        };

        Ok(FinalConfig {
            flavor,
            worker_threads,
            start_paused,
        })
    }
}

fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
    match int {
        syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
            Ok(value) => Ok(value),
            Err(e) => Err(syn::Error::new(
                span,
                format!("Failed to parse value of `{}` as integer: {}", field, e),
            )),
        },
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as integer.", field),
        )),
    }
}

fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
    match int {
        syn::Lit::Str(s) => Ok(s.value()),
        syn::Lit::Verbatim(s) => Ok(s.to_string()),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as string.", field),
        )),
    }
}

fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
    match bool {
        syn::Lit::Bool(b) => Ok(b.value),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as bool.", field),
        )),
    }
}

fn parse_knobs(
    mut input: syn::ItemFn,
    args: syn::AttributeArgs,
    is_test: bool,
    rt_multi_thread: bool,
) -> Result<TokenStream, syn::Error> {
    if input.sig.asyncness.take().is_none() {
        let msg = "the `async` keyword is missing from the function declaration";
        return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
    }

    let mut config = Configuration::new(is_test, rt_multi_thread);
    let macro_name = config.macro_name();

    for arg in args {
        match arg {
            syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
                let ident = namevalue
                    .path
                    .get_ident()
                    .ok_or_else(|| {
                        syn::Error::new_spanned(&namevalue, "Must have specified ident")
                    })?
                    .to_string()
                    .to_lowercase();
                match ident.as_str() {
                    "worker_threads" => {
                        config.set_worker_threads(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "flavor" => {
                        config.set_flavor(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "start_paused" => {
                        config.set_start_paused(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "core_threads" => {
                        let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                    name => {
                        let msg = format!(
                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
                            name,
                        );
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                }
            }
            syn::NestedMeta::Meta(syn::Meta::Path(path)) => {
                let name = path
                    .get_ident()
                    .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
                    .to_string()
                    .to_lowercase();
                let msg = match name.as_str() {
                    "threaded_scheduler" | "multi_thread" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"multi_thread\")].",
                            macro_name
                        )
                    }
                    "basic_scheduler" | "current_thread" | "single_threaded" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"current_thread\")].",
                            macro_name
                        )
                    }
                    "flavor" | "worker_threads" | "start_paused" => {
                        format!("The `{}` attribute requires an argument.", name)
                    }
                    name => {
                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
                    }
                };
                return Err(syn::Error::new_spanned(path, msg));
            }
            other => {
                return Err(syn::Error::new_spanned(
                    other,
                    "Unknown attribute inside the macro",
                ));
            }
        }
    }

    let config = config.build()?;

    // If type mismatch occurs, the current rustc points to the last statement.
    let (last_stmt_start_span, last_stmt_end_span) = {
        let mut last_stmt = input
            .block
            .stmts
            .last()
            .map(ToTokens::into_token_stream)
            .unwrap_or_default()
            .into_iter();
        // `Span` on stable Rust has a limitation that only points to the first
        // token, not the whole tokens. We can work around this limitation by
        // using the first/last span of the tokens like
        // `syn::Error::new_spanned` does.
        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
        let end = last_stmt.last().map_or(start, |t| t.span());
        (start, end)
    };

    let mut rt = match config.flavor {
        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_current_thread()
        },
        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_multi_thread()
        },
    };
    if let Some(v) = config.worker_threads {
        rt = quote! { #rt.worker_threads(#v) };
    }
    if let Some(v) = config.start_paused {
        rt = quote! { #rt.start_paused(#v) };
    }

    let header = if is_test {
        quote! {
            #[::core::prelude::v1::test]
        }
    } else {
        quote! {}
    };

    let body = &input.block;
    let brace_token = input.block.brace_token;
    input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
        {
            #rt
                .enable_all()
                .build()
                .expect("Failed building the Runtime")
                .block_on(async #body)
        }
    })
    .expect("Parsing failure");
    input.block.brace_token = brace_token;

    let result = quote! {
        #header
        #input
    };

    Ok(result.into())
}

#[cfg(not(test))] // Work around for rust-lang/rust#62127
pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
        let msg = "the main function cannot accept arguments";
        return syn::Error::new_spanned(&input.sig.ident, msg)
            .to_compile_error()
            .into();
    }

    parse_knobs(input, args, false, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}

pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    for attr in &input.attrs {
        if attr.path.is_ident("test") {
            let msg = "second test attribute is supplied";
            return syn::Error::new_spanned(&attr, msg)
                .to_compile_error()
                .into();
        }
    }

    parse_knobs(input, args, true, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}
<commit_msg>chore: explicitly relaxed clippy lint for runtime entry macro (#4030)<commit_after>use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, quote_spanned, ToTokens};

#[derive(Clone, Copy, PartialEq)]
enum RuntimeFlavor {
    CurrentThread,
    Threaded,
}

impl RuntimeFlavor {
    fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
        match s {
            "current_thread" => Ok(RuntimeFlavor::CurrentThread),
            "multi_thread" => Ok(RuntimeFlavor::Threaded),
            "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
            "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
            "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
            _ => Err(format!("No such runtime flavor `{}`. The runtime flavors are `current_thread` and `multi_thread`.", s)),
        }
    }
}

struct FinalConfig {
    flavor: RuntimeFlavor,
    worker_threads: Option<usize>,
    start_paused: Option<bool>,
}

struct Configuration {
    rt_multi_thread_available: bool,
    default_flavor: RuntimeFlavor,
    flavor: Option<RuntimeFlavor>,
    worker_threads: Option<(usize, Span)>,
    start_paused: Option<(bool, Span)>,
    is_test: bool,
}

impl Configuration {
    fn new(is_test: bool, rt_multi_thread: bool) -> Self {
        Configuration {
            rt_multi_thread_available: rt_multi_thread,
            default_flavor: match is_test {
                true => RuntimeFlavor::CurrentThread,
                false => RuntimeFlavor::Threaded,
            },
            flavor: None,
            worker_threads: None,
            start_paused: None,
            is_test,
        }
    }

    fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.flavor.is_some() {
            return Err(syn::Error::new(span, "`flavor` set multiple times."));
        }

        let runtime_str = parse_string(runtime, span, "flavor")?;
        let runtime =
            RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
        self.flavor = Some(runtime);
        Ok(())
    }

    fn set_worker_threads(
        &mut self,
        worker_threads: syn::Lit,
        span: Span,
    ) -> Result<(), syn::Error> {
        if self.worker_threads.is_some() {
            return Err(syn::Error::new(
                span,
                "`worker_threads` set multiple times.",
            ));
        }

        let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
        if worker_threads == 0 {
            return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
        }
        self.worker_threads = Some((worker_threads, span));
        Ok(())
    }

    fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.start_paused.is_some() {
            return Err(syn::Error::new(span, "`start_paused` set multiple times."));
        }

        let start_paused = parse_bool(start_paused, span, "start_paused")?;
        self.start_paused = Some((start_paused, span));
        Ok(())
    }

    fn macro_name(&self) -> &'static str {
        if self.is_test {
            "tokio::test"
        } else {
            "tokio::main"
        }
    }

    fn build(&self) -> Result<FinalConfig, syn::Error> {
        let flavor = self.flavor.unwrap_or(self.default_flavor);
        use RuntimeFlavor::*;

        let worker_threads = match (flavor, self.worker_threads) {
            (CurrentThread, Some((_, worker_threads_span))) => {
                let msg = format!(
                    "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(worker_threads_span, msg));
            }
            (CurrentThread, None) => None,
            (Threaded, worker_threads) if self.rt_multi_thread_available => {
                worker_threads.map(|(val, _span)| val)
            }
            (Threaded, _) => {
                let msg = if self.flavor.is_none() {
                    "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
                } else {
                    "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
                };
                return Err(syn::Error::new(Span::call_site(), msg));
            }
        };

        let start_paused = match (flavor, self.start_paused) {
            (Threaded, Some((_, start_paused_span))) => {
                let msg = format!(
                    "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(start_paused_span, msg));
            }
            (CurrentThread, Some((start_paused, _))) => Some(start_paused),
            (_, None) => None,
        };

        Ok(FinalConfig {
            flavor,
            worker_threads,
            start_paused,
        })
    }
}

fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
    match int {
        syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
            Ok(value) => Ok(value),
            Err(e) => Err(syn::Error::new(
                span,
                format!("Failed to parse value of `{}` as integer: {}", field, e),
            )),
        },
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as integer.", field),
        )),
    }
}

fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
    match int {
        syn::Lit::Str(s) => Ok(s.value()),
        syn::Lit::Verbatim(s) => Ok(s.to_string()),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as string.", field),
        )),
    }
}

fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
    match bool {
        syn::Lit::Bool(b) => Ok(b.value),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as bool.", field),
        )),
    }
}

fn parse_knobs(
    mut input: syn::ItemFn,
    args: syn::AttributeArgs,
    is_test: bool,
    rt_multi_thread: bool,
) -> Result<TokenStream, syn::Error> {
    if input.sig.asyncness.take().is_none() {
        let msg = "the `async` keyword is missing from the function declaration";
        return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
    }

    let mut config = Configuration::new(is_test, rt_multi_thread);
    let macro_name = config.macro_name();

    for arg in args {
        match arg {
            syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
                let ident = namevalue
                    .path
                    .get_ident()
                    .ok_or_else(|| {
                        syn::Error::new_spanned(&namevalue, "Must have specified ident")
                    })?
                    .to_string()
                    .to_lowercase();
                match ident.as_str() {
                    "worker_threads" => {
                        config.set_worker_threads(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "flavor" => {
                        config.set_flavor(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "start_paused" => {
                        config.set_start_paused(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "core_threads" => {
                        let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                    name => {
                        let msg = format!(
                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
                            name,
                        );
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                }
            }
            syn::NestedMeta::Meta(syn::Meta::Path(path)) => {
                let name = path
                    .get_ident()
                    .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
                    .to_string()
                    .to_lowercase();
                let msg = match name.as_str() {
                    "threaded_scheduler" | "multi_thread" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"multi_thread\")].",
                            macro_name
                        )
                    }
                    "basic_scheduler" | "current_thread" | "single_threaded" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"current_thread\")].",
                            macro_name
                        )
                    }
                    "flavor" | "worker_threads" | "start_paused" => {
                        format!("The `{}` attribute requires an argument.", name)
                    }
                    name => {
                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
                    }
                };
                return Err(syn::Error::new_spanned(path, msg));
            }
            other => {
                return Err(syn::Error::new_spanned(
                    other,
                    "Unknown attribute inside the macro",
                ));
            }
        }
    }

    let config = config.build()?;

    // If type mismatch occurs, the current rustc points to the last statement.
    let (last_stmt_start_span, last_stmt_end_span) = {
        let mut last_stmt = input
            .block
            .stmts
            .last()
            .map(ToTokens::into_token_stream)
            .unwrap_or_default()
            .into_iter();
        // `Span` on stable Rust has a limitation that only points to the first
        // token, not the whole tokens. We can work around this limitation by
        // using the first/last span of the tokens like
        // `syn::Error::new_spanned` does.
        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
        let end = last_stmt.last().map_or(start, |t| t.span());
        (start, end)
    };

    let mut rt = match config.flavor {
        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_current_thread()
        },
        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_multi_thread()
        },
    };
    if let Some(v) = config.worker_threads {
        rt = quote! { #rt.worker_threads(#v) };
    }
    if let Some(v) = config.start_paused {
        rt = quote! { #rt.start_paused(#v) };
    }

    let header = if is_test {
        quote! {
            #[::core::prelude::v1::test]
        }
    } else {
        quote! {}
    };

    let body = &input.block;
    let brace_token = input.block.brace_token;
    input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
        {
            let body = async #body;
            #[allow(clippy::expect_used)]
            #rt
                .enable_all()
                .build()
                .expect("Failed building the Runtime")
                .block_on(body)
        }
    })
    .expect("Parsing failure");
    input.block.brace_token = brace_token;

    let result = quote! {
        #header
        #input
    };

    Ok(result.into())
}

#[cfg(not(test))] // Work around for rust-lang/rust#62127
pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
        let msg = "the main function cannot accept arguments";
        return syn::Error::new_spanned(&input.sig.ident, msg)
            .to_compile_error()
            .into();
    }

    parse_knobs(input, args, false, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}

pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    for attr in &input.attrs {
        if attr.path.is_ident("test") {
            let msg = "second test attribute is supplied";
            return syn::Error::new_spanned(&attr, msg)
                .to_compile_error()
                .into();
        }
    }

    parse_knobs(input, args, true, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}
<|endoftext|>


<commit_before>use crate::io::{AsyncRead, AsyncWrite};

use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

cfg_io_util! {
    /// A future that asynchronously copies the entire contents of a reader into a
    /// writer.
    ///
    /// This struct is generally created by calling [`copy`][copy]. Please
    /// see the documentation of `copy()` for more details.
    ///
    /// [copy]: copy()
    #[derive(Debug)]
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    pub struct Copy<'a, R: ?Sized, W: ?Sized> {
        reader: &'a mut R,
        read_done: bool,
        writer: &'a mut W,
        pos: usize,
        cap: usize,
        amt: u64,
        buf: Box<[u8]>,
    }

    /// Asynchronously copies the entire contents of a reader into a writer.
    ///
    /// This function returns a future that will continuously read data from
    /// `reader` and then write it into `writer` in a streaming fashion until
    /// `reader` returns EOF.
    ///
    /// On success, the total number of bytes that were copied from `reader` to
    /// `writer` is returned.
    ///
    /// This is an asynchronous version of [`std::io::copy`][std].
    ///
    /// [std]: std::io::copy
    ///
    /// # Errors
    ///
    /// The returned future will finish with an error will return an error
    /// immediately if any call to `poll_read` or `poll_write` returns an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::io;
    ///
    /// # async fn dox() -> std::io::Result<()> {
    /// let mut reader: &[u8] = b"hello";
    /// let mut writer: Vec<u8> = vec![];
    ///
    /// io::copy(&mut reader, &mut writer).await?;
    ///
    /// assert_eq!(&b"hello"[..], &writer[..]);
    /// # Ok(())
    /// # }
    /// ```
    pub fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> Copy<'a, R, W>
    where
        R: AsyncRead + Unpin + ?Sized,
        W: AsyncWrite + Unpin + ?Sized,
    {
        Copy {
            reader,
            read_done: false,
            writer,
            amt: 0,
            pos: 0,
            cap: 0,
            buf: Box::new([0; 2048]),
        }
    }
}

impl<R, W> Future for Copy<'_, R, W>
where
    R: AsyncRead + Unpin + ?Sized,
    W: AsyncWrite + Unpin + ?Sized,
{
    type Output = io::Result<u64>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        loop {
            // If our buffer is empty, then we need to read some data to
            // continue.
            if self.pos == self.cap && !self.read_done {
                let me = &mut *self;
                let n = ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf))?;
                if n == 0 {
                    self.read_done = true;
                } else {
                    self.pos = 0;
                    self.cap = n;
                }
            }

            // If our buffer has some data, let's write it out!
            while self.pos < self.cap {
                let me = &mut *self;
                let i = ready!(Pin::new(&mut *me.writer).poll_write(cx, &me.buf[me.pos..me.cap]))?;
                if i == 0 {
                    return Poll::Ready(Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "write zero byte into writer",
                    )));
                } else {
                    self.pos += i;
                    self.amt += i as u64;
                }
            }

            // If we've written all the data and we've seen EOF, flush out the
            // data and finish the transfer.
            if self.pos == self.cap && self.read_done {
                let me = &mut *self;
                ready!(Pin::new(&mut *me.writer).poll_flush(cx))?;
                return Poll::Ready(Ok(self.amt));
            }
        }
    }
}

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

    #[test]
    fn assert_unpin() {
        use std::marker::PhantomPinned;
        crate::is_unpin::<Copy<'_, PhantomPinned, PhantomPinned>>();
    }
}
<commit_msg>io: allocate buffer directly on heap (#2634)<commit_after>use crate::io::{AsyncRead, AsyncWrite};

use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

cfg_io_util! {
    /// A future that asynchronously copies the entire contents of a reader into a
    /// writer.
    ///
    /// This struct is generally created by calling [`copy`][copy]. Please
    /// see the documentation of `copy()` for more details.
    ///
    /// [copy]: copy()
    #[derive(Debug)]
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    pub struct Copy<'a, R: ?Sized, W: ?Sized> {
        reader: &'a mut R,
        read_done: bool,
        writer: &'a mut W,
        pos: usize,
        cap: usize,
        amt: u64,
        buf: Box<[u8]>,
    }

    /// Asynchronously copies the entire contents of a reader into a writer.
    ///
    /// This function returns a future that will continuously read data from
    /// `reader` and then write it into `writer` in a streaming fashion until
    /// `reader` returns EOF.
    ///
    /// On success, the total number of bytes that were copied from `reader` to
    /// `writer` is returned.
    ///
    /// This is an asynchronous version of [`std::io::copy`][std].
    ///
    /// [std]: std::io::copy
    ///
    /// # Errors
    ///
    /// The returned future will finish with an error will return an error
    /// immediately if any call to `poll_read` or `poll_write` returns an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::io;
    ///
    /// # async fn dox() -> std::io::Result<()> {
    /// let mut reader: &[u8] = b"hello";
    /// let mut writer: Vec<u8> = vec![];
    ///
    /// io::copy(&mut reader, &mut writer).await?;
    ///
    /// assert_eq!(&b"hello"[..], &writer[..]);
    /// # Ok(())
    /// # }
    /// ```
    pub fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> Copy<'a, R, W>
    where
        R: AsyncRead + Unpin + ?Sized,
        W: AsyncWrite + Unpin + ?Sized,
    {
        Copy {
            reader,
            read_done: false,
            writer,
            amt: 0,
            pos: 0,
            cap: 0,
            buf: vec![0; 2048].into_boxed_slice(),
        }
    }
}

impl<R, W> Future for Copy<'_, R, W>
where
    R: AsyncRead + Unpin + ?Sized,
    W: AsyncWrite + Unpin + ?Sized,
{
    type Output = io::Result<u64>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        loop {
            // If our buffer is empty, then we need to read some data to
            // continue.
            if self.pos == self.cap && !self.read_done {
                let me = &mut *self;
                let n = ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf))?;
                if n == 0 {
                    self.read_done = true;
                } else {
                    self.pos = 0;
                    self.cap = n;
                }
            }

            // If our buffer has some data, let's write it out!
            while self.pos < self.cap {
                let me = &mut *self;
                let i = ready!(Pin::new(&mut *me.writer).poll_write(cx, &me.buf[me.pos..me.cap]))?;
                if i == 0 {
                    return Poll::Ready(Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "write zero byte into writer",
                    )));
                } else {
                    self.pos += i;
                    self.amt += i as u64;
                }
            }

            // If we've written all the data and we've seen EOF, flush out the
            // data and finish the transfer.
            if self.pos == self.cap && self.read_done {
                let me = &mut *self;
                ready!(Pin::new(&mut *me.writer).poll_flush(cx))?;
                return Poll::Ready(Ok(self.amt));
            }
        }
    }
}

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

    #[test]
    fn assert_unpin() {
        use std::marker::PhantomPinned;
        crate::is_unpin::<Copy<'_, PhantomPinned, PhantomPinned>>();
    }
}
<|endoftext|>