<commit_before>extern crate rs6502;

use rs6502::*;

#[test]
fn can_instantiate_cpu() {
    let cpu = Cpu::new();

    assert!(0 == 0);
}

#[test]
fn can_load_code_segment_into_memory() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);

    let memory_sum: u32 = cpu.memory.iter().map(|n| *n as u32).sum();
    assert_eq!(46, memory_sum);
}

#[test]
fn can_load_code_segment_at_default_address() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);

    assert_eq!(0x0D, cpu.memory.read_byte(0xC003));
    assert_eq!(0x0C, cpu.memory.read_byte(0xC002));
    assert_eq!(0x0B, cpu.memory.read_byte(0xC001));
    assert_eq!(0x0A, cpu.memory.read_byte(0xC000));
}

#[test]
fn can_load_code_segment_at_specific_address() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], 0xF000);

    assert_eq!(0x0D, cpu.memory.read_byte(0xF003));
    assert_eq!(0x0C, cpu.memory.read_byte(0xF002));
    assert_eq!(0x0B, cpu.memory.read_byte(0xF001));
    assert_eq!(0x0A, cpu.memory.read_byte(0xF000));
}

#[test]
fn errors_when_code_segment_extends_past_memory_bounds() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    let load_result = cpu.load(&fake_code[..], 0xFFFD);

    assert_eq!(Err(CpuError::code_segment_out_of_range(0xFFFD)),
               load_result);
}

#[test]
fn errors_on_unknown_opcode() {
    let fake_code = vec![0xC3];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);
    let step_result: CpuStepResult = cpu.step();

    assert_eq!(Err(CpuError::unknown_opcode(0xC000, 0xC3)), step_result);// This is the unofficial DCP (d,X) opcode
}

#[test]
fn can_get_operand_from_opcode() {
    let fake_code = vec![0xC3];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);
    let step_result: CpuStepResult = cpu.step();
}

#[test]
fn adc_can_set_decimal_flag() {
    let code = vec![0xF8];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step();

    assert_eq!(true, cpu.flags.decimal);
}

#[test]
fn adc_can_disable_decimal_flag() {
    let code = vec![0xD8];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step();

    assert_eq!(false, cpu.flags.decimal);
}

#[test]
fn adc_can_add_basic_numbers() {
    let code = vec![0xA9, 0x05, 0x69, 0x03];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(8, cpu.registers.A);
}

#[test]
fn adc_can_add_basic_numbers_set_carry_and_wrap_around() {
    let code = vec![0xA9, 0xFD, 0x69, 0x05];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(2, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn adc_can_add_numbers_in_binary_coded_decimal() {
    let code = vec![0xF8, 0xA9, 0x05, 0x69, 0x05];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(true, cpu.flags.decimal);
    assert_eq!(0x10, cpu.registers.A);
}

#[test]
fn adc_can_add_numbers_in_binary_coded_decimal_and_set_carry() {
    let code = vec![0xF8, 0xA9, 0x95, 0x69, 0x10];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(true, cpu.flags.carry);
    assert_eq!(true, cpu.flags.decimal);
    assert_eq!(0x05, cpu.registers.A);
}

#[test]
fn sta_can_store_bytes_in_memory() {
    let code = vec![0xA9, 0x20, 0x8D, 0x00, 0x20];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x20, cpu.registers.A);
    assert_eq!(0x20, cpu.memory[0x2000]);
}

#[test]
fn and_can_apply_logical_and_operation() {
    // Load 255 into A and mask it against 0x0F
    let code = vec![0xA9, 0xFF, 0x29, 0x0F];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x0F, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn and_can_apply_logical_and_operation_and_set_sign_flag() {
    // Load 2 into the A register and shift it left
    let code = vec![0xA9, 0x02, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x04, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn asl_can_shift_bits_left() {
    let code = vec![0xA9, 0x02, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x04, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn asl_shifts_last_bit_into_carry() {
    let code = vec![0xA9, 0x80, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x00, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn bcc_can_jump_forward() {
    let code = vec![0xA9, 0xFE, 0x69, 0x01, 0x90, 0x03, 0xA9, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(false, cpu.flags.carry);
    assert_eq!(0xC009, cpu.registers.PC);
}

#[test]
fn bcc_can_jump_backward() {
    let code = vec![0xA9, 0xF0, 0x69, 0x01, 0x90, 0xFC];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(50);

    assert_eq!(0x00, cpu.registers.A);
}

#[test]
fn bcs_can_jump_forward() {
    let code = vec![0xA9, 0xFF, 0x69, 0x01, 0xB0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x00, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn beq_can_jump_forward() {
    let code = vec![0xA9, 0xF0, 0x69, 0x10, 0xF0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x00, cpu.registers.A);
}

#[test]
fn bit_can_set_flags_and_preserve_registers() {
    let code = vec![0xA9, 0xF0, 0x24, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(true, cpu.flags.zero);
    assert_eq!(0xF0, cpu.registers.A);  // Preserves A
}

#[test]
fn bit_can_set_overflow_flag() {
    let code = vec![0xA9, 0xF0, 0x85, 0x44, 0x24, 0x44];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(false, cpu.flags.zero);
    assert_eq!(true, cpu.flags.overflow);
    assert_eq!(true, cpu.flags.sign);
    assert_eq!(0xF0, cpu.registers.A);  // Preserves A
}

#[test]
fn bmi_can_jump_forward() {
    let code = vec![0xA9, 0x7F, 0x69, 0x01, 0x30, 0x03, 0xA9, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x80, cpu.registers.A);
    assert_eq!(true, cpu.flags.sign);
}

#[test]
fn bne_jumps_on_non_zero() {
    let code = vec![0xA9, 0xFE, 0x69, 0x01, 0xD0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(false, cpu.flags.zero);
}

#[test]
fn bne_does_not_jump_on_zero() {
    let code = vec![0xA9, 0xFF, 0x69, 0x01, 0xD0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xAA, cpu.registers.A);
}

#[test]
fn bpl_does_not_jump_on_sign_set() {
    let code = vec![0xA9, 0xFE, 0x10, 0x03, 0xA9, 0xF3];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xF3, cpu.registers.A);
    assert_eq!(true, cpu.flags.sign);
}

#[test]
fn bpl_does_jump_on_sign_not_set() {
    let code = vec![0xA9, 0x0E, 0x10, 0x03, 0xA9, 0xF3];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x0E, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn brk_does_store_pc_and_status_flags_on_stack() {
    let code = vec![0xA9, 0x0E, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xC0, cpu.memory[0x1FE]);
    assert_eq!(0x03, cpu.memory[0x1FD]);
}

#[test]
fn bvc_does_not_jump_on_overflow_set() {
    let code = vec![0xA9, 0x7F, 0x69, 0x01, 0x50, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(true, cpu.flags.overflow);
}

#[test]
fn bvc_does_jump_on_overflow_clear() {
    let code = vec![0xA9, 0x7E, 0x69, 0x01, 0x50, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x7F, cpu.registers.A);
    assert_eq!(false, cpu.flags.overflow);
}<commit_msg>Add tests for BVS instruction<commit_after>extern crate rs6502;

use rs6502::*;

#[test]
fn can_instantiate_cpu() {
    let cpu = Cpu::new();

    assert!(0 == 0);
}

#[test]
fn can_load_code_segment_into_memory() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);

    let memory_sum: u32 = cpu.memory.iter().map(|n| *n as u32).sum();
    assert_eq!(46, memory_sum);
}

#[test]
fn can_load_code_segment_at_default_address() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);

    assert_eq!(0x0D, cpu.memory.read_byte(0xC003));
    assert_eq!(0x0C, cpu.memory.read_byte(0xC002));
    assert_eq!(0x0B, cpu.memory.read_byte(0xC001));
    assert_eq!(0x0A, cpu.memory.read_byte(0xC000));
}

#[test]
fn can_load_code_segment_at_specific_address() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], 0xF000);

    assert_eq!(0x0D, cpu.memory.read_byte(0xF003));
    assert_eq!(0x0C, cpu.memory.read_byte(0xF002));
    assert_eq!(0x0B, cpu.memory.read_byte(0xF001));
    assert_eq!(0x0A, cpu.memory.read_byte(0xF000));
}

#[test]
fn errors_when_code_segment_extends_past_memory_bounds() {
    let fake_code = vec![0x0A, 0x0B, 0x0C, 0x0D];
    let mut cpu = Cpu::new();
    let load_result = cpu.load(&fake_code[..], 0xFFFD);

    assert_eq!(Err(CpuError::code_segment_out_of_range(0xFFFD)),
               load_result);
}

#[test]
fn errors_on_unknown_opcode() {
    let fake_code = vec![0xC3];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);
    let step_result: CpuStepResult = cpu.step();

    assert_eq!(Err(CpuError::unknown_opcode(0xC000, 0xC3)), step_result);// This is the unofficial DCP (d,X) opcode
}

#[test]
fn can_get_operand_from_opcode() {
    let fake_code = vec![0xC3];
    let mut cpu = Cpu::new();
    cpu.load(&fake_code[..], None);
    let step_result: CpuStepResult = cpu.step();
}

#[test]
fn adc_can_set_decimal_flag() {
    let code = vec![0xF8];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step();

    assert_eq!(true, cpu.flags.decimal);
}

#[test]
fn adc_can_disable_decimal_flag() {
    let code = vec![0xD8];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step();

    assert_eq!(false, cpu.flags.decimal);
}

#[test]
fn adc_can_add_basic_numbers() {
    let code = vec![0xA9, 0x05, 0x69, 0x03];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(8, cpu.registers.A);
}

#[test]
fn adc_can_add_basic_numbers_set_carry_and_wrap_around() {
    let code = vec![0xA9, 0xFD, 0x69, 0x05];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(2, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn adc_can_add_numbers_in_binary_coded_decimal() {
    let code = vec![0xF8, 0xA9, 0x05, 0x69, 0x05];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(true, cpu.flags.decimal);
    assert_eq!(0x10, cpu.registers.A);
}

#[test]
fn adc_can_add_numbers_in_binary_coded_decimal_and_set_carry() {
    let code = vec![0xF8, 0xA9, 0x95, 0x69, 0x10];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(true, cpu.flags.carry);
    assert_eq!(true, cpu.flags.decimal);
    assert_eq!(0x05, cpu.registers.A);
}

#[test]
fn sta_can_store_bytes_in_memory() {
    let code = vec![0xA9, 0x20, 0x8D, 0x00, 0x20];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x20, cpu.registers.A);
    assert_eq!(0x20, cpu.memory[0x2000]);
}

#[test]
fn and_can_apply_logical_and_operation() {
    // Load 255 into A and mask it against 0x0F
    let code = vec![0xA9, 0xFF, 0x29, 0x0F];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x0F, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn and_can_apply_logical_and_operation_and_set_sign_flag() {
    // Load 2 into the A register and shift it left
    let code = vec![0xA9, 0x02, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x04, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn asl_can_shift_bits_left() {
    let code = vec![0xA9, 0x02, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x04, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn asl_shifts_last_bit_into_carry() {
    let code = vec![0xA9, 0x80, 0x0A];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(2);

    assert_eq!(0x00, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn bcc_can_jump_forward() {
    let code = vec![0xA9, 0xFE, 0x69, 0x01, 0x90, 0x03, 0xA9, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(3);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(false, cpu.flags.carry);
    assert_eq!(0xC009, cpu.registers.PC);
}

#[test]
fn bcc_can_jump_backward() {
    let code = vec![0xA9, 0xF0, 0x69, 0x01, 0x90, 0xFC];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(50);

    assert_eq!(0x00, cpu.registers.A);
}

#[test]
fn bcs_can_jump_forward() {
    let code = vec![0xA9, 0xFF, 0x69, 0x01, 0xB0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x00, cpu.registers.A);
    assert_eq!(true, cpu.flags.carry);
}

#[test]
fn beq_can_jump_forward() {
    let code = vec![0xA9, 0xF0, 0x69, 0x10, 0xF0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x00, cpu.registers.A);
}

#[test]
fn bit_can_set_flags_and_preserve_registers() {
    let code = vec![0xA9, 0xF0, 0x24, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(true, cpu.flags.zero);
    assert_eq!(0xF0, cpu.registers.A);  // Preserves A
}

#[test]
fn bit_can_set_overflow_flag() {
    let code = vec![0xA9, 0xF0, 0x85, 0x44, 0x24, 0x44];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(false, cpu.flags.zero);
    assert_eq!(true, cpu.flags.overflow);
    assert_eq!(true, cpu.flags.sign);
    assert_eq!(0xF0, cpu.registers.A);  // Preserves A
}

#[test]
fn bmi_can_jump_forward() {
    let code = vec![0xA9, 0x7F, 0x69, 0x01, 0x30, 0x03, 0xA9, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x80, cpu.registers.A);
    assert_eq!(true, cpu.flags.sign);
}

#[test]
fn bne_jumps_on_non_zero() {
    let code = vec![0xA9, 0xFE, 0x69, 0x01, 0xD0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(false, cpu.flags.zero);
}

#[test]
fn bne_does_not_jump_on_zero() {
    let code = vec![0xA9, 0xFF, 0x69, 0x01, 0xD0, 0x03, 0xA9, 0xAA];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xAA, cpu.registers.A);
}

#[test]
fn bpl_does_not_jump_on_sign_set() {
    let code = vec![0xA9, 0xFE, 0x10, 0x03, 0xA9, 0xF3];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xF3, cpu.registers.A);
    assert_eq!(true, cpu.flags.sign);
}

#[test]
fn bpl_does_jump_on_sign_not_set() {
    let code = vec![0xA9, 0x0E, 0x10, 0x03, 0xA9, 0xF3];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x0E, cpu.registers.A);
    assert_eq!(false, cpu.flags.sign);
}

#[test]
fn brk_does_store_pc_and_status_flags_on_stack() {
    let code = vec![0xA9, 0x0E, 0x00];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xC0, cpu.memory[0x1FE]);
    assert_eq!(0x03, cpu.memory[0x1FD]);
}

#[test]
fn bvc_does_not_jump_on_overflow_set() {
    let code = vec![0xA9, 0x7F, 0x69, 0x01, 0x50, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(true, cpu.flags.overflow);
}

#[test]
fn bvc_does_jump_on_overflow_clear() {
    let code = vec![0xA9, 0x7E, 0x69, 0x01, 0x50, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x7F, cpu.registers.A);
    assert_eq!(false, cpu.flags.overflow);
}

#[test]
fn bvs_does_not_jump_on_overflow_clear() {
    let code = vec![0xA9, 0x7E, 0x69, 0x01, 0x70, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0xFF, cpu.registers.A);
    assert_eq!(false, cpu.flags.overflow);
}

#[test]
fn bvs_does_jump_on_overflow_set() {
    let code = vec![0xA9, 0x7F, 0x69, 0x01, 0x70, 0x03, 0xA9, 0xFF];
    let mut cpu = Cpu::new();
    cpu.load(&code[..], None);

    cpu.step_n(10);

    assert_eq!(0x80, cpu.registers.A);
    assert_eq!(true, cpu.flags.overflow);
}<|endoftext|>


<commit_before>use std::process::Command;
use std::fs;

fn test_output(result: &str, expected: &str) {
    assert!(Command::new("cmp")
            .arg(result)
            .arg(expected)
            .spawn().unwrap().wait().unwrap().success());
    fs::remove_file(result).unwrap();
}

fn run_cmd(cmd: &str, prefix: Option<&str>) -> bool {
    let build = if cfg!(debug_assertions) { "debug" } else { "release" };

    println!("{}", cmd);

    Command::new("bash")
            .arg("-c")
            .arg(format!(
                "{prefix}target/{build}/{cmd}",
                prefix=prefix.unwrap_or(""),
                cmd=cmd,
                build=build
            ))
            .spawn().unwrap().wait().unwrap().success()
}

fn run_exp(dataset: &str, codebook: &str, params: &str) -> bool {
    run_cmd(
        &format!(
            "merfishtools -v exp -t 1 {codebook} --estimate {est} --stats {stats} {params} < {raw} > {cdf}",
            codebook=format!("tests/codebook/{}.txt", codebook),
            raw=format!("tests/data/{}.txt", dataset),
            cdf=format!("tests/results/{}.txt", dataset),
            est=format!("tests/results/{}.est.txt", dataset),
            stats=format!("tests/results/{}.stats.txt", dataset),
            params=params
        ),
        None
    )
}

#[test]
fn test_exp_mhd4_cell0() {
    assert!(run_exp("140genesData.1.cell0", "140genesData.1", "--p0 0.016 0.015 0.021 0.015 0.051 0.065 0.062 0.026 0.019 0.021 0.023 0.034 0.02 0.029 0.07 0.027 --p1 0.05 0.04 0.065 0.052 0.177 0.135 0.155 0.07 0.053 0.057 0.066 0.09 0.069 0.07 0.145 0.103"));
    //assert!(run_exp("140genesData.1.cell0", "140genesData.1", "--p0 0.1 --p1 0.6"));
}

#[test]
fn test_exp_mhd4_thbs_bias() {
    assert!(run_exp("140genesData.1.cell24", "140genesData.1", "--p0 0.016 0.015 0.021 0.015 0.051 0.065 0.062 0.026 0.019 0.021 0.023 0.034 0.02 0.029 0.07 0.027 --p1 0.05 0.04 0.065 0.052 0.177 0.135 0.155 0.07 0.053 0.057 0.066 0.09 0.069 0.07 0.145 0.103"));
}

#[test]
fn test_exp_mhd2() {
    assert!(run_exp("1001genesData.3.cell0", "1001genesData", ""));
}

/// This is a longrunning test that should be activated only on purpose.
fn test_exp_mhd2_8() {
    assert!(run_exp("simulated-MHD2-8.25.all", "simulated-MHD2-8", "--p0 0.005 --p1 0.01"));
}

#[test]
fn test_exp_mhd4_sim() {
    assert!(run_exp("simulated-MHD4.35.1", "simulated-MHD4", "--p0 0.04 --p1 0.1"));
}

#[test]
fn test_estimate_error_rates_real() {
    assert!(run_cmd(
        "merfishtools -v est-error-rates tests/codebook/140genesData.1.txt > tests/results/140genesData.error-rates.tsv",
        Some("grep -P '^3\\t' tests/data/140genesData.readouts.txt | cut -f2,3,4 | ")
    ));
}

#[test]
fn test_estimate_error_rates_simulated() {
    assert!(run_cmd(
        "merfishtools -v est-error-rates tests/codebook/140genesData.1.txt < tests/data/simulated-MHD4.35.readouts.txt > tests/results/simulated-MHD4.35.error-rates.tsv",
        None
    ));
}
<commit_msg>Fixed error rates.<commit_after>use std::process::Command;
use std::fs;

fn test_output(result: &str, expected: &str) {
    assert!(Command::new("cmp")
            .arg(result)
            .arg(expected)
            .spawn().unwrap().wait().unwrap().success());
    fs::remove_file(result).unwrap();
}

fn run_cmd(cmd: &str, prefix: Option<&str>) -> bool {
    let build = if cfg!(debug_assertions) { "debug" } else { "release" };

    println!("{}", cmd);

    Command::new("bash")
            .arg("-c")
            .arg(format!(
                "{prefix}target/{build}/{cmd}",
                prefix=prefix.unwrap_or(""),
                cmd=cmd,
                build=build
            ))
            .spawn().unwrap().wait().unwrap().success()
}

fn run_exp(dataset: &str, codebook: &str, params: &str) -> bool {
    run_cmd(
        &format!(
            "merfishtools -v exp -t 1 {codebook} --estimate {est} --stats {stats} {params} < {raw} > {cdf}",
            codebook=format!("tests/codebook/{}.txt", codebook),
            raw=format!("tests/data/{}.txt", dataset),
            cdf=format!("tests/results/{}.txt", dataset),
            est=format!("tests/results/{}.est.txt", dataset),
            stats=format!("tests/results/{}.stats.txt", dataset),
            params=params
        ),
        None
    )
}

#[test]
fn test_exp_mhd4_cell0() {
    assert!(run_exp("140genesData.1.cell0", "140genesData.1", "--p0 0.01258132102084181 0.013360438064572837 0.019941605919281354 0.016581551788592854 0.22860571489439577 0.22269171721737882 0.2073701214077358 0.02231576033300221 0.017689514891827997 0.02054328192011981 0.02101414200654054 0.03938992313394839 0.023995219270935196 0.03557793748730565 0.17007370565085334 0.03569158826173908 --p1 0.5912118122316706 0.5063863207089024 0.42338772914919315 0.43282699715628403 0.17843503208026912 0.18766044306365068 0.1989449380102563 0.42262699427628003 0.4930707032134955 0.4034046655581933 0.4610572590564446 0.30553305919201673 0.2931349183480258 0.2634814096058736 0.16619112507647252 0.36410395236487053"));
    //assert!(run_exp("140genesData.1.cell0", "140genesData.1", "--p0 0.016 0.015 0.021 0.015 0.051 0.065 0.062 0.026 0.019 0.021 0.023 0.034 0.02 0.029 0.07 0.027 --p1 0.05 0.04 0.065 0.052 0.177 0.135 0.155 0.07 0.053 0.057 0.066 0.09 0.069 0.07 0.145 0.103"));
    //assert!(run_exp("140genesData.1.cell0", "140genesData.1", "--p0 0.1 --p1 0.6"));
}

#[test]
fn test_exp_mhd4_thbs_bias() {
    assert!(run_exp("140genesData.1.cell24", "140genesData.1", "--p0 0.016 0.015 0.021 0.015 0.051 0.065 0.062 0.026 0.019 0.021 0.023 0.034 0.02 0.029 0.07 0.027 --p1 0.05 0.04 0.065 0.052 0.177 0.135 0.155 0.07 0.053 0.057 0.066 0.09 0.069 0.07 0.145 0.103"));
}

#[test]
fn test_exp_mhd2() {
    assert!(run_exp("1001genesData.3.cell0", "1001genesData", ""));
}

/// This is a longrunning test that should be activated only on purpose.
fn test_exp_mhd2_8() {
    assert!(run_exp("simulated-MHD2-8.25.all", "simulated-MHD2-8", "--p0 0.005 --p1 0.01"));
}

#[test]
fn test_exp_mhd4_sim() {
    assert!(run_exp("simulated-MHD4.35.1", "simulated-MHD4", "--p0 0.04 --p1 0.1"));
}

#[test]
fn test_estimate_error_rates_real() {
    assert!(run_cmd(
        "merfishtools -v est-error-rates tests/codebook/140genesData.1.txt > tests/results/140genesData.error-rates.tsv",
        Some("grep -P '^3\\t' tests/data/140genesData.readouts.txt | cut -f2,3,4 | ")
    ));
}

#[test]
fn test_estimate_error_rates_simulated() {
    assert!(run_cmd(
        "merfishtools -v est-error-rates tests/codebook/140genesData.1.txt < tests/data/simulated-MHD4.35.readouts.txt > tests/results/simulated-MHD4.35.error-rates.tsv",
        None
    ));
}
<|endoftext|>


<commit_before>extern crate opentype;
extern crate truetype;

use opentype::glyph_transformation::Scripts;
use std::fs::File;
use truetype::{Tag, Value};

mod fixture;

use fixture::Fixture;

macro_rules! ok(($result:expr) => ($result.unwrap()));

macro_rules! setup(
    ($table:expr) => (setup(Fixture::CFF, Some($table)));
    ($fixture:ident, $table:expr) => (setup(Fixture::$fixture, Some($table)));
);

macro_rules! tags(($($name:expr),*) => (vec![$(Tag(*$name)),*]));

#[test]
fn glyph_definition() {
    use opentype::GlyphDefinition;

    let GlyphDefinition { .. } = ok!(Value::read(&mut setup!(TTF, "GDEF")));
}

#[test]
fn glyph_positioning_features() {
    use opentype::GlyphPositioning;

    let GlyphPositioning { features, .. } = ok!(Value::read(&mut setup!("GPOS")));
    let tags = features.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"kern", b"kern", b"kern", b"kern", b"kern",
                           b"size", b"size", b"size", b"size", b"size"]);
    let lookups = features.records.iter().map(|record| record.lookup_count).collect::<Vec<_>>();
    assert_eq!(lookups, &[1, 1, 1, 1, 1, 0, 0, 0, 0, 0]);
}

#[test]
fn glyph_positioning_lookups() {
    use opentype::GlyphPositioning;
    use opentype::glyph_positioning::{PairAdjustment, Table};

    let GlyphPositioning { lookups, .. } = ok!(Value::read(&mut setup!("GPOS")));
    assert_eq!(lookups.records.len(), 1);
    let record = &lookups.records[0];
    assert!(record.mark_filtering_set.is_none());
    assert_eq!(record.tables.len(), 2);
    match &record.tables[0] {
        &Table::PairAdjustment(PairAdjustment::Format1(ref table)) => {
            assert_eq!(table.set_count, 65);
        },
        _ => unreachable!(),
    }
    match &record.tables[1] {
        &Table::PairAdjustment(PairAdjustment::Format2(ref table)) => {
            assert_eq!(table.class1_count, 99);
            assert_eq!(table.class2_count, 95);
        },
        _ => unreachable!(),
    }
}

#[test]
fn glyph_positioning_scripts() {
    use opentype::GlyphPositioning;
    scripts(&ok!(GlyphPositioning::read(&mut setup!("GPOS"))).scripts);
}

#[test]
fn glyph_substitution_features() {
    use opentype::GlyphSubstitution;

    let GlyphSubstitution { features, .. } = ok!(Value::read(&mut setup!("GSUB")));
    let tags = features.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"aalt", b"aalt", b"aalt", b"aalt", b"aalt",
                           b"case", b"case", b"case", b"case", b"case",
                           b"dnom", b"dnom", b"dnom", b"dnom", b"dnom",
                           b"frac", b"frac", b"frac", b"frac", b"frac",
                           b"liga", b"liga", b"liga", b"liga", b"liga",
                           b"lnum", b"lnum", b"lnum", b"lnum", b"lnum",
                           b"locl", b"locl", b"locl",
                           b"numr", b"numr", b"numr", b"numr", b"numr",
                           b"onum", b"onum", b"onum", b"onum", b"onum",
                           b"ordn", b"ordn", b"ordn", b"ordn", b"ordn",
                           b"pnum", b"pnum", b"pnum", b"pnum", b"pnum",
                           b"sinf", b"sinf", b"sinf", b"sinf", b"sinf",
                           b"subs", b"subs", b"subs", b"subs", b"subs",
                           b"sups", b"sups", b"sups", b"sups", b"sups",
                           b"tnum", b"tnum", b"tnum", b"tnum", b"tnum",
                           b"zero", b"zero", b"zero", b"zero", b"zero"]);
    let lookups = features.records.iter().map(|record| record.lookup_count).collect::<Vec<_>>();
    assert_eq!(lookups, vec![2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 1, 1, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1]);
}

#[test]
fn glyph_substitution_lookups() {
    use opentype::GlyphSubstitution;
    use opentype::glyph_substitution::{SingleSubstitution, Table};

    let GlyphSubstitution { lookups, .. } = ok!(Value::read(&mut setup!("GSUB")));
    let kinds = lookups.records.iter().map(|record| record.kind).collect::<Vec<_>>();
    assert_eq!(kinds, &[1, 3, 1, 1, 1, 1, 1, 6, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4, 1]);
    let record = &lookups.records[0];
    assert_eq!(record.tables.len(), 1);
    match &record.tables[0] {
        &Table::SingleSubstitution(SingleSubstitution::Format2(ref table)) => {
            assert_eq!(table.glyph_count, 61);
        },
        _ => unreachable!(),
    }
    let record = &lookups.records[17];
    assert_eq!(record.tables.len(), 1);
    match &record.tables[0] {
        &Table::LigatureSubstitution(ref table) => {
            assert_eq!(table.set_count, 1);
            let table = &table.sets[0];
            assert_eq!(table.count, 3);
            let table = &table.records[0];
            assert_eq!(table.component_count, 2);
        },
        _ => unreachable!(),
    }
}

#[test]
fn glyph_substitution_scripts() {
    use opentype::GlyphSubstitution;
    scripts(&ok!(GlyphSubstitution::read(&mut setup!("GSUB"))).scripts);
}

fn scripts(scripts: &Scripts) {
    use opentype::glyph_transformation::script::{Language, Script};

    let tags = scripts.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"DFLT", b"latn"]);
    assert!(scripts.get(Script::Default).is_some());
    assert!(scripts.get(Script::Latin).is_some());
    let tags = scripts.records.iter()
                              .map(|record| record.language_headers.iter()
                                                                   .map(|header| header.tag)
                                                                   .collect::<Vec<_>>())
                              .collect::<Vec<_>>();
    assert_eq!(tags, &[vec![], tags![b"AZE ", b"CRT ", b"TRK "]]);
    let record = &scripts.records[0];
    assert!(record.default_language.is_some());
    assert_eq!(record.language_count, 0);
    let record = &scripts.records[1];
    assert_eq!(record.language_count, 3);
    assert!(record.get(Language::Turkish).is_some());
}

fn setup(fixture: Fixture, table: Option<&str>) -> File {
    use std::io::{Seek, SeekFrom};

    let mut file = ok!(File::open(fixture.path()));
    ok!(file.seek(SeekFrom::Start(table.map(|table| fixture.offset(table)).unwrap_or(0))));
    file
}
<commit_msg>Add a test for GDEF<commit_after>extern crate opentype;
extern crate truetype;

use opentype::glyph_transformation::Scripts;
use std::fs::File;
use truetype::{Tag, Value};

mod fixture;

use fixture::Fixture;

macro_rules! ok(($result:expr) => ($result.unwrap()));

macro_rules! setup(
    ($table:expr) => (setup(Fixture::CFF, Some($table)));
    ($fixture:ident, $table:expr) => (setup(Fixture::$fixture, Some($table)));
);

macro_rules! tags(($($name:expr),*) => (vec![$(Tag(*$name)),*]));

#[test]
fn glyph_definition() {
    use opentype::glyph_definition::{GlyphDefinition, Header};
    use opentype::glyph_transformation::Class;

    let table: GlyphDefinition = ok!(Value::read(&mut setup!(TTF, "GDEF")));
    match &table.header {
        &Header::Version1(..) => {},
        _ => unreachable!(),
    }
    match &table.glyph_class {
        &Some(Class::Format2(ref table)) => {
            assert_eq!(table.range_count, 1);
            assert_eq!(table.ranges[0].start, 0);
            assert_eq!(table.ranges[0].end, 937);
        },
        _ => unreachable!(),
    }
    assert!(table.attachments.is_none());
    match &table.ligatures {
        &Some(ref table) => assert_eq!(table.count, 0),
        _ => unreachable!(),
    }
}

#[test]
fn glyph_positioning_features() {
    use opentype::GlyphPositioning;

    let GlyphPositioning { features, .. } = ok!(Value::read(&mut setup!("GPOS")));
    let tags = features.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"kern", b"kern", b"kern", b"kern", b"kern",
                           b"size", b"size", b"size", b"size", b"size"]);
    let lookups = features.records.iter().map(|record| record.lookup_count).collect::<Vec<_>>();
    assert_eq!(lookups, &[1, 1, 1, 1, 1, 0, 0, 0, 0, 0]);
}

#[test]
fn glyph_positioning_lookups() {
    use opentype::GlyphPositioning;
    use opentype::glyph_positioning::{PairAdjustment, Table};

    let GlyphPositioning { lookups, .. } = ok!(Value::read(&mut setup!("GPOS")));
    assert_eq!(lookups.records.len(), 1);
    let record = &lookups.records[0];
    assert!(record.mark_filtering_set.is_none());
    assert_eq!(record.tables.len(), 2);
    match &record.tables[0] {
        &Table::PairAdjustment(PairAdjustment::Format1(ref table)) => {
            assert_eq!(table.set_count, 65);
        },
        _ => unreachable!(),
    }
    match &record.tables[1] {
        &Table::PairAdjustment(PairAdjustment::Format2(ref table)) => {
            assert_eq!(table.class1_count, 99);
            assert_eq!(table.class2_count, 95);
        },
        _ => unreachable!(),
    }
}

#[test]
fn glyph_positioning_scripts() {
    use opentype::GlyphPositioning;
    scripts(&ok!(GlyphPositioning::read(&mut setup!("GPOS"))).scripts);
}

#[test]
fn glyph_substitution_features() {
    use opentype::GlyphSubstitution;

    let GlyphSubstitution { features, .. } = ok!(Value::read(&mut setup!("GSUB")));
    let tags = features.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"aalt", b"aalt", b"aalt", b"aalt", b"aalt",
                           b"case", b"case", b"case", b"case", b"case",
                           b"dnom", b"dnom", b"dnom", b"dnom", b"dnom",
                           b"frac", b"frac", b"frac", b"frac", b"frac",
                           b"liga", b"liga", b"liga", b"liga", b"liga",
                           b"lnum", b"lnum", b"lnum", b"lnum", b"lnum",
                           b"locl", b"locl", b"locl",
                           b"numr", b"numr", b"numr", b"numr", b"numr",
                           b"onum", b"onum", b"onum", b"onum", b"onum",
                           b"ordn", b"ordn", b"ordn", b"ordn", b"ordn",
                           b"pnum", b"pnum", b"pnum", b"pnum", b"pnum",
                           b"sinf", b"sinf", b"sinf", b"sinf", b"sinf",
                           b"subs", b"subs", b"subs", b"subs", b"subs",
                           b"sups", b"sups", b"sups", b"sups", b"sups",
                           b"tnum", b"tnum", b"tnum", b"tnum", b"tnum",
                           b"zero", b"zero", b"zero", b"zero", b"zero"]);
    let lookups = features.records.iter().map(|record| record.lookup_count).collect::<Vec<_>>();
    assert_eq!(lookups, vec![2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 1, 1, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 1,
                             1, 1, 1, 1, 1, 1, 1, 1, 1]);
}

#[test]
fn glyph_substitution_lookups() {
    use opentype::GlyphSubstitution;
    use opentype::glyph_substitution::{SingleSubstitution, Table};

    let GlyphSubstitution { lookups, .. } = ok!(Value::read(&mut setup!("GSUB")));
    let kinds = lookups.records.iter().map(|record| record.kind).collect::<Vec<_>>();
    assert_eq!(kinds, &[1, 3, 1, 1, 1, 1, 1, 6, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4, 1]);
    let record = &lookups.records[0];
    assert_eq!(record.tables.len(), 1);
    match &record.tables[0] {
        &Table::SingleSubstitution(SingleSubstitution::Format2(ref table)) => {
            assert_eq!(table.glyph_count, 61);
        },
        _ => unreachable!(),
    }
    let record = &lookups.records[17];
    assert_eq!(record.tables.len(), 1);
    match &record.tables[0] {
        &Table::LigatureSubstitution(ref table) => {
            assert_eq!(table.set_count, 1);
            let table = &table.sets[0];
            assert_eq!(table.count, 3);
            let table = &table.records[0];
            assert_eq!(table.component_count, 2);
        },
        _ => unreachable!(),
    }
}

#[test]
fn glyph_substitution_scripts() {
    use opentype::GlyphSubstitution;
    scripts(&ok!(GlyphSubstitution::read(&mut setup!("GSUB"))).scripts);
}

fn scripts(scripts: &Scripts) {
    use opentype::glyph_transformation::script::{Language, Script};

    let tags = scripts.headers.iter().map(|header| header.tag).collect::<Vec<_>>();
    assert_eq!(tags, tags![b"DFLT", b"latn"]);
    assert!(scripts.get(Script::Default).is_some());
    assert!(scripts.get(Script::Latin).is_some());
    let tags = scripts.records.iter()
                              .map(|record| record.language_headers.iter()
                                                                   .map(|header| header.tag)
                                                                   .collect::<Vec<_>>())
                              .collect::<Vec<_>>();
    assert_eq!(tags, &[vec![], tags![b"AZE ", b"CRT ", b"TRK "]]);
    let record = &scripts.records[0];
    assert!(record.default_language.is_some());
    assert_eq!(record.language_count, 0);
    let record = &scripts.records[1];
    assert_eq!(record.language_count, 3);
    assert!(record.get(Language::Turkish).is_some());
}

fn setup(fixture: Fixture, table: Option<&str>) -> File {
    use std::io::{Seek, SeekFrom};

    let mut file = ok!(File::open(fixture.path()));
    ok!(file.seek(SeekFrom::Start(table.map(|table| fixture.offset(table)).unwrap_or(0))));
    file
}
<|endoftext|>


<commit_before>use anyhow::{bail, Result};
use heck::{CamelCase, KebabCase, SnakeCase};
use itertools::Itertools;
use proc_macro2::TokenStream;
use proc_macro_error::abort_call_site;
use quote::{format_ident, quote, quote_spanned};
use syn::{
    punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, DeriveInput, Ident,
    Lit, Meta, MetaList, NestedMeta, Variant,
};

// Generate i32 to enum mapping, used to implement glib::translate::FromGlib<i32>, such as:
//   if value == Animal::Goat as i32 {
//       return Animal::Goat;
//   }
fn gen_from_glib(enum_name: &Ident, enum_variants: &Punctuated<Variant, Comma>) -> TokenStream {
    // FIXME: can we express this with a match()?
    let recurse = enum_variants.iter().map(|v| {
        let name = &v.ident;
        quote_spanned! {v.span()=>
            if value == #enum_name::#name as i32 {
                return #enum_name::#name;
            }
        }
    });
    quote! {
        #(#recurse)*
    }
}

// find the #[genum] attribute in @attrs
fn find_genum_meta(attrs: &[Attribute]) -> Result<Option<MetaList>> {
    let meta = match attrs.iter().find(|a| a.path.is_ident("genum")) {
        Some(a) => a.parse_meta(),
        _ => return Ok(None),
    };
    match meta? {
        Meta::List(n) => Ok(Some(n)),
        _ => bail!("wrong meta type"),
    }
}

// parse a single meta like: ident = "value"
fn parse_attribute(meta: &NestedMeta) -> Result<(String, String)> {
    let meta = match &meta {
        NestedMeta::Meta(m) => m,
        _ => bail!("wrong meta type"),
    };
    let meta = match meta {
        Meta::NameValue(n) => n,
        _ => bail!("wrong meta type"),
    };
    let value = match &meta.lit {
        Lit::Str(s) => s.value(),
        _ => bail!("wrong meta type"),
    };

    let ident = match meta.path.get_ident() {
        None => bail!("missing ident"),
        Some(ident) => ident,
    };

    Ok((ident.to_string(), value))
}

#[derive(Debug)]
enum EnumAttribute {
    TypeName(String),
}

fn parse_enum_attribute(meta: &NestedMeta) -> Result<EnumAttribute> {
    let (ident, v) = parse_attribute(meta)?;

    match ident.as_ref() {
        "type_name" => Ok(EnumAttribute::TypeName(v)),
        s => bail!("Unknown enum meta {}", s),
    }
}

#[derive(Debug)]
enum ItemAttribute {
    Name(String),
    Nick(String),
}

fn parse_item_attribute(meta: &NestedMeta) -> Result<ItemAttribute> {
    let (ident, v) = parse_attribute(meta)?;

    match ident.as_ref() {
        "name" => Ok(ItemAttribute::Name(v)),
        "nick" => Ok(ItemAttribute::Nick(v)),
        s => bail!("Unknown item meta {}", s),
    }
}

// Parse optional enum item attributes such as:
// #[genum(name = "My Name", nick = "my-nick")]
fn parse_item_attributes(attrs: &[Attribute]) -> Result<Vec<ItemAttribute>> {
    let meta = find_genum_meta(attrs)?;

    let v = match meta {
        Some(meta) => meta
            .nested
            .iter()
            .map(|m| parse_item_attribute(&m))
            .fold_results(Vec::new(), |mut v, a| {
                v.push(a);
                v
            })?,
        None => Vec::new(),
    };

    Ok(v)
}

// Generate gobject_sys::GEnumValue structs mapping the enum such as:
//     gobject_sys::GEnumValue {
//         value: Animal::Goat as i32,
//         value_name: "Goat\0" as *const _ as *const _,
//         value_nick: "goat\0" as *const _ as *const _,
//     },
fn gen_genum_values(
    enum_name: &Ident,
    enum_variants: &Punctuated<Variant, Comma>,
) -> (TokenStream, usize) {
    // start at one as GEnumValue array is null-terminated
    let mut n = 1;
    let recurse = enum_variants.iter().map(|v| {
        let name = &v.ident;
        let mut value_name = name.to_string().to_camel_case();
        let mut value_nick = name.to_string().to_kebab_case();

        let attrs = parse_item_attributes(&v.attrs);
        let attrs = match attrs {
            Ok(attrs) => attrs,
            Err(e) => abort_call_site!(
                "{}: GEnum enum supports only the following optional attributes: #[genum(name = \"The Cat\", nick = \"chat\")]",
                e
            ),
        };

        attrs.into_iter().for_each(|attr|
            match attr {
                ItemAttribute::Name(n) => value_name = n,
                ItemAttribute::Nick(n) => value_nick = n,
            }
        );

        let value_name = format!("{}\0", value_name);
        let value_nick = format!("{}\0", value_nick);

        n += 1;
        quote_spanned! {v.span()=>
            gobject_sys::GEnumValue {
                value: #enum_name::#name as i32,
                value_name: #value_name as *const _ as *const _,
                value_nick: #value_nick as *const _ as *const _,
            },
        }
    });
    (
        quote! {
            #(#recurse)*
        },
        n,
    )
}

// Parse enum attribute such as:
// #[genum(type_name = "TestAnimalType")]
fn parse_type_name(input: &DeriveInput) -> Result<String> {
    let meta = match find_genum_meta(&input.attrs)? {
        Some(meta) => meta,
        _ => bail!("Missing 'genum' attribute"),
    };

    let meta = match meta.nested.first() {
        Some(meta) => meta,
        _ => bail!("Missing meta 'type_name'"),
    };

    match parse_enum_attribute(&meta)? {
        EnumAttribute::TypeName(n) => Ok(n),
    }
}

pub fn impl_genum(input: &syn::DeriveInput) -> TokenStream {
    let name = &input.ident;

    let enum_variants = match input.data {
        Data::Enum(ref e) => &e.variants,
        _ => abort_call_site!("GEnum only supports enums"),
    };

    let gtype_name = match parse_type_name(&input) {
        Ok(v) => v,
        Err(e) => abort_call_site!(
            "{}: derive(GEnum) requires #[genum(type_name = \"EnumTypeName\")]",
            e
        ),
    };
    let get_type = format_ident!("{}_get_type", name.to_string().to_snake_case());
    let from_glib = gen_from_glib(name, enum_variants);
    let (genum_values, nb_genum_values) = gen_genum_values(name, enum_variants);

    quote! {
        impl glib::translate::ToGlib for #name {
            type GlibType = i32;

            fn to_glib(&self) -> i32 {
                *self as i32
            }
        }

        impl glib::translate::FromGlib<i32> for #name {
            fn from_glib(value: i32) -> Self {
                #from_glib
                unreachable!();
            }
        }

        impl<'a> glib::value::FromValueOptional<'a> for #name {
            unsafe fn from_value_optional(value: &glib::Value) -> Option<Self> {
                Some(glib::value::FromValue::from_value(value))
            }
        }

        impl<'a> glib::value::FromValue<'a> for #name {
            unsafe fn from_value(value: &glib::Value) -> Self {
                glib::translate::from_glib(
                    gobject_sys::g_value_get_enum(
                        glib::translate::ToGlibPtr::to_glib_none(value).0))
            }
        }

        impl glib::value::SetValue for #name {
            unsafe fn set_value(value: &mut glib::Value, this: &Self) {
                gobject_sys::g_value_set_enum(
                    glib::translate::ToGlibPtrMut::to_glib_none_mut(value).0,
                    glib::translate::ToGlib::to_glib(this))
            }
        }

        impl StaticType for #name {
            fn static_type() -> glib::Type {
                #get_type()
            }
        }

        fn #get_type() -> glib::Type {
            static ONCE: std::sync::Once = std::sync::Once::new();
            static mut TYPE: glib::Type = glib::Type::Invalid;

            ONCE.call_once(|| {
                static mut VALUES: [gobject_sys::GEnumValue; #nb_genum_values] = [
                    #genum_values
                    gobject_sys::GEnumValue {
                        value: 0,
                        value_name: std::ptr::null(),
                        value_nick: std::ptr::null(),
                    },
                ];

                let name = std::ffi::CString::new(#gtype_name).expect("CString::new failed");
                unsafe {
                    let type_ = gobject_sys::g_enum_register_static(name.as_ptr(), VALUES.as_ptr());
                    TYPE = glib::translate::from_glib(type_);
                }
            });

            unsafe {
                assert_ne!(TYPE, glib::Type::Invalid);
                TYPE
            }
        }
    }
}
<commit_msg>glib-macros: genum_derive: add missing copyright header<commit_after>// Copyright 2020, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use anyhow::{bail, Result};
use heck::{CamelCase, KebabCase, SnakeCase};
use itertools::Itertools;
use proc_macro2::TokenStream;
use proc_macro_error::abort_call_site;
use quote::{format_ident, quote, quote_spanned};
use syn::{
    punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, DeriveInput, Ident,
    Lit, Meta, MetaList, NestedMeta, Variant,
};

// Generate i32 to enum mapping, used to implement glib::translate::FromGlib<i32>, such as:
//   if value == Animal::Goat as i32 {
//       return Animal::Goat;
//   }
fn gen_from_glib(enum_name: &Ident, enum_variants: &Punctuated<Variant, Comma>) -> TokenStream {
    // FIXME: can we express this with a match()?
    let recurse = enum_variants.iter().map(|v| {
        let name = &v.ident;
        quote_spanned! {v.span()=>
            if value == #enum_name::#name as i32 {
                return #enum_name::#name;
            }
        }
    });
    quote! {
        #(#recurse)*
    }
}

// find the #[genum] attribute in @attrs
fn find_genum_meta(attrs: &[Attribute]) -> Result<Option<MetaList>> {
    let meta = match attrs.iter().find(|a| a.path.is_ident("genum")) {
        Some(a) => a.parse_meta(),
        _ => return Ok(None),
    };
    match meta? {
        Meta::List(n) => Ok(Some(n)),
        _ => bail!("wrong meta type"),
    }
}

// parse a single meta like: ident = "value"
fn parse_attribute(meta: &NestedMeta) -> Result<(String, String)> {
    let meta = match &meta {
        NestedMeta::Meta(m) => m,
        _ => bail!("wrong meta type"),
    };
    let meta = match meta {
        Meta::NameValue(n) => n,
        _ => bail!("wrong meta type"),
    };
    let value = match &meta.lit {
        Lit::Str(s) => s.value(),
        _ => bail!("wrong meta type"),
    };

    let ident = match meta.path.get_ident() {
        None => bail!("missing ident"),
        Some(ident) => ident,
    };

    Ok((ident.to_string(), value))
}

#[derive(Debug)]
enum EnumAttribute {
    TypeName(String),
}

fn parse_enum_attribute(meta: &NestedMeta) -> Result<EnumAttribute> {
    let (ident, v) = parse_attribute(meta)?;

    match ident.as_ref() {
        "type_name" => Ok(EnumAttribute::TypeName(v)),
        s => bail!("Unknown enum meta {}", s),
    }
}

#[derive(Debug)]
enum ItemAttribute {
    Name(String),
    Nick(String),
}

fn parse_item_attribute(meta: &NestedMeta) -> Result<ItemAttribute> {
    let (ident, v) = parse_attribute(meta)?;

    match ident.as_ref() {
        "name" => Ok(ItemAttribute::Name(v)),
        "nick" => Ok(ItemAttribute::Nick(v)),
        s => bail!("Unknown item meta {}", s),
    }
}

// Parse optional enum item attributes such as:
// #[genum(name = "My Name", nick = "my-nick")]
fn parse_item_attributes(attrs: &[Attribute]) -> Result<Vec<ItemAttribute>> {
    let meta = find_genum_meta(attrs)?;

    let v = match meta {
        Some(meta) => meta
            .nested
            .iter()
            .map(|m| parse_item_attribute(&m))
            .fold_results(Vec::new(), |mut v, a| {
                v.push(a);
                v
            })?,
        None => Vec::new(),
    };

    Ok(v)
}

// Generate gobject_sys::GEnumValue structs mapping the enum such as:
//     gobject_sys::GEnumValue {
//         value: Animal::Goat as i32,
//         value_name: "Goat\0" as *const _ as *const _,
//         value_nick: "goat\0" as *const _ as *const _,
//     },
fn gen_genum_values(
    enum_name: &Ident,
    enum_variants: &Punctuated<Variant, Comma>,
) -> (TokenStream, usize) {
    // start at one as GEnumValue array is null-terminated
    let mut n = 1;
    let recurse = enum_variants.iter().map(|v| {
        let name = &v.ident;
        let mut value_name = name.to_string().to_camel_case();
        let mut value_nick = name.to_string().to_kebab_case();

        let attrs = parse_item_attributes(&v.attrs);
        let attrs = match attrs {
            Ok(attrs) => attrs,
            Err(e) => abort_call_site!(
                "{}: GEnum enum supports only the following optional attributes: #[genum(name = \"The Cat\", nick = \"chat\")]",
                e
            ),
        };

        attrs.into_iter().for_each(|attr|
            match attr {
                ItemAttribute::Name(n) => value_name = n,
                ItemAttribute::Nick(n) => value_nick = n,
            }
        );

        let value_name = format!("{}\0", value_name);
        let value_nick = format!("{}\0", value_nick);

        n += 1;
        quote_spanned! {v.span()=>
            gobject_sys::GEnumValue {
                value: #enum_name::#name as i32,
                value_name: #value_name as *const _ as *const _,
                value_nick: #value_nick as *const _ as *const _,
            },
        }
    });
    (
        quote! {
            #(#recurse)*
        },
        n,
    )
}

// Parse enum attribute such as:
// #[genum(type_name = "TestAnimalType")]
fn parse_type_name(input: &DeriveInput) -> Result<String> {
    let meta = match find_genum_meta(&input.attrs)? {
        Some(meta) => meta,
        _ => bail!("Missing 'genum' attribute"),
    };

    let meta = match meta.nested.first() {
        Some(meta) => meta,
        _ => bail!("Missing meta 'type_name'"),
    };

    match parse_enum_attribute(&meta)? {
        EnumAttribute::TypeName(n) => Ok(n),
    }
}

pub fn impl_genum(input: &syn::DeriveInput) -> TokenStream {
    let name = &input.ident;

    let enum_variants = match input.data {
        Data::Enum(ref e) => &e.variants,
        _ => abort_call_site!("GEnum only supports enums"),
    };

    let gtype_name = match parse_type_name(&input) {
        Ok(v) => v,
        Err(e) => abort_call_site!(
            "{}: derive(GEnum) requires #[genum(type_name = \"EnumTypeName\")]",
            e
        ),
    };
    let get_type = format_ident!("{}_get_type", name.to_string().to_snake_case());
    let from_glib = gen_from_glib(name, enum_variants);
    let (genum_values, nb_genum_values) = gen_genum_values(name, enum_variants);

    quote! {
        impl glib::translate::ToGlib for #name {
            type GlibType = i32;

            fn to_glib(&self) -> i32 {
                *self as i32
            }
        }

        impl glib::translate::FromGlib<i32> for #name {
            fn from_glib(value: i32) -> Self {
                #from_glib
                unreachable!();
            }
        }

        impl<'a> glib::value::FromValueOptional<'a> for #name {
            unsafe fn from_value_optional(value: &glib::Value) -> Option<Self> {
                Some(glib::value::FromValue::from_value(value))
            }
        }

        impl<'a> glib::value::FromValue<'a> for #name {
            unsafe fn from_value(value: &glib::Value) -> Self {
                glib::translate::from_glib(
                    gobject_sys::g_value_get_enum(
                        glib::translate::ToGlibPtr::to_glib_none(value).0))
            }
        }

        impl glib::value::SetValue for #name {
            unsafe fn set_value(value: &mut glib::Value, this: &Self) {
                gobject_sys::g_value_set_enum(
                    glib::translate::ToGlibPtrMut::to_glib_none_mut(value).0,
                    glib::translate::ToGlib::to_glib(this))
            }
        }

        impl StaticType for #name {
            fn static_type() -> glib::Type {
                #get_type()
            }
        }

        fn #get_type() -> glib::Type {
            static ONCE: std::sync::Once = std::sync::Once::new();
            static mut TYPE: glib::Type = glib::Type::Invalid;

            ONCE.call_once(|| {
                static mut VALUES: [gobject_sys::GEnumValue; #nb_genum_values] = [
                    #genum_values
                    gobject_sys::GEnumValue {
                        value: 0,
                        value_name: std::ptr::null(),
                        value_nick: std::ptr::null(),
                    },
                ];

                let name = std::ffi::CString::new(#gtype_name).expect("CString::new failed");
                unsafe {
                    let type_ = gobject_sys::g_enum_register_static(name.as_ptr(), VALUES.as_ptr());
                    TYPE = glib::translate::from_glib(type_);
                }
            });

            unsafe {
                assert_ne!(TYPE, glib::Type::Invalid);
                TYPE
            }
        }
    }
}
<|endoftext|>


<commit_before>//! An IRC transport that wraps an IRC-framed stream to provide a number of features including
//! automatic PING replies, automatic sending of PINGs, and message rate-limiting.
use std::collections::VecDeque;
use std::sync::{Arc, RwLock, RwLockReadGuard};
use std::time::{Duration, Instant};

use futures::{Async, AsyncSink, Future, Poll, Sink, StartSend, Stream};
use chrono::prelude::*;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_io::codec::Framed;
use tokio_timer;
use tokio_timer::{Interval, Sleep, Timer};

use error;
use client::data::Config;
use proto::{Command, IrcCodec, Message};

/// An IRC transport that handles core functionality.
pub struct IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    inner: Framed<T, IrcCodec>,
    burst_timer: Timer,
    rolling_burst_window: VecDeque<Sleep>,
    burst_window_length: u64,
    max_burst_messages: u64,
    current_burst_messages: u64,
    ping_timer: Interval,
    ping_timeout: u64,
    last_ping_data: String,
    last_ping_sent: Instant,
    last_pong_received: Instant,
}

impl<T> IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    /// Creates a new `IrcTransport` from the given IRC stream.
    pub fn new(config: &Config, inner: Framed<T, IrcCodec>) -> IrcTransport<T> {
        let timer = tokio_timer::wheel().build();
        IrcTransport {
            inner: inner,
            burst_timer: tokio_timer::wheel().build(),
            rolling_burst_window: VecDeque::new(),
            burst_window_length: u64::from(config.burst_window_length()),
            max_burst_messages: u64::from(config.max_messages_in_burst()),
            current_burst_messages: 0,
            ping_timer: timer.interval(Duration::from_secs(u64::from(config.ping_time()))),
            ping_timeout: u64::from(config.ping_timeout()),
            last_ping_data: String::new(),
            last_ping_sent: Instant::now(),
            last_pong_received: Instant::now(),
        }
    }

    /// Gets the inner stream underlying the `IrcTransport`.
    pub fn into_inner(self) -> Framed<T, IrcCodec> {
        self.inner
    }

    fn ping_timed_out(&self) -> bool {
        if self.last_pong_received < self.last_ping_sent {
            self.last_ping_sent.elapsed().as_secs() >= self.ping_timeout
        } else {
            false
        }
    }

    fn send_ping(&mut self) -> error::Result<()> {
        self.last_ping_sent = Instant::now();
        self.last_ping_data = format!("{}", Local::now().timestamp());
        let data = self.last_ping_data.clone();
        let result = self.start_send(Command::PING(data, None).into())?;
        assert!(result.is_ready());
        self.poll_complete()?;
        Ok(())
    }

    fn rolling_burst_window_front(&mut self) -> Result<Async<()>, tokio_timer::TimerError> {
        self.rolling_burst_window.front_mut().map(|w| w.poll()).unwrap_or(Ok(Async::NotReady))
    }
}

impl<T> Stream for IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    type Item = Message;
    type Error = error::IrcError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.ping_timed_out() {
            self.close()?;
            return Err(error::IrcError::PingTimeout)
        }

        let timer_poll = self.ping_timer.poll()?;
        let inner_poll = self.inner.poll()?;

        match (inner_poll, timer_poll) {
            (Async::NotReady, Async::NotReady) => Ok(Async::NotReady),
            (Async::NotReady, Async::Ready(msg)) => {
                assert!(msg.is_some());
                self.send_ping()?;
                Ok(Async::NotReady)
            }
            (Async::Ready(None), _) => Ok(Async::Ready(None)),
            (Async::Ready(Some(msg)), _) => {
                match timer_poll {
                    Async::Ready(msg) => {
                        assert!(msg.is_some());
                        self.send_ping()?;
                    }
                    Async::NotReady => (),
                }

                match msg.command {
                    // Automatically respond to PINGs from the server.
                    Command::PING(ref data, _) => {
                        let result = self.start_send(Command::PONG(data.to_owned(), None).into())?;
                        assert!(result.is_ready());
                        self.poll_complete()?;
                    }
                    // Check PONG responses from the server.
                    Command::PONG(ref data, None) |
                    Command::PONG(_, Some(ref data)) => {
                        if self.last_ping_data == data[..] {
                            self.last_pong_received = Instant::now();
                        }
                    }
                    _ => (),
                }

                Ok(Async::Ready(Some(msg)))
            }
        }
    }
}

impl<T> Sink for IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    type SinkItem = Message;
    type SinkError = error::IrcError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        if self.ping_timed_out() {
            self.close()?;
            Err(error::IrcError::PingTimeout)
        } else {
            // Check if the oldest message in the rolling window is discounted.
            if let Async::Ready(()) = self.rolling_burst_window_front()? {
                self.current_burst_messages -= 1;
                self.rolling_burst_window.pop_front();
            }

            // Throttling if too many messages have been sent recently.
            if self.current_burst_messages >= self.max_burst_messages {
                // When throttled, we know we need to finish sending what's already queued up.
                self.poll_complete()?;
                return Ok(AsyncSink::NotReady(item))
            }

            match self.inner.start_send(item)? {
                AsyncSink::NotReady(item) => Ok(AsyncSink::NotReady(item)),
                AsyncSink::Ready => {
                    self.current_burst_messages += 1;
                    self.rolling_burst_window.push_back(self.burst_timer.sleep(Duration::from_secs(
                        self.burst_window_length
                    )));
                    Ok(AsyncSink::Ready)
                }
            }
        }
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        if self.ping_timed_out() {
            self.close()?;
            Err(error::IrcError::PingTimeout)
        } else {
            Ok(self.inner.poll_complete()?)
        }
    }

    fn close(&mut self) -> Poll<(), Self::SinkError> {
        self.inner.close()
    }
}

/// A view of the logs from a particular `Logged` transport.
#[derive(Clone, Debug)]
pub struct LogView {
    sent: Arc<RwLock<Vec<Message>>>,
    received: Arc<RwLock<Vec<Message>>>,
}

impl LogView {
    /// Gets a read guard for all the messages sent on the transport.
    pub fn sent(&self) -> error::Result<RwLockReadGuard<Vec<Message>>> {
        self.sent.read().map_err(|_| error::IrcError::PoisonedLog)
    }

    /// Gets a read guard for all the messages received on the transport.
    pub fn received(&self) -> error::Result<RwLockReadGuard<Vec<Message>>> {
        self.received.read().map_err(|_| error::IrcError::PoisonedLog)
    }
}

/// A logged version of the `IrcTransport` that records all sent and received messages.
/// Note: this will introduce some performance overhead by cloning all messages.
pub struct Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    inner: IrcTransport<T>,
    view: LogView,
}

impl<T> Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    /// Wraps the given `IrcTransport` in logging.
    pub fn wrap(inner: IrcTransport<T>) -> Logged<T> {
        Logged {
            inner: inner,
            view: LogView {
                sent: Arc::new(RwLock::new(vec![])),
                received: Arc::new(RwLock::new(vec![])),
            },
        }
    }

    /// Gets a view of the logging for this transport.
    pub fn view(&self) -> LogView {
        self.view.clone()
    }
}

impl<T> Stream for Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    type Item = Message;
    type Error = error::IrcError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match try_ready!(self.inner.poll()) {
            Some(msg) => {
                let recv: error::Result<_> = self.view.received.write().map_err(|_| {
                    error::IrcError::PoisonedLog
                });
                recv?.push(msg.clone());
                Ok(Async::Ready(Some(msg)))
            }
            None => Ok(Async::Ready(None)),
        }
    }
}

impl<T> Sink for Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    type SinkItem = Message;
    type SinkError = error::IrcError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let res = self.inner.start_send(item.clone())?;
        let sent: error::Result<_> = self.view.sent.write().map_err(|_| {
            error::IrcError::PoisonedLog
        });
        sent?.push(item);
        Ok(res)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(self.inner.poll_complete()?)
    }
}
<commit_msg>IrcTransport::poll_complete now polls the ping timer (fixes #115).<commit_after>//! An IRC transport that wraps an IRC-framed stream to provide a number of features including
//! automatic PING replies, automatic sending of PINGs, and message rate-limiting.
use std::collections::VecDeque;
use std::sync::{Arc, RwLock, RwLockReadGuard};
use std::time::{Duration, Instant};

use futures::{Async, AsyncSink, Future, Poll, Sink, StartSend, Stream};
use chrono::prelude::*;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_io::codec::Framed;
use tokio_timer;
use tokio_timer::{Interval, Sleep, Timer};

use error;
use client::data::Config;
use proto::{Command, IrcCodec, Message};

/// An IRC transport that handles core functionality.
pub struct IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    inner: Framed<T, IrcCodec>,
    burst_timer: Timer,
    rolling_burst_window: VecDeque<Sleep>,
    burst_window_length: u64,
    max_burst_messages: u64,
    current_burst_messages: u64,
    ping_timer: Interval,
    ping_timeout: u64,
    last_ping_data: String,
    last_ping_sent: Instant,
    last_pong_received: Instant,
}

impl<T> IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    /// Creates a new `IrcTransport` from the given IRC stream.
    pub fn new(config: &Config, inner: Framed<T, IrcCodec>) -> IrcTransport<T> {
        let timer = tokio_timer::wheel().build();
        IrcTransport {
            inner: inner,
            burst_timer: tokio_timer::wheel().build(),
            rolling_burst_window: VecDeque::new(),
            burst_window_length: u64::from(config.burst_window_length()),
            max_burst_messages: u64::from(config.max_messages_in_burst()),
            current_burst_messages: 0,
            ping_timer: timer.interval(Duration::from_secs(u64::from(config.ping_time()))),
            ping_timeout: u64::from(config.ping_timeout()),
            last_ping_data: String::new(),
            last_ping_sent: Instant::now(),
            last_pong_received: Instant::now(),
        }
    }

    /// Gets the inner stream underlying the `IrcTransport`.
    pub fn into_inner(self) -> Framed<T, IrcCodec> {
        self.inner
    }

    fn ping_timed_out(&self) -> bool {
        if self.last_pong_received < self.last_ping_sent {
            self.last_ping_sent.elapsed().as_secs() >= self.ping_timeout
        } else {
            false
        }
    }

    fn send_ping(&mut self) -> error::Result<()> {
        self.last_ping_sent = Instant::now();
        self.last_ping_data = format!("{}", Local::now().timestamp());
        let data = self.last_ping_data.clone();
        let result = self.start_send(Command::PING(data, None).into())?;
        assert!(result.is_ready());
        self.poll_complete()?;
        Ok(())
    }

    fn rolling_burst_window_front(&mut self) -> Result<Async<()>, tokio_timer::TimerError> {
        self.rolling_burst_window.front_mut().map(|w| w.poll()).unwrap_or(Ok(Async::NotReady))
    }
}

impl<T> Stream for IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    type Item = Message;
    type Error = error::IrcError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.ping_timed_out() {
            self.close()?;
            return Err(error::IrcError::PingTimeout)
        }

        let timer_poll = self.ping_timer.poll()?;
        let inner_poll = self.inner.poll()?;

        match (inner_poll, timer_poll) {
            (Async::NotReady, Async::NotReady) => Ok(Async::NotReady),
            (Async::NotReady, Async::Ready(msg)) => {
                assert!(msg.is_some());
                self.send_ping()?;
                Ok(Async::NotReady)
            }
            (Async::Ready(None), _) => Ok(Async::Ready(None)),
            (Async::Ready(Some(msg)), _) => {
                if let Async::Ready(msg) = timer_poll {
                    assert!(msg.is_some());
                    self.send_ping()?;
                }

                match msg.command {
                    // Automatically respond to PINGs from the server.
                    Command::PING(ref data, _) => {
                        let result = self.start_send(Command::PONG(data.to_owned(), None).into())?;
                        assert!(result.is_ready());
                        self.poll_complete()?;
                    }
                    // Check PONG responses from the server.
                    Command::PONG(ref data, None) |
                    Command::PONG(_, Some(ref data)) => {
                        if self.last_ping_data == data[..] {
                            self.last_pong_received = Instant::now();
                        }
                    }
                    _ => (),
                }

                Ok(Async::Ready(Some(msg)))
            }
        }
    }
}

impl<T> Sink for IrcTransport<T>
where
    T: AsyncRead + AsyncWrite,
{
    type SinkItem = Message;
    type SinkError = error::IrcError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        if self.ping_timed_out() {
            self.close()?;
            Err(error::IrcError::PingTimeout)
        } else {
            // Check if the oldest message in the rolling window is discounted.
            if let Async::Ready(()) = self.rolling_burst_window_front()? {
                self.current_burst_messages -= 1;
                self.rolling_burst_window.pop_front();
            }

            // Throttling if too many messages have been sent recently.
            if self.current_burst_messages >= self.max_burst_messages {
                // When throttled, we know we need to finish sending what's already queued up.
                self.poll_complete()?;
                return Ok(AsyncSink::NotReady(item))
            }

            match self.inner.start_send(item)? {
                AsyncSink::NotReady(item) => Ok(AsyncSink::NotReady(item)),
                AsyncSink::Ready => {
                    self.current_burst_messages += 1;
                    self.rolling_burst_window.push_back(self.burst_timer.sleep(Duration::from_secs(
                        self.burst_window_length
                    )));
                    Ok(AsyncSink::Ready)
                }
            }
        }
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        if self.ping_timed_out() {
            self.close()?;
            Err(error::IrcError::PingTimeout)
        } else {
            // If it's time to send a ping, we should do it! This is necessary to ensure that the
            // sink half will close even if the stream half closed without a ping timeout.
            if let Async::Ready(msg) = self.ping_timer.poll()? {
                assert!(msg.is_some());
                self.send_ping()?;
            }

            Ok(self.inner.poll_complete()?)
        }
    }

    fn close(&mut self) -> Poll<(), Self::SinkError> {
        self.inner.close()
    }
}

/// A view of the logs from a particular `Logged` transport.
#[derive(Clone, Debug)]
pub struct LogView {
    sent: Arc<RwLock<Vec<Message>>>,
    received: Arc<RwLock<Vec<Message>>>,
}

impl LogView {
    /// Gets a read guard for all the messages sent on the transport.
    pub fn sent(&self) -> error::Result<RwLockReadGuard<Vec<Message>>> {
        self.sent.read().map_err(|_| error::IrcError::PoisonedLog)
    }

    /// Gets a read guard for all the messages received on the transport.
    pub fn received(&self) -> error::Result<RwLockReadGuard<Vec<Message>>> {
        self.received.read().map_err(|_| error::IrcError::PoisonedLog)
    }
}

/// A logged version of the `IrcTransport` that records all sent and received messages.
/// Note: this will introduce some performance overhead by cloning all messages.
pub struct Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    inner: IrcTransport<T>,
    view: LogView,
}

impl<T> Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    /// Wraps the given `IrcTransport` in logging.
    pub fn wrap(inner: IrcTransport<T>) -> Logged<T> {
        Logged {
            inner: inner,
            view: LogView {
                sent: Arc::new(RwLock::new(vec![])),
                received: Arc::new(RwLock::new(vec![])),
            },
        }
    }

    /// Gets a view of the logging for this transport.
    pub fn view(&self) -> LogView {
        self.view.clone()
    }
}

impl<T> Stream for Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    type Item = Message;
    type Error = error::IrcError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match try_ready!(self.inner.poll()) {
            Some(msg) => {
                let recv: error::Result<_> = self.view.received.write().map_err(|_| {
                    error::IrcError::PoisonedLog
                });
                recv?.push(msg.clone());
                Ok(Async::Ready(Some(msg)))
            }
            None => Ok(Async::Ready(None)),
        }
    }
}

impl<T> Sink for Logged<T>
where
    T: AsyncRead + AsyncWrite,
{
    type SinkItem = Message;
    type SinkError = error::IrcError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let res = self.inner.start_send(item.clone())?;
        let sent: error::Result<_> = self.view.sent.write().map_err(|_| {
            error::IrcError::PoisonedLog
        });
        sent?.push(item);
        Ok(res)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(self.inner.poll_complete()?)
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![no_main]

use std::mem::size_of;

use cros_fuzz::fuzz_target;
use cros_fuzz::rand::FuzzRng;
use data_model::VolatileMemory;
use devices::virtio::{DescriptorChain, Queue};
use rand::{Rng, RngCore};
use sys_util::{GuestAddress, GuestMemory};

const MAX_QUEUE_SIZE: u16 = 256;
const MEM_SIZE: u64 = 1024 * 1024;

thread_local! {
    static GUEST_MEM: GuestMemory = GuestMemory::new(&[(GuestAddress(0), MEM_SIZE)]).unwrap();
}

// These are taken from the virtio spec and can be used as a reference for the size calculations in
// the fuzzer.
#[repr(C, packed)]
struct virtq_desc {
    addr: u64,
    len: u32,
    flags: u16,
    next: u16,
}

#[repr(C, packed)]
struct virtq_avail {
    flags: u16,
    idx: u16,
    ring: [u16; MAX_QUEUE_SIZE as usize],
    used_event: u16,
}

#[repr(C, packed)]
struct virtq_used_elem {
    id: u32,
    len: u32,
}

#[repr(C, packed)]
struct virtq_used {
    flags: u16,
    idx: u16,
    ring: [virtq_used_elem; MAX_QUEUE_SIZE as usize],
    avail_event: u16,
}

fuzz_target!(|data: &[u8]| {
    let mut q = Queue::new(MAX_QUEUE_SIZE);
    let mut rng = FuzzRng::new(data);
    q.size = rng.gen();
    q.ready = true;

    // For each of {desc_table,avail_ring,used_ring} generate a random address that includes enough
    // space to hold the relevant struct with the largest possible queue size.
    let max_table_size = MAX_QUEUE_SIZE as u64 * size_of::<virtq_desc>() as u64;
    q.desc_table = GuestAddress(rng.gen_range(0, MEM_SIZE - max_table_size));
    q.avail_ring = GuestAddress(rng.gen_range(0, MEM_SIZE - size_of::<virtq_avail>() as u64));
    q.used_ring = GuestAddress(rng.gen_range(0, MEM_SIZE - size_of::<virtq_used>() as u64));

    GUEST_MEM.with(|mem| {
        if !q.is_valid(mem) {
            return;
        }

        // First zero out all of the memory.
        let vs = mem.get_slice(0, MEM_SIZE).unwrap();
        vs.write_bytes(0);

        // Fill in the descriptor table.
        let queue_size = q.size as usize;
        let mut buf = vec![0u8; queue_size * size_of::<virtq_desc>()];

        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.desc_table).unwrap();

        // Fill in the available ring. See the definition of virtq_avail above for the source of
        // these numbers.
        let avail_size = 4 + (queue_size * 2) + 2;
        buf.resize(avail_size, 0);
        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.avail_ring).unwrap();

        // Fill in the used ring. See the definition of virtq_used above for the source of
        // these numbers.
        let used_size = 4 + (queue_size * size_of::<virtq_used_elem>()) + 2;
        buf.resize(used_size, 0);
        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.used_ring).unwrap();

        while let Some(avail_desc) = q.pop(mem) {
            let idx = avail_desc.index;
            let total = avail_desc
                .into_iter()
                .filter(DescriptorChain::is_write_only)
                .try_fold(0u32, |sum, cur| sum.checked_add(cur.len));
            if let Some(len) = total {
                q.add_used(mem, idx, len);
            } else {
                q.add_used(mem, idx, 0);
            }
        }
    });
});
<commit_msg>fuzz: update virtqueue fuzzer for new VolatileSlice API<commit_after>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![no_main]

use std::mem::size_of;

use cros_fuzz::fuzz_target;
use cros_fuzz::rand::FuzzRng;
use devices::virtio::{DescriptorChain, Queue};
use rand::{Rng, RngCore};
use sys_util::{GuestAddress, GuestMemory};

const MAX_QUEUE_SIZE: u16 = 256;
const MEM_SIZE: u64 = 1024 * 1024;

thread_local! {
    static GUEST_MEM: GuestMemory = GuestMemory::new(&[(GuestAddress(0), MEM_SIZE)]).unwrap();
}

// These are taken from the virtio spec and can be used as a reference for the size calculations in
// the fuzzer.
#[repr(C, packed)]
struct virtq_desc {
    addr: u64,
    len: u32,
    flags: u16,
    next: u16,
}

#[repr(C, packed)]
struct virtq_avail {
    flags: u16,
    idx: u16,
    ring: [u16; MAX_QUEUE_SIZE as usize],
    used_event: u16,
}

#[repr(C, packed)]
struct virtq_used_elem {
    id: u32,
    len: u32,
}

#[repr(C, packed)]
struct virtq_used {
    flags: u16,
    idx: u16,
    ring: [virtq_used_elem; MAX_QUEUE_SIZE as usize],
    avail_event: u16,
}

fuzz_target!(|data: &[u8]| {
    let mut q = Queue::new(MAX_QUEUE_SIZE);
    let mut rng = FuzzRng::new(data);
    q.size = rng.gen();
    q.ready = true;

    // For each of {desc_table,avail_ring,used_ring} generate a random address that includes enough
    // space to hold the relevant struct with the largest possible queue size.
    let max_table_size = MAX_QUEUE_SIZE as u64 * size_of::<virtq_desc>() as u64;
    q.desc_table = GuestAddress(rng.gen_range(0, MEM_SIZE - max_table_size));
    q.avail_ring = GuestAddress(rng.gen_range(0, MEM_SIZE - size_of::<virtq_avail>() as u64));
    q.used_ring = GuestAddress(rng.gen_range(0, MEM_SIZE - size_of::<virtq_used>() as u64));

    GUEST_MEM.with(|mem| {
        if !q.is_valid(mem) {
            return;
        }

        // First zero out all of the memory.
        let vs = mem
            .get_slice_at_addr(GuestAddress(0), MEM_SIZE as usize)
            .unwrap();
        vs.write_bytes(0);

        // Fill in the descriptor table.
        let queue_size = q.size as usize;
        let mut buf = vec![0u8; queue_size * size_of::<virtq_desc>()];

        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.desc_table).unwrap();

        // Fill in the available ring. See the definition of virtq_avail above for the source of
        // these numbers.
        let avail_size = 4 + (queue_size * 2) + 2;
        buf.resize(avail_size, 0);
        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.avail_ring).unwrap();

        // Fill in the used ring. See the definition of virtq_used above for the source of
        // these numbers.
        let used_size = 4 + (queue_size * size_of::<virtq_used_elem>()) + 2;
        buf.resize(used_size, 0);
        rng.fill_bytes(&mut buf[..]);
        mem.write_all_at_addr(&buf[..], q.used_ring).unwrap();

        while let Some(avail_desc) = q.pop(mem) {
            let idx = avail_desc.index;
            let total = avail_desc
                .into_iter()
                .filter(DescriptorChain::is_write_only)
                .try_fold(0u32, |sum, cur| sum.checked_add(cur.len));
            if let Some(len) = total {
                q.add_used(mem, idx, len);
            } else {
                q.add_used(mem, idx, 0);
            }
        }
    });
});
<|endoftext|>


<commit_before>#![allow(unused_variables)]

fn main() {
enum Option<T> {
    Some(T),
    None,
}

}
<commit_msg>Removed allowed warnings<commit_after>fn main() {
enum Option<T> {
    Some(T),
    None,
}

}
<|endoftext|>