<commit_before>extern crate ansi_term;
extern crate rand;

use ansi_term::Colour::*;
use rand::Rng;
use std::io;

enum Card {
    ClubsAce,
    ClubsTen,
    ClubsKing,
    ClubsQueen,
    ClubsJack,
    ClubsNine,
    ClubsEight,
    ClubsSeven,
    SpadesAce,
    SpadesTen,
    SpadesKing,
    SpadesQueen,
    SpadesJack,
    SpadesNine,
    SpadesEight,
    SpadesSeven,
    HeartsAce,
    HeartsTen,
    HeartsKing,
    HeartsQueen,
    HeartsJack,
    HeartsNine,
    HeartsEight,
    HeartsSeven,
    DiamondsAce,
    DiamondsTen,
    DiamondsKing,
    DiamondsQueen,
    DiamondsJack,
    DiamondsNine,
    DiamondsEight,
    DiamondsSeven,
}

enum Game {
    Suit,
    Grand,
    Null,
}

enum Player {
    A,
    B,
    C,
}

fn announce(player: u8) -> Game {
    Game::Suit
}

fn bid(dealer: u8) -> (u8, u8) {
    let winner: u8 = dealer;
    let highest: u8 = 18;
    // TODO
    (winner, highest)
}

fn deal(dealer: u8) {
    let mut cards: Vec<u8> = Vec::new();
    // print sorted cards
    for n in 0..32 {
        cards.push(n);
    }
    for n in 0..8 {
        print_card(n);
    };
    println!("");
    for n in 8..16 {
        print_card(n);
    }
    println!("");
    for n in 16..24 {
        print_card(n);
    };
    println!("");
    for n in 24..32 {
        print_card(n);
    };
    println!("");
}

fn print_card(card: u8) {
    let club    = "♣";
    let spade   = "♠";
    let heart   = "♥";
    let diamond = "♦";
    match card {
        0 => {
            let mut ace = "A".to_string();
            ace.push_str(club);
            print!("{} ", Black.bold().paint(&ace));
        },
        1 => {
            let mut ten = "10".to_string();
            ten.push_str(club);
            print!("{} ", Black.bold().paint(&ten));
        },
        2 => {
            let mut king = "K".to_string();
            king.push_str(club);
            print!("{} ", Black.bold().paint(&king));
        },
        3 => {
            let mut queen = "Q".to_string();
            queen.push_str(club);
            print!("{} ", Black.bold().paint(&queen));
        },
        4 => {
            let mut jack = "J".to_string();
            jack.push_str(club);
            print!("{} ", Black.bold().paint(&jack));
        },
        5 => {
            let mut nine = "9".to_string();
            nine.push_str(club);
            print!("{} ", Black.bold().paint(&nine));
        },
        6 => {
            let mut eight = "8".to_string();
            eight.push_str(club);
            print!("{} ", Black.bold().paint(&eight));
        },
        7 => {
            let mut seven = "7".to_string();
            seven.push_str(club);
            print!("{} ", Black.bold().paint(&seven));
        },
        8 => {
            let mut ace = "A".to_string();
            ace.push_str(spade);
            print!("{} ", Green.bold().paint(&ace));
        },
        9 => {
            let mut ten = "10".to_string();
            ten.push_str(spade);
            print!("{} ", Green.bold().paint(&ten));
        },
        10 => {
            let mut king = "K".to_string();
            king.push_str(spade);
            print!("{} ", Green.bold().paint(&king));
        },
        11 => {
            let mut queen = "Q".to_string();
            queen.push_str(spade);
            print!("{} ", Green.bold().paint(&queen));
        },
        12 => {
            let mut jack = "J".to_string();
            jack.push_str(spade);
            print!("{} ", Green.bold().paint(&jack));
        },
        13 => {
            let mut nine = "9".to_string();
            nine.push_str(spade);
            print!("{} ", Green.bold().paint(&nine));
        },
        14 => {
            let mut eight = "8".to_string();
            eight.push_str(spade);
            print!("{} ", Green.bold().paint(&eight));
        },
        15 => {
            let mut seven = "7".to_string();
            seven.push_str(spade);
            print!("{} ", Green.bold().paint(&seven));
        },
        16 => {
            let mut ace = "A".to_string();
            ace.push_str(heart);
            print!("{} ", Red.bold().paint(&ace));
        },
        17 => {
            let mut ten = "10".to_string();
            ten.push_str(heart);
            print!("{} ", Red.bold().paint(&ten));
        },
        18 => {
            let mut king = "K".to_string();
            king.push_str(heart);
            print!("{} ", Red.bold().paint(&king));
        },
        19 => {
            let mut queen = "Q".to_string();
            queen.push_str(heart);
            print!("{} ", Red.bold().paint(&queen));
        },
        20 => {
            let mut jack = "J".to_string();
            jack.push_str(heart);
            print!("{} ", Red.bold().paint(&jack));
        },
        21 => {
            let mut nine = "9".to_string();
            nine.push_str(heart);
            print!("{} ", Red.bold().paint(&nine));
        },
        22 => {
            let mut eight = "8".to_string();
            eight.push_str(heart);
            print!("{} ", Red.bold().paint(&eight));
        },
        23 => {
            let mut seven = "7".to_string();
            seven.push_str(heart);
            print!("{} ", Red.bold().paint(&seven));
        },
        24 => {
            let mut ace = "A".to_string();
            ace.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&ace));
        },
        25 => {
            let mut ten = "10".to_string();
            ten.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&ten));
        },
        26 => {
            let mut king = "K".to_string();
            king.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&king));
        },
        27 => {
            let mut queen = "Q".to_string();
            queen.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&queen));
        },
        28 => {
            let mut jack = "J".to_string();
            jack.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&jack));
        },
        29 => {
            let mut nine = "9".to_string();
            nine.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&nine));
        },
        30 => {
            let mut eight = "8".to_string();
            eight.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&eight));
        },
        31 => {
            let mut seven = "7".to_string();
            seven.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&seven));
        },
        _ => panic!("Unknown card"),
    }
}

fn main() {
    // randomly select player
    let player_number: u8 = rand::thread_rng().gen_range(0, 3);
    match player_number {
        0 => println!("player A:"),
        1 => println!("player B:"),
        2 => println!("player C:"),
        _ => panic!("Unknown player {}", player_number),
    }
    loop {
        // player player_number is dealing
        deal(player_number);
        // bid
        let (player, bid) = bid(player_number);
        // announce
        let game = announce(player);
        // play 10 tricks in a row
        for trick in 0..10 {
            println!("trick #{}:", trick);
            // ask for input
            let mut input = String::new();
            io::stdin().read_line(&mut input)
                .ok()
                .expect("failed to read line");
            let input: u8 = match input.trim().parse() {
                Ok(num) => num,
                Err(_) => break,
            };
            // is input valid?
        };
        println!("New game?");
        let mut input = String::new();
        io::stdin().read_line(&mut input)
            .ok()
            .expect("failed to read line");
        let input: u8 = match input.trim().parse() {
            Ok(num) => num,
            Err(_) => break,
        };
        println!("your input: {}", input);
        if input != 1 {
            break;
        }
    }
}
<commit_msg>Print shuffled cards (per player plus Skat).<commit_after>extern crate ansi_term;
extern crate rand;

use ansi_term::Colour::*;
use rand::Rng;
use std::io;

enum Card {
    ClubsAce,
    ClubsTen,
    ClubsKing,
    ClubsQueen,
    ClubsJack,
    ClubsNine,
    ClubsEight,
    ClubsSeven,
    SpadesAce,
    SpadesTen,
    SpadesKing,
    SpadesQueen,
    SpadesJack,
    SpadesNine,
    SpadesEight,
    SpadesSeven,
    HeartsAce,
    HeartsTen,
    HeartsKing,
    HeartsQueen,
    HeartsJack,
    HeartsNine,
    HeartsEight,
    HeartsSeven,
    DiamondsAce,
    DiamondsTen,
    DiamondsKing,
    DiamondsQueen,
    DiamondsJack,
    DiamondsNine,
    DiamondsEight,
    DiamondsSeven,
}

enum Game {
    Suit,
    Grand,
    Null,
}

enum Player {
    A,
    B,
    C,
}

fn announce(player: u8) -> Game {
    Game::Suit
}

fn bid(dealer: u8) -> (u8, u8) {
    let winner: u8 = dealer;
    let highest: u8 = 18;
    // TODO
    (winner, highest)
}

fn deal(dealer: u8) {
    let mut cards: Vec<u8> = Vec::new();
    // print sorted cards
    for n in 0..32 {
        cards.push(n);
    }
    for n in 0..8 {
        print_card(n);
    };
    println!("");
    for n in 8..16 {
        print_card(n);
    }
    println!("");
    for n in 16..24 {
        print_card(n);
    };
    println!("");
    for n in 24..32 {
        print_card(n);
    };
    println!("");
    println!("");
    // TODO: shuffle cards
    // print shuffled cards (per player plus Skat)
    for n in 0..10 {
        print_card(n);
    };
    println!("");
    for n in 10..20 {
        print_card(n);
    }
    println!("");
    for n in 20..30 {
        print_card(n);
    };
    println!("");
    for n in 30..32 {
        print_card(n);
    };
    println!("");
}

fn print_card(card: u8) {
    let club    = "♣";
    let spade   = "♠";
    let heart   = "♥";
    let diamond = "♦";
    match card {
        0 => {
            let mut ace = " A".to_string();
            ace.push_str(club);
            print!("{} ", Black.bold().paint(&ace));
        },
        1 => {
            let mut ten = "10".to_string();
            ten.push_str(club);
            print!("{} ", Black.bold().paint(&ten));
        },
        2 => {
            let mut king = " K".to_string();
            king.push_str(club);
            print!("{} ", Black.bold().paint(&king));
        },
        3 => {
            let mut queen = " Q".to_string();
            queen.push_str(club);
            print!("{} ", Black.bold().paint(&queen));
        },
        4 => {
            let mut jack = " J".to_string();
            jack.push_str(club);
            print!("{} ", Black.bold().paint(&jack));
        },
        5 => {
            let mut nine = " 9".to_string();
            nine.push_str(club);
            print!("{} ", Black.bold().paint(&nine));
        },
        6 => {
            let mut eight = " 8".to_string();
            eight.push_str(club);
            print!("{} ", Black.bold().paint(&eight));
        },
        7 => {
            let mut seven = " 7".to_string();
            seven.push_str(club);
            print!("{} ", Black.bold().paint(&seven));
        },
        8 => {
            let mut ace = " A".to_string();
            ace.push_str(spade);
            print!("{} ", Green.bold().paint(&ace));
        },
        9 => {
            let mut ten = "10".to_string();
            ten.push_str(spade);
            print!("{} ", Green.bold().paint(&ten));
        },
        10 => {
            let mut king = " K".to_string();
            king.push_str(spade);
            print!("{} ", Green.bold().paint(&king));
        },
        11 => {
            let mut queen = " Q".to_string();
            queen.push_str(spade);
            print!("{} ", Green.bold().paint(&queen));
        },
        12 => {
            let mut jack = " J".to_string();
            jack.push_str(spade);
            print!("{} ", Green.bold().paint(&jack));
        },
        13 => {
            let mut nine = " 9".to_string();
            nine.push_str(spade);
            print!("{} ", Green.bold().paint(&nine));
        },
        14 => {
            let mut eight = " 8".to_string();
            eight.push_str(spade);
            print!("{} ", Green.bold().paint(&eight));
        },
        15 => {
            let mut seven = " 7".to_string();
            seven.push_str(spade);
            print!("{} ", Green.bold().paint(&seven));
        },
        16 => {
            let mut ace = " A".to_string();
            ace.push_str(heart);
            print!("{} ", Red.bold().paint(&ace));
        },
        17 => {
            let mut ten = "10".to_string();
            ten.push_str(heart);
            print!("{} ", Red.bold().paint(&ten));
        },
        18 => {
            let mut king = " K".to_string();
            king.push_str(heart);
            print!("{} ", Red.bold().paint(&king));
        },
        19 => {
            let mut queen = " Q".to_string();
            queen.push_str(heart);
            print!("{} ", Red.bold().paint(&queen));
        },
        20 => {
            let mut jack = " J".to_string();
            jack.push_str(heart);
            print!("{} ", Red.bold().paint(&jack));
        },
        21 => {
            let mut nine = " 9".to_string();
            nine.push_str(heart);
            print!("{} ", Red.bold().paint(&nine));
        },
        22 => {
            let mut eight = " 8".to_string();
            eight.push_str(heart);
            print!("{} ", Red.bold().paint(&eight));
        },
        23 => {
            let mut seven = " 7".to_string();
            seven.push_str(heart);
            print!("{} ", Red.bold().paint(&seven));
        },
        24 => {
            let mut ace = " A".to_string();
            ace.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&ace));
        },
        25 => {
            let mut ten = "10".to_string();
            ten.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&ten));
        },
        26 => {
            let mut king = " K".to_string();
            king.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&king));
        },
        27 => {
            let mut queen = " Q".to_string();
            queen.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&queen));
        },
        28 => {
            let mut jack = " J".to_string();
            jack.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&jack));
        },
        29 => {
            let mut nine = " 9".to_string();
            nine.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&nine));
        },
        30 => {
            let mut eight = " 8".to_string();
            eight.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&eight));
        },
        31 => {
            let mut seven = " 7".to_string();
            seven.push_str(diamond);
            print!("{} ", Yellow.bold().paint(&seven));
        },
        _ => panic!("Unknown card"),
    }
}

fn main() {
    // randomly select player
    let player_number: u8 = rand::thread_rng().gen_range(0, 3);
    match player_number {
        0 => println!("player A:"),
        1 => println!("player B:"),
        2 => println!("player C:"),
        _ => panic!("Unknown player {}", player_number),
    }
    loop {
        // player player_number is dealing
        deal(player_number);
        // bid
        let (player, bid) = bid(player_number);
        // announce
        let game = announce(player);
        // play 10 tricks in a row
        for trick in 0..10 {
            println!("trick #{}:", trick);
            // ask for input
            let mut input = String::new();
            io::stdin().read_line(&mut input)
                .ok()
                .expect("failed to read line");
            let input: u8 = match input.trim().parse() {
                Ok(num) => num,
                Err(_) => break,
            };
            // is input valid?
        };
        println!("New game?");
        let mut input = String::new();
        io::stdin().read_line(&mut input)
            .ok()
            .expect("failed to read line");
        let input: u8 = match input.trim().parse() {
            Ok(num) => num,
            Err(_) => break,
        };
        println!("your input: {}", input);
        if input != 1 {
            break;
        }
    }
}
<|endoftext|>


<commit_before>#[test]
fn pathmatch_test_fail()
{
    assert_eq!(2i, 3i);
}

<commit_msg>exact equality tests<commit_after>pub fn pathmatch(pattern: &str, pathstring: &str) -> bool
{
    return pattern == pathstring
}

#[test]
fn pathmatch_test_exact_equality()
{
    assert!(pathmatch("", ""));
    assert!(pathmatch("a", "a"));
    assert!(pathmatch("BC", "BC"));
    assert!(!pathmatch("a", ""));
    assert!(!pathmatch("", "BC"));
}

<|endoftext|>


<commit_before>#![allow(deprecated)]
use utils::data as D;
use std::ops::Index;
use std::path::PathBuf;
use std::{io, fs};
use curl::easy::Easy;
use flate2::{Flush, Decompress};
use std::io::Write;
use utils::data::{Dataset, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind};
use torch;

type Sample = (Tensor<f32>, Tensor<i64>);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    // XXX how to ignore EEXIST?
    fs::create_dir(raw_path.clone())?;
    fs::create_dir(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let fname: Vec<_> = paths[5].split("/").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len());
        // do we have a header?
        Decompress::new(false)
            .decompress_vec(data.as_slice(), &mut output, Flush::Finish)?;
        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set = vec![read_image_file(raw_path.join("train-images-idx3-ubyte")),
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte")),
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_label_file(path: PathBuf) -> TensorKind {
    unimplemented!()
}
fn read_image_file(path: PathBuf) -> TensorKind {
    unimplemented!()
}

pub struct MNIST {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<i64>,
//    pub transform: Option<fn(Vec<u8>)>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
//    #[builder(default="None")]
//    transform: Option<Box<fn(&Vec<u8>)>>,
}

impl MNISTArgsBuilder {
    pub fn done(self) -> Dataset<Sample> {
        let args = self.build().unwrap();
        Dataset::new(Rc::new(MNIST::new(&args)))
    }
}

impl MNIST {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: &MNISTArgs) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        let mut v = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
//            transform: args.transform,
        }
    }
}

impl Index<usize> for MNIST {
    type Output = Sample;
    fn index(&self, idx: usize) -> &Self::Output {
        unimplemented!()
    }
}

impl D::DatasetIntf<Sample> for MNIST {
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn iter(&self) -> Box<Iterator<Item = Sample>> {
        unimplemented!();
    }
}
<commit_msg>finish MNIST loader implementation<commit_after>#![allow(deprecated)]
use std::ops::Index;
use std::path::PathBuf;
use std::{io, fs, slice, cell};
use curl::easy::Easy;
use memmap::{Mmap, Protection};
use flate2::{Flush, Decompress};
use std::io::Write;
use utils::data::{Dataset, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind};
use torch;

type Sample = (TensorKind, i64);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";

fn create_dir_f(arg: PathBuf) -> io::Result<()> {
    let result = fs::create_dir(arg);
    if let Some(err) = result.err() {
        match err.kind() {
            io::ErrorKind::AlreadyExists => Ok(()),
            _ => Err(err),
        }
    } else {
        Ok(())
    }
}

fn check_exists(arg: &PathBuf) -> io::Result<bool> {
    let train = fs::metadata(arg.join(TRAINING_FILE))?;
    let test = fs::metadata(arg.join(TEST_FILE))?;

    return Ok(train.is_file() && test.is_file());
}

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    create_dir_f(raw_path.clone())?;
    create_dir_f(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let fname: Vec<_> = paths[5].split("/").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len());
        // do we have a header?
        Decompress::new(false)
            .decompress_vec(data.as_slice(), &mut output, Flush::Finish)?;
        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set = vec![read_image_file(raw_path.join("train-images-idx3-ubyte"))?,
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))?];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte"))?,
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))?];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_image_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let (length, nrows, ncols);
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 2051);
        length = data[1] as usize;
        nrows = data[2] as usize;
        ncols = data[3] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut images: Vec<Vec<u8>> = Vec::with_capacity(length);
    let mut idx = 16;
    for _ in 0..length {
        let mut img = Vec::with_capacity(nrows * ncols);
        for _ in 0..nrows * ncols {
            img.push(data[idx]);
            idx += 1;
        }
        images.push(img)
    }
    Ok(torch::byte_tensor(images).view(&[-1, 28, 28]))
}

fn read_label_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let length;
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 20549);
        length = data[1] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut labels: Vec<i64> = Vec::with_capacity(length);
    for i in 0..length {
        labels.push(data[8 + i] as i64)
    }
    Ok(torch::long_tensor(labels))
}

pub struct MNIST {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<i64>,
    pub transformed: cell::RefCell<Vec<(TensorKind, i64)>>,
    pub transform: Option<Xfrm>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
}

type Xfrm = Box<fn(&TensorKind) -> TensorKind>;
impl MNISTArgsBuilder {
    pub fn done(self, xfrm: Option<Xfrm>) -> Dataset<Sample> {
        let args = self.build().unwrap();
        Dataset::new(Rc::new(MNIST::new(&args, xfrm)))
    }
}

impl MNIST {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: &MNISTArgs, xfrm: Option<Xfrm>) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        if !check_exists(&processed_path).expect("Dataset not found, try downloading") {
            panic!("Dataset not found, use download=true to download it");
        }
        let mut v: Vec<TensorKind> = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));
        let count = data.size()[0];
        let mut mutable: Vec<(TensorKind, i64)> = Vec::with_capacity(count);
        for i in 0..count {
            mutable.push((data.s(i).copy().into(), i as i64))
        }

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
            transformed: cell::RefCell::new(mutable),
            transform: xfrm,
        }
    }
}

impl Index<usize> for MNIST {
    type Output = Sample;
    fn index(&self, idx: usize) -> &Self::Output {
        let img = self.data.s(idx);
        let xfrmp = self.transformed.as_ptr();
        let mut xfrm = unsafe { &mut *xfrmp };

        xfrm[idx] = (img.copy().into(), self.labels[idx] as i64);
        &xfrm[idx]
    }
}

impl DatasetIntf<Sample> for MNIST {
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn iter(&self) -> Box<Iterator<Item = Sample>> {
        unimplemented!();
    }
}
<|endoftext|>


<commit_before>use crate::schema::TestBackend;
use diesel::*;

#[test]
fn test_debug_count_output() {
    use crate::schema::users::dsl::*;
    let sql = debug_query::<TestBackend, _>(&users.count()).to_string();
    if cfg!(feature = "postgres") {
        assert_eq!(sql, r#"SELECT COUNT(*) FROM "users" -- binds: []"#);
    } else {
        assert_eq!(sql, "SELECT COUNT(*) FROM `users` -- binds: []");
    }
}

#[test]
fn test_debug_output() {
    use crate::schema::users::dsl::*;
    let command = update(users.filter(id.eq(1))).set(name.eq("new_name"));
    let sql = debug_query::<TestBackend, _>(&command).to_string();
    if cfg!(feature = "postgres") {
        assert_eq!(
            sql,
            r#"UPDATE "users" SET "name" = $1 WHERE "users"."id" = $2 -- binds: ["new_name", 1]"#
        )
    } else {
        assert_eq!(
            sql,
            r#"UPDATE `users` SET `name` = ? WHERE `users`.`id` = ? -- binds: ["new_name", 1]"#
        )
    }
}

#[test]
fn test_debug_batch_insert() {
    // This test ensures that we've implemented `debug_query` for batch insert
    // on sqlite
    // This requires a separate impl because it's more than one sql statement that
    // is executed

    use crate::schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq(Some("black"))),
        (name.eq("Tess"), hair_color.eq(None::<&str>)),
    ];
    let borrowed_command = insert_into(users).values(&values);
    let borrowed_sql_display = debug_query::<TestBackend, _>(&borrowed_command).to_string();
    let borrowed_sql_debug = format!("{:?}", debug_query::<TestBackend, _>(&borrowed_command));

    let owned_command = insert_into(users).values(values);
    let owned_sql_display = debug_query::<TestBackend, _>(&owned_command).to_string();
    let owned_sql_debug = format!("{:?}", debug_query::<TestBackend, _>(&owned_command));

    if cfg!(feature = "postgres") {
        assert_eq!(
            borrowed_sql_display,
            r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: "INSERT INTO \"users\" (\"name\", \"hair_color\") VALUES ($1, $2), ($3, $4)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: "INSERT INTO \"users\" (\"name\", \"hair_color\") VALUES ($1, $2), ($3, $4)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );
    } else if cfg!(feature = "sqlite") {
        assert_eq!(
            borrowed_sql_display,
            r#"BEGIN;
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Sean", Some("black")]
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Tess", None]
COMMIT;
"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: ["BEGIN", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Sean\", Some(\"black\")]", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Tess\", None]", "COMMIT"], binds: [] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"BEGIN;
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Sean", Some("black")]
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Tess", None]
COMMIT;
"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: ["BEGIN", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Sean\", Some(\"black\")]", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Tess\", None]", "COMMIT"], binds: [] }"#
        );
    } else {
        assert_eq!(
            borrowed_sql_display,
            r#"INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );
    }
}

#[test]
#[cfg(feature = "postgres")]
fn test_upsert() {
    // this test ensures we get the right debug string for upserts
    use crate::schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq(Some("black"))),
        (name.eq("Tess"), hair_color.eq(None::<&str>)),
    ];

    let upsert_command_single_where = insert_into(users)
        .values(values)
        .on_conflict(hair_color)
        .filter_target(hair_color.eq(Some("black")))
        .do_nothing();
    let upsert_single_where_sql_display =
        debug_query::<TestBackend, _>(&upsert_command_single_where).to_string();

    assert_eq!(
        upsert_single_where_sql_display,
        r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) ON CONFLICT ("hair_color") WHERE "users"."hair_color" = $5 DO NOTHING -- binds: ["Sean", Some("black"), "Tess", None, Some("black")]"#
    );

    // let upsert_command_second_where = insert_into(users)
    //     .values(values)
    //     .on_conflict(hair_color)
    //     .filter_target(name.eq(Some("Sean")))
    //     .filter_target(hair_color.eq(Some("black")))
    //     .do_nothing();

    // let upsert_second_where_sql_display =
    //     debug_query::<TestBackend, _>(&upsert_command_second_where).to_string();

    // assert_eq!(
    //     upsert_second_where_sql_display,
    //     r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) ON CONFLICT ("hair_color") WHERE "users"."hair_color" = $5 and WHERE "users"."name" = $6 DO NOTHING -- binds: ["Sean", Some("black"), "Tess", None, Some("black"), Some("Sean")]"#
    // );
}
<commit_msg>add test for repeated target filtering<commit_after>use crate::schema::TestBackend;
use diesel::*;

#[test]
fn test_debug_count_output() {
    use crate::schema::users::dsl::*;
    let sql = debug_query::<TestBackend, _>(&users.count()).to_string();
    if cfg!(feature = "postgres") {
        assert_eq!(sql, r#"SELECT COUNT(*) FROM "users" -- binds: []"#);
    } else {
        assert_eq!(sql, "SELECT COUNT(*) FROM `users` -- binds: []");
    }
}

#[test]
fn test_debug_output() {
    use crate::schema::users::dsl::*;
    let command = update(users.filter(id.eq(1))).set(name.eq("new_name"));
    let sql = debug_query::<TestBackend, _>(&command).to_string();
    if cfg!(feature = "postgres") {
        assert_eq!(
            sql,
            r#"UPDATE "users" SET "name" = $1 WHERE "users"."id" = $2 -- binds: ["new_name", 1]"#
        )
    } else {
        assert_eq!(
            sql,
            r#"UPDATE `users` SET `name` = ? WHERE `users`.`id` = ? -- binds: ["new_name", 1]"#
        )
    }
}

#[test]
fn test_debug_batch_insert() {
    // This test ensures that we've implemented `debug_query` for batch insert
    // on sqlite
    // This requires a separate impl because it's more than one sql statement that
    // is executed

    use crate::schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq(Some("black"))),
        (name.eq("Tess"), hair_color.eq(None::<&str>)),
    ];
    let borrowed_command = insert_into(users).values(&values);
    let borrowed_sql_display = debug_query::<TestBackend, _>(&borrowed_command).to_string();
    let borrowed_sql_debug = format!("{:?}", debug_query::<TestBackend, _>(&borrowed_command));

    let owned_command = insert_into(users).values(values);
    let owned_sql_display = debug_query::<TestBackend, _>(&owned_command).to_string();
    let owned_sql_debug = format!("{:?}", debug_query::<TestBackend, _>(&owned_command));

    if cfg!(feature = "postgres") {
        assert_eq!(
            borrowed_sql_display,
            r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: "INSERT INTO \"users\" (\"name\", \"hair_color\") VALUES ($1, $2), ($3, $4)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: "INSERT INTO \"users\" (\"name\", \"hair_color\") VALUES ($1, $2), ($3, $4)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );
    } else if cfg!(feature = "sqlite") {
        assert_eq!(
            borrowed_sql_display,
            r#"BEGIN;
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Sean", Some("black")]
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Tess", None]
COMMIT;
"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: ["BEGIN", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Sean\", Some(\"black\")]", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Tess\", None]", "COMMIT"], binds: [] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"BEGIN;
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Sean", Some("black")]
INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: ["Tess", None]
COMMIT;
"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: ["BEGIN", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Sean\", Some(\"black\")]", "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) -- binds: [\"Tess\", None]", "COMMIT"], binds: [] }"#
        );
    } else {
        assert_eq!(
            borrowed_sql_display,
            r#"INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            borrowed_sql_debug,
            r#"Query { sql: "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );

        assert_eq!(
            owned_sql_display,
            r#"INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?) -- binds: ["Sean", Some("black"), "Tess", None]"#
        );
        assert_eq!(
            owned_sql_debug,
            r#"Query { sql: "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?), (?, ?)", binds: ["Sean", Some("black"), "Tess", None] }"#
        );
    }
}

#[test]
#[cfg(feature = "postgres")]
fn test_upsert() {
    // this test ensures we get the right debug string for upserts
    use crate::schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq(Some("black"))),
        (name.eq("Tess"), hair_color.eq(None::<&str>)),
    ];

    let upsert_command_single_where = insert_into(users)
        .values(&values)
        .on_conflict(hair_color)
        .filter_target(hair_color.eq("black"))
        .do_nothing();
    let upsert_single_where_sql_display =
        debug_query::<TestBackend, _>(&upsert_command_single_where).to_string();

    assert_eq!(
        upsert_single_where_sql_display,
        r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) ON CONFLICT ("hair_color") WHERE "users"."hair_color" = $5 DO NOTHING -- binds: ["Sean", Some("black"), "Tess", None, "black"]"#
    );

    let upsert_command_second_where = insert_into(users)
        .values(&values)
        .on_conflict(hair_color)
        .filter_target(hair_color.eq("black"))
        .filter_target(name.eq("Sean"))
        .do_nothing();

    let upsert_second_where_sql_display =
        debug_query::<TestBackend, _>(&upsert_command_second_where).to_string();

    assert_eq!(
        upsert_second_where_sql_display,
        r#"INSERT INTO "users" ("name", "hair_color") VALUES ($1, $2), ($3, $4) ON CONFLICT ("hair_color") WHERE "users"."hair_color" = $5 AND "users"."name" = $6 DO NOTHING -- binds: ["Sean", Some("black"), "Tess", None, "black", "Sean"]"#
    );
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };

    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;

    let bold: fn(w: &str) -> String = if should_colorize {
        |w| Style::new().bold().paint(w).to_string()
    } else {
        |w| w.to_string()
    };

    println!("{} {}", bold("Vendor:"), info.vendor);
    println!("{} {}", bold("Product:"), info.product);
    println!("{} {}", bold("Version:"), info.version);
    println!("{} {}", bold("URL:"), info.url);
    println!("{}", bold("Interfaces:"));

    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    line_len: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        line_len.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

fn varlink_bridge(address: Option<&str>) -> Result<()> {
    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn main() -> Result<()> {
    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );
    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<commit_msg>varlink-cli: make json brackets non-bold<commit_after>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };

    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;

    let bold: fn(w: &str) -> String = if should_colorize {
        |w| Style::new().bold().paint(w).to_string()
    } else {
        |w| w.to_string()
    };

    println!("{} {}", bold("Vendor:"), info.vendor);
    println!("{} {}", bold("Product:"), info.product);
    println!("{} {}", bold("Version:"), info.version);
    println!("{} {}", bold("URL:"), info.url);
    println!("{}", bold("Interfaces:"));

    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    line_len: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        line_len.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            array_brackets: Style::new(),
            object_brackets: Style::new(),
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

fn varlink_bridge(address: Option<&str>) -> Result<()> {
    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn main() -> Result<()> {
    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );
    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use crate::libchain::chain::Chain;
use crate::types::block_number::{BlockNumber, BlockTag, Tag};
use crate::types::filter::Filter as FilterType;
use cita_types::H256;
use jsonrpc_types::rpc_types::{Filter, FilterChanges, Log};

/// The RPC interfaces about filter.
///     * newFilter
///     * newBlockFilter
///     * getFilterChanges
///     * getFilterLogs
///     * uninstallFilter
/// *Not include `getLogs`*.
pub trait RpcFilter {
    // Create a new filter and return the filter id
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#newfilter
    fn new_filter(&self, filter: Filter) -> usize;
    // Create a new filter that can listen the new block.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#newblockfilter
    fn new_block_filter(&self) -> usize;
    // Get the logs for the filter with the given id since last time it was called.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#getfilterchanges
    fn get_filter_changes(&self, id: usize) -> Option<FilterChanges>;
    // Get the logs for the filter with the given id.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#getfilterlogs
    fn get_filter_logs(&self, id: usize) -> Option<Vec<Log>>;
    // Remove the filter with the given id.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#uninstallfilter
    fn uninstall_filter(&self, id: usize) -> bool;
}

/// Helper for RpcFilter
trait FilterHelper {
    // Get the block filter with the given id
    fn get_block_filter(&self, id: usize) -> BlockNumber;
    // Get logs with given filter
    fn get_logs_with_filter(&self, filter: Filter, block_filter: BlockNumber) -> Vec<Log>;
}

impl FilterHelper for Chain {
    fn get_block_filter(&self, id: usize) -> BlockNumber {
        let filterdb = self.filter_db();
        let mut block_filter = BlockNumber::min_value();

        if let Some(block_number) = filterdb.try_lock().unwrap().get_block_filter(id) {
            block_filter = *block_number;
        }

        drop(filterdb);
        block_filter
    }

    fn get_logs_with_filter(&self, filter: Filter, block_filter: BlockNumber) -> Vec<Log> {
        let mut filter: FilterType = filter.into();
        filter.from_block = BlockTag::Height(block_filter);
        filter.to_block = BlockTag::Tag(Tag::Latest);
        let limit = filter.limit;
        split_logs(
            self.get_logs(&filter).into_iter().map(Into::into).collect(),
            limit,
        )
    }
}

impl RpcFilter for Chain {
    fn new_filter(&self, filter: Filter) -> usize {
        let filterdb = self.filter_db();
        let id = filterdb.try_lock().unwrap().gen_id();
        filterdb.try_lock().unwrap().gen_logs_filter(id, filter);
        drop(filterdb);
        id
    }

    fn new_block_filter(&self) -> usize {
        let filterdb = self.filter_db();
        let block_number = self.get_current_height();
        let id = filterdb.try_lock().unwrap().gen_id();
        filterdb
            .try_lock()
            .unwrap()
            .gen_block_filter(id, block_number);
        drop(filterdb);
        id
    }

    fn get_filter_changes(&self, id: usize) -> Option<FilterChanges> {
        let filterdb = self.filter_db();
        let mut changes = Some(FilterChanges::Empty);
        let current_number = self.get_current_height();
        let block_filter = self.get_block_filter(id);

        if !filterdb.try_lock().unwrap().is_filter(id) {
            drop(filterdb);
            return changes;
        }

        // Check the logs
        if let Some(filter) = filterdb.try_lock().unwrap().get_logs_filter(id) {
            trace!("Into filter changes: logs");
            changes = Some(FilterChanges::Logs(
                self.get_logs_with_filter(filter.clone(), block_filter),
            ));
        };

        // Check the block
        if filterdb.try_lock().unwrap().is_block_filter(id) {
            trace!("Into filter changes: block");
            // Return hashes from next block after the filter was created to the current block.
            let hashes = ((block_filter + 1)..=current_number)
                .filter_map(|_id| self.block_hash_by_height(_id))
                .collect::<Vec<H256>>();
            trace!("Block filter changes: {:?}", hashes);
            changes = Some(FilterChanges::Hashes(
                hashes.into_iter().map(Into::into).collect(),
            ));
        };

        // Update the block filter: use the current number
        filterdb
            .try_lock()
            .unwrap()
            .gen_block_filter(id, current_number + 1);
        drop(filterdb);
        changes
    }

    fn get_filter_logs(&self, id: usize) -> Option<Vec<Log>> {
        let filterdb = self.filter_db();
        let block_filter = self.get_block_filter(id);
        let logs = match filterdb.try_lock().unwrap().get_logs_filter(id) {
            Some(filter) => Some(self.get_logs_with_filter(filter.clone(), block_filter)),
            _ => None,
        };
        drop(filterdb);
        logs
    }

    fn uninstall_filter(&self, id: usize) -> bool {
        let filterdb = self.filter_db();
        let uninstall_ok = filterdb.try_lock().unwrap().uninstall(id);
        drop(filterdb);
        uninstall_ok
    }
}

/// Split the logs with the limit filter.
fn split_logs(mut logs: Vec<Log>, limit: Option<usize>) -> Vec<Log> {
    let len = logs.len();
    match limit {
        Some(limit) if len >= limit => logs.split_off(len - limit),
        _ => logs,
    }
}
<commit_msg>Fix newFilter: should add current block.<commit_after>use crate::libchain::chain::Chain;
use crate::types::block_number::{BlockNumber, BlockTag, Tag};
use crate::types::filter::Filter as FilterType;
use cita_types::H256;
use jsonrpc_types::rpc_types::{Filter, FilterChanges, Log};

/// The RPC interfaces about filter.
///     * newFilter
///     * newBlockFilter
///     * getFilterChanges
///     * getFilterLogs
///     * uninstallFilter
/// *Not include `getLogs`*.
pub trait RpcFilter {
    // Create a new filter and return the filter id
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#newfilter
    fn new_filter(&self, filter: Filter) -> usize;
    // Create a new filter that can listen the new block.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#newblockfilter
    fn new_block_filter(&self) -> usize;
    // Get the logs for the filter with the given id since last time it was called.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#getfilterchanges
    fn get_filter_changes(&self, id: usize) -> Option<FilterChanges>;
    // Get the logs for the filter with the given id.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#getfilterlogs
    fn get_filter_logs(&self, id: usize) -> Option<Vec<Log>>;
    // Remove the filter with the given id.
    // https://docs.citahub.com/zh-CN/cita/rpc-guide/rpc#uninstallfilter
    fn uninstall_filter(&self, id: usize) -> bool;
}

/// Helper for RpcFilter
trait FilterHelper {
    // Get the block filter with the given id
    fn get_block_filter(&self, id: usize) -> BlockNumber;
    // Get logs with given filter
    fn get_logs_with_filter(&self, filter: Filter, block_filter: BlockNumber) -> Vec<Log>;
}

impl FilterHelper for Chain {
    fn get_block_filter(&self, id: usize) -> BlockNumber {
        let filterdb = self.filter_db();
        let mut block_filter = BlockNumber::min_value();

        if let Some(block_number) = filterdb.try_lock().unwrap().get_block_filter(id) {
            block_filter = *block_number;
        }

        drop(filterdb);
        block_filter
    }

    fn get_logs_with_filter(&self, filter: Filter, block_filter: BlockNumber) -> Vec<Log> {
        let mut filter: FilterType = filter.into();
        filter.from_block = BlockTag::Height(block_filter);
        filter.to_block = BlockTag::Tag(Tag::Latest);
        let limit = filter.limit;
        split_logs(
            self.get_logs(&filter).into_iter().map(Into::into).collect(),
            limit,
        )
    }
}

impl RpcFilter for Chain {
    fn new_filter(&self, filter: Filter) -> usize {
        let filterdb = self.filter_db();
        let id = filterdb.try_lock().unwrap().gen_id();
        let block_number = self.get_current_height();
        filterdb.try_lock().unwrap().gen_logs_filter(id, filter);
        filterdb
            .try_lock()
            .unwrap()
            .gen_block_filter(id, block_number);
        drop(filterdb);
        id
    }

    fn new_block_filter(&self) -> usize {
        let filterdb = self.filter_db();
        let block_number = self.get_current_height();
        let id = filterdb.try_lock().unwrap().gen_id();
        filterdb
            .try_lock()
            .unwrap()
            .gen_block_filter(id, block_number);
        drop(filterdb);
        id
    }

    fn get_filter_changes(&self, id: usize) -> Option<FilterChanges> {
        let filterdb = self.filter_db();
        let mut changes = Some(FilterChanges::Empty);
        let current_number = self.get_current_height();
        let block_filter = self.get_block_filter(id);

        if !filterdb.try_lock().unwrap().is_filter(id) {
            drop(filterdb);
            return changes;
        }

        // Check the logs
        if let Some(filter) = filterdb.try_lock().unwrap().get_logs_filter(id) {
            trace!("Into filter changes: logs");
            changes = Some(FilterChanges::Logs(
                self.get_logs_with_filter(filter.clone(), block_filter),
            ));
        };

        // Check the block
        if filterdb.try_lock().unwrap().is_block_filter(id) {
            trace!("Into filter changes: block");
            // Return hashes from next block after the filter was created to the current block.
            let hashes = ((block_filter + 1)..=current_number)
                .filter_map(|_id| self.block_hash_by_height(_id))
                .collect::<Vec<H256>>();
            trace!("Block filter changes: {:?}", hashes);
            changes = Some(FilterChanges::Hashes(
                hashes.into_iter().map(Into::into).collect(),
            ));
        };

        // Update the block filter: use the current number
        filterdb
            .try_lock()
            .unwrap()
            .gen_block_filter(id, current_number + 1);
        drop(filterdb);
        changes
    }

    fn get_filter_logs(&self, id: usize) -> Option<Vec<Log>> {
        let filterdb = self.filter_db();
        let block_filter = self.get_block_filter(id);
        let logs = match filterdb.try_lock().unwrap().get_logs_filter(id) {
            Some(filter) => Some(self.get_logs_with_filter(filter.clone(), block_filter)),
            _ => None,
        };
        drop(filterdb);
        logs
    }

    fn uninstall_filter(&self, id: usize) -> bool {
        let filterdb = self.filter_db();
        let uninstall_ok = filterdb.try_lock().unwrap().uninstall(id);
        drop(filterdb);
        uninstall_ok
    }
}

/// Split the logs with the limit filter.
fn split_logs(mut logs: Vec<Log>, limit: Option<usize>) -> Vec<Log> {
    let len = logs.len();
    match limit {
        Some(limit) if len >= limit => logs.split_off(len - limit),
        _ => logs,
    }
}
<|endoftext|>