<commit_before>mod addressing;
mod instructions;
mod registers;

use self::addressing::*;
use bus::{Bus, Addressable};
use byteorder::{ByteOrder, LittleEndian};
use self::instructions as inst;
use self::registers::*;
use std::fmt;

pub enum Condition {
    Z,
    C,
    Nz,
    Nc
}

impl fmt::Display for Condition {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Condition::Z => write!(f, "Z"),
            Condition::C => write!(f, "C"),
            Condition::Nz => write!(f, "NZ"),
            Condition::Nc => write!(f, "NC")
        }
    }
}

fn condition_table(idx: u8) -> Condition {
    match idx {
        0 => Condition::Nz,
        1 => Condition::Z,
        2 => Condition::Nc,
        3 => Condition::C,
        _ => unreachable!()
    }
}

fn register_addr_table(idx: u8) -> Box<AddressingMode<u8>> {
    match idx {
        0 => Box::new(RegisterAddressing(Register::B)),
        1 => Box::new(RegisterAddressing(Register::C)),
        2 => Box::new(RegisterAddressing(Register::D)),
        3 => Box::new(RegisterAddressing(Register::E)),
        4 => Box::new(RegisterAddressing(Register::H)),
        5 => Box::new(RegisterAddressing(Register::L)),
        6 => Box::new(RegisterIndirectAddressing(Register::HL)),
        7 => Box::new(RegisterAddressing(Register::A)),
        _ => unreachable!()
    }
}

fn register_pair_table(idx: u8) -> Register {
    match idx {
        0 => Register::BC,
        1 => Register::DE,
        2 => Register::HL,
        3 => Register::SP,
        _ => unreachable!()
    }
}

pub struct Cpu {
    regs: Registers,
    ime: bool,
    prefixed: bool
}

impl Cpu {
    pub fn new() -> Self {
        Cpu {
            regs: Registers::default(),
            ime: false,
            prefixed: false
        }
    }

    pub fn power_on(&mut self) {
        self.regs.set_af(0x01B0);
        self.regs.set_bc(0x1300);
        self.regs.set_de(0xD800);
        self.regs.set_hl(0x4D01);
        self.regs.set_sp(0xFFFE);
        self.regs.set_pc(0x100);
        self.ime = true;
        self.prefixed = false;
    }

    pub fn step(&mut self, bus: &mut Bus) -> usize {
        let pc = self.regs.pc();
        let opcode = self.next_byte(&bus);

        println!("{:#X}: {:#X}", pc, opcode);

        self.decode_and_execute(bus, opcode);

        0
    }

    fn next_byte(&mut self, bus: &Bus) -> u8 {
        let pc = self.regs.pc();
        let byte = bus.read(pc);
        self.regs.set_pc(pc.wrapping_add(1));

        byte
    }

    fn next_word(&mut self, bus: &Bus) -> u16 {
        let lb = self.next_byte(bus);
        let hb = self.next_byte(bus);

        LittleEndian::read_u16(&[lb, hb])
    }

    fn pop_stack(&mut self, bus: &Bus) -> u16 {
        let addr = self.regs.sp();
        let pc = self.regs.pc();
        let word = &[bus.read(addr), bus.read(addr + 1)];
        self.regs.set_pc(pc + 2);

        LittleEndian::read_u16(word)
    }

    fn push_stack(&mut self, bus: &mut Bus, val: u16) {
        let addr = self.regs.sp();

        bus.write(addr - 2, (val & 0x00FF) as u8);
        bus.write(addr - 1, ((val >> 8) & 0x00FF) as u8);

        self.regs.set_sp(addr - 2);
    }

    fn condition_met(&self, cond: Condition) -> bool {
        match cond {
            Condition::C => (self.regs.f() & FLAG_C) == FLAG_C,
            Condition::Z => (self.regs.f() & FLAG_Z) == FLAG_Z,
            Condition::Nc => (self.regs.f() & FLAG_C) != FLAG_C,
            Condition::Nz => (self.regs.f() & FLAG_Z) != FLAG_Z
        }
    }

    fn decode_and_execute(&mut self, bus: &mut Bus, opcode: u8) {
        let x = (0b11000000 & opcode) >> 6;
        let y = (0b00111000 & opcode) >> 3;
        let z = 0b00000111 & opcode;
        let p = y >> 1;
        let q = (0b00001000 & opcode) >> 3;

        match x {
            0 => match z {
                0 => self.decode_x_0_z_0(bus, y),
                1 => self.decode_x_0_z_1(bus, q, p),
                2 => self.decode_x_0_z_2(bus, q, p),
                3 => self.decode_x_0_z_3(bus, q, p),
                4 => {
                    // INC r
                    let dest = register_addr_table(y);
                    inst::inc_8(self, bus, dest.as_ref());
                },
                5 => {
                    // DEC r
                    let dest = register_addr_table(y);
                    inst::dec_8(self, bus, dest.as_ref());
                },
                6 => {
                    // LD r, n
                    let dest = register_addr_table(y);
                    let src = ImmediateAddressing(self.next_byte(bus));
                    inst::ld(self, bus, dest.as_ref(), &src);
                },
                7 => self.decode_x_0_z_7(bus, y),
                _ => unreachable!()
            },
            1 => {
                if y == 6 {
                    // HALT
                    inst::halt();
                } else {
                    // LD r, r
                    let dest = register_addr_table(y);
                    let src = register_addr_table(z);
                    inst::ld(self, bus, dest.as_ref(), src.as_ref());
                }
            },
            2 => self.decode_x_2(bus, y, z),
            3 => self.decode_x_3(bus, y, z, q, p),
            _ => unreachable!()
        };
    }

    fn decode_x_0_z_0(&mut self, bus: &mut Bus, y: u8) {
        match y {
            0 => inst::nop(), // NOP
            1 => {
                // LD (nn),SP
                let dest = ExtendedAddressing(self.next_word(bus));
                let src = RegisterAddressing(Register::SP);
                inst::ld(self, bus, &dest, &src);
            },
            2 => {
                // STOP 0
                self.next_byte(bus);
                inst::stop();
            },
            3 => {
                // JR i8
                let offset = ImmediateAddressing(self.next_byte(bus));
                inst::jr(self, bus, &offset);
            },
            4...7 => {
                // JR cc, i8
                let condition = condition_table(y - 4);
                if self.condition_met(condition) {
                    let offset = ImmediateAddressing(self.next_byte(bus));
                    inst::jr(self, bus, &offset);
                }
            },
            _ => unreachable!()
        }
    }

    fn decode_x_0_z_1(&mut self, bus: &mut Bus, q: u8, p: u8) {
        match q {
            0 => {
                // LD rr, nn
                let register = register_pair_table(p);
                let dest = RegisterAddressing(register);
                let src = ImmediateAddressing(self.next_word(bus));
                inst::ld(self, bus, &dest, &src);
            },
            1 => {
                // ADD HL, rr
                let register = register_pair_table(p);
                let dest = RegisterAddressing(Register::HL);
                let src = RegisterAddressing(register);
                inst::add_16(self, bus, &dest, &src);
            },
            _ => unreachable!()
        }
    }

    fn decode_x_0_z_2(&mut self, bus: &mut Bus, q: u8, p: u8) {
        match q {
            0 => match p {
                0 => {
                    // LD (BC), A
                    let dest = RegisterIndirectAddressing(Register::BC);
                    let src = RegisterAddressing(Register::A);
                    inst::ld(self, bus, &dest, &src);
                },
                1 => {
                    // LD (DE), A
                    let dest = RegisterIndirectAddressing(Register::DE);
                    let src = RegisterAddressing(Register::A);
                    inst::ld(self, bus, &dest, &src);
                },
                2 => {
                    // LD (nn), HL
                    let dest = ExtendedAddressing(self.next_word(bus));
                    let src = RegisterAddressing(Register::HL);
                    inst::ld(self, bus, &dest, &src);
                },
                3 => {
                    // LD (nn), A
                    let dest = ExtendedAddressing(self.next_word(bus));
                    let src = RegisterAddressing(Register::A);
                    inst::ld(self, bus, &dest, &src);
                },
                _ => unreachable!()
            },
            1 => match p {
                0 => {
                    // LD A, (BC)
                    let dest = RegisterAddressing(Register::A);
                    let src = RegisterIndirectAddressing(Register::BC);
                    inst::ld(self, bus, &dest, &src);
                },
                1 => {
                    // LD A, (DE)
                    let dest = RegisterAddressing(Register::A);
                    let src = RegisterIndirectAddressing(Register::DE);
                    inst::ld(self, bus, &dest, &src);
                },
                2 => {
                    // LD HL, (nn)
                    let dest = RegisterAddressing(Register::HL);
                    let src = ExtendedAddressing(self.next_word(bus));
                    inst::ld(self, bus, &dest, &src);
                },
                3 => {
                    // LD A, (nn)
                    let dest = RegisterAddressing(Register::A);
                    let src = ExtendedAddressing(self.next_word(bus));
                    inst::ld(self, bus, &dest, &src);
                },
                _ => unreachable!()
            },
            _ => unreachable!()
        }
    }

    fn decode_x_0_z_3(&mut self, bus: &mut Bus, q: u8, p: u8) {
        let register = register_pair_table(p);
        let dest = RegisterAddressing(register);

        match q {
            0 => inst::inc_16(self, bus, &dest),    // INC rr
            1 => inst::dec_16(self, bus, &dest),    // DEC rr
            _ => unreachable!()
        }
    }

    fn decode_x_0_z_7(&mut self, bus: &mut Bus, y: u8) {
        match y {
            0 => inst::rlca(self),  // RLCA
            1 => inst::rrca(self),  // RRCA
            2 => inst::rla(self),   // RLA
            3 => inst::rra(self),   // RRA
            4 => inst::daa(self),   // DAA
            5 => inst::cpl(self),   // CPL
            6 => inst::scf(self),   // SCF
            7 => inst::ccf(self),   // CCF
            _ => unreachable!()
        }
    }

    fn decode_x_2(&mut self, bus: &mut Bus, y: u8, z: u8) {
        let src = register_addr_table(z);

        match y {
            0 => inst::add_8(self, bus, src.as_ref()),  // ADD A, r
            1 => inst::adc(self, bus, src.as_ref()),    // ADC A, r
            2 => inst::sub(self, bus, src.as_ref()),    // SUB r
            3 => inst::sbc(self, bus, src.as_ref()),    // SBC A, r
            4 => inst::and(self, bus, src.as_ref()),    // AND
            5 => inst::xor(self, bus, src.as_ref()),    // XOR
            6 => inst::or(self, bus, src.as_ref()),     // OR
            7 => inst::cp(self, bus, src.as_ref()),     // CP
            _ => unreachable!()
        }
    }

    fn decode_x_3(&mut self, bus: &mut Bus, y: u8, z: u8, q: u8, p: u8) {
        match z {
            0 => {
                let condition = condition_table(y);
                if self.condition_met(condition) {
                    inst::ret(self, bus);
                }
            },
            _ => unreachable!()
        }
    }
}<commit_msg>Condensed decode functions Wasn't really liking splitting them all out into separate functions. Changed it to just some nested matches.<commit_after>mod addressing;
mod instructions;
mod registers;

use self::addressing::*;
use bus::{Bus, Addressable};
use byteorder::{ByteOrder, LittleEndian};
use self::instructions as inst;
use self::registers::*;
use std::fmt;

pub enum Condition {
    Z,
    C,
    Nz,
    Nc
}

impl fmt::Display for Condition {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Condition::Z => write!(f, "Z"),
            Condition::C => write!(f, "C"),
            Condition::Nz => write!(f, "NZ"),
            Condition::Nc => write!(f, "NC")
        }
    }
}

fn condition_table(idx: u8) -> Condition {
    match idx {
        0 => Condition::Nz,
        1 => Condition::Z,
        2 => Condition::Nc,
        3 => Condition::C,
        _ => unreachable!()
    }
}

fn register_addr_table(idx: u8) -> Box<AddressingMode<u8>> {
    match idx {
        0 => Box::new(RegisterAddressing(Register::B)),
        1 => Box::new(RegisterAddressing(Register::C)),
        2 => Box::new(RegisterAddressing(Register::D)),
        3 => Box::new(RegisterAddressing(Register::E)),
        4 => Box::new(RegisterAddressing(Register::H)),
        5 => Box::new(RegisterAddressing(Register::L)),
        6 => Box::new(RegisterIndirectAddressing(Register::HL)),
        7 => Box::new(RegisterAddressing(Register::A)),
        _ => unreachable!()
    }
}

fn register_pair_table(idx: u8) -> Register {
    match idx {
        0 => Register::BC,
        1 => Register::DE,
        2 => Register::HL,
        3 => Register::SP,
        _ => unreachable!()
    }
}

pub struct Cpu {
    regs: Registers,
    ime: bool,
    prefixed: bool
}

impl Cpu {
    pub fn new() -> Self {
        Cpu {
            regs: Registers::default(),
            ime: false,
            prefixed: false
        }
    }

    pub fn power_on(&mut self) {
        self.regs.set_af(0x01B0);
        self.regs.set_bc(0x1300);
        self.regs.set_de(0xD800);
        self.regs.set_hl(0x4D01);
        self.regs.set_sp(0xFFFE);
        self.regs.set_pc(0x100);
        self.ime = true;
        self.prefixed = false;
    }

    pub fn step(&mut self, bus: &mut Bus) -> usize {
        let pc = self.regs.pc();
        let opcode = self.next_byte(&bus);

        println!("{:#X}: {:#X}", pc, opcode);

        self.decode_and_execute(bus, opcode);

        0
    }

    fn next_byte(&mut self, bus: &Bus) -> u8 {
        let pc = self.regs.pc();
        let byte = bus.read(pc);
        self.regs.set_pc(pc.wrapping_add(1));

        byte
    }

    fn next_word(&mut self, bus: &Bus) -> u16 {
        let lb = self.next_byte(bus);
        let hb = self.next_byte(bus);

        LittleEndian::read_u16(&[lb, hb])
    }

    fn pop_stack(&mut self, bus: &Bus) -> u16 {
        let addr = self.regs.sp();
        let pc = self.regs.pc();
        let word = &[bus.read(addr), bus.read(addr + 1)];
        self.regs.set_pc(pc + 2);

        LittleEndian::read_u16(word)
    }

    fn push_stack(&mut self, bus: &mut Bus, val: u16) {
        let addr = self.regs.sp();

        bus.write(addr - 2, (val & 0x00FF) as u8);
        bus.write(addr - 1, ((val >> 8) & 0x00FF) as u8);

        self.regs.set_sp(addr - 2);
    }

    fn condition_met(&self, cond: Condition) -> bool {
        match cond {
            Condition::C => (self.regs.f() & FLAG_C) == FLAG_C,
            Condition::Z => (self.regs.f() & FLAG_Z) == FLAG_Z,
            Condition::Nc => (self.regs.f() & FLAG_C) != FLAG_C,
            Condition::Nz => (self.regs.f() & FLAG_Z) != FLAG_Z
        }
    }

    fn decode_and_execute(&mut self, bus: &mut Bus, opcode: u8) {
        let x = (0b11000000 & opcode) >> 6;
        let y = (0b00111000 & opcode) >> 3;
        let z = 0b00000111 & opcode;
        let p = y >> 1;
        let q = (0b00001000 & opcode) >> 3;

        match x {
            // X=0
            0 => match z {
                // X=0, Z=0
                0 => {
                    match y {
                        // X=0, Z=0, Y=0
                        // NOP
                        0 => inst::nop(),
                        1 => {
                            // X=0, Z=0, Y=1
                            // LD (nn),SP
                            let dest = ExtendedAddressing(self.next_word(bus));
                            let src = RegisterAddressing(Register::SP);
                            inst::ld(self, bus, &dest, &src);
                        },
                        2 => {
                            // X=0, Z=0, Y=2
                            // STOP 0
                            self.next_byte(bus);
                            inst::stop();
                        },
                        3 => {
                            // X=0, Z=0, Y=3
                            // JR i8
                            let offset = ImmediateAddressing(self.next_byte(bus));
                            inst::jr(self, bus, &offset);
                        },
                        4...7 => {
                            // X=0, Z=0, Y=4-7
                            // JR cc, i8
                            let condition = condition_table(y - 4);
                            if self.condition_met(condition) {
                                let offset = ImmediateAddressing(self.next_byte(bus));
                                inst::jr(self, bus, &offset);
                            }
                        },
                        _ => unreachable!()
                    }
                },
                // X=0, Z=1
                1 => {
                    match q {
                        0 => {
                            // X=0, Z=1, Q=0
                            // LD rr, nn
                            let register = register_pair_table(p);
                            let dest = RegisterAddressing(register);
                            let src = ImmediateAddressing(self.next_word(bus));
                            inst::ld(self, bus, &dest, &src);
                        },
                        1 => {
                            // X=0, Z=1, Q=1
                            // ADD HL, rr
                            let register = register_pair_table(p);
                            let dest = RegisterAddressing(Register::HL);
                            let src = RegisterAddressing(register);
                            inst::add_16(self, bus, &dest, &src);
                        },
                        _ => unreachable!()
                    }
                },
                // X=0, Z=2
                2 => {
                    match q {
                        // X=0, Z=2, Q=0
                        0 => match p {
                            0 => {
                                // X=0, Z=2, Q=0, P=0
                                // LD (BC), A
                                let dest = RegisterIndirectAddressing(Register::BC);
                                let src = RegisterAddressing(Register::A);
                                inst::ld(self, bus, &dest, &src);
                            },
                            1 => {
                                // X=0, Z=2, Q=0, P=1
                                // LD (DE), A
                                let dest = RegisterIndirectAddressing(Register::DE);
                                let src = RegisterAddressing(Register::A);
                                inst::ld(self, bus, &dest, &src);
                            },
                            2 => {
                                // X=0, Z=2, Q=0, P=2
                                // LD (nn), HL
                                let dest = ExtendedAddressing(self.next_word(bus));
                                let src = RegisterAddressing(Register::HL);
                                inst::ld(self, bus, &dest, &src);
                            },
                            3 => {
                                // X=0, Z=2, Q=0, P=3
                                // LD (nn), A
                                let dest = ExtendedAddressing(self.next_word(bus));
                                let src = RegisterAddressing(Register::A);
                                inst::ld(self, bus, &dest, &src);
                            },
                            _ => unreachable!()
                        },
                        // X=0, Z=2, Q=1
                        1 => match p {
                            0 => {
                                // X=0, Z=2, Q=1, P=0
                                // LD A, (BC)
                                let dest = RegisterAddressing(Register::A);
                                let src = RegisterIndirectAddressing(Register::BC);
                                inst::ld(self, bus, &dest, &src);
                            },
                            1 => {
                                // X=0, Z=2, Q=1, P=1
                                // LD A, (DE)
                                let dest = RegisterAddressing(Register::A);
                                let src = RegisterIndirectAddressing(Register::DE);
                                inst::ld(self, bus, &dest, &src);
                            },
                            2 => {
                                // X=0, Z=2, Q=1, P=2
                                // LD HL, (nn)
                                let dest = RegisterAddressing(Register::HL);
                                let src = ExtendedAddressing(self.next_word(bus));
                                inst::ld(self, bus, &dest, &src);
                            },
                            3 => {
                                // X=0, Z=2, Q=1, P=3
                                // LD A, (nn)
                                let dest = RegisterAddressing(Register::A);
                                let src = ExtendedAddressing(self.next_word(bus));
                                inst::ld(self, bus, &dest, &src);
                            },
                            _ => unreachable!()
                        },
                        _ => unreachable!()
                    }
                },
                // X=0, Z=3
                3 => {
                    let register = register_pair_table(p);
                    let dest = RegisterAddressing(register);

                    match q {
                        // X=0, Z=3, Q=0
                        // INC rr
                        0 => inst::inc_16(self, bus, &dest),
                        // X=0, Z=3, Q=1
                        // DEC rr
                        1 => inst::dec_16(self, bus, &dest),
                        _ => unreachable!()
                    }
                },
                // X=0, Z=4
                // INC r
                4 => {
                    let dest = register_addr_table(y);
                    inst::inc_8(self, bus, dest.as_ref());
                },
                // X=0, Z=5
                // DEC r
                5 => {
                    let dest = register_addr_table(y);
                    inst::dec_8(self, bus, dest.as_ref());
                },
                // X=0, Z=6
                // LD r, n
                6 => {
                    let dest = register_addr_table(y);
                    let src = ImmediateAddressing(self.next_byte(bus));
                    inst::ld(self, bus, dest.as_ref(), &src);
                },
                // X=0, Z=7
                7 => {
                    match y {
                        // X=0, Z=7, Y=0
                        // RLCA
                        0 => inst::rlca(self),
                        // X=0, Z=7, Y=1
                        // RRCA
                        1 => inst::rrca(self),
                        // X=0, Z=7, Y=2
                        // RLA
                        2 => inst::rla(self),
                        // X=0, Z=7, Y=3
                        // RRA
                        3 => inst::rra(self),
                        // X=0, Z=7, Y=4
                        // DAA
                        4 => inst::daa(self),
                        // X=0, Z=7, Y=5
                        // CPL
                        5 => inst::cpl(self),
                        // X=0, Z=7, Y=6
                        // SCF
                        6 => inst::scf(self),
                        // X=0, Z=7, Y=7
                        // CCF
                        7 => inst::ccf(self),
                        _ => unreachable!()
                    }
                },
                _ => unreachable!()
            },
            // X=1
            1 => {
                if y == 6 {
                    // X=1, Y=6
                    // HALT
                    inst::halt();
                } else {
                    // X=1, Y=1..5,7
                    // LD r, r
                    let dest = register_addr_table(y);
                    let src = register_addr_table(z);
                    inst::ld(self, bus, dest.as_ref(), src.as_ref());
                }
            },
            // X=2
            2 => {
                let src = register_addr_table(z);
                match y {
                    // X=2, Y=0
                    // ADD A, r
                    0 => inst::add_8(self, bus, src.as_ref()),
                    // X=2, Y=1
                    // ADC A, r
                    1 => inst::adc(self, bus, src.as_ref()),
                    // X=2, Y=2
                    // SUB r
                    2 => inst::sub(self, bus, src.as_ref()),
                    // X=2, Y=3
                    // SBC A, r
                    3 => inst::sbc(self, bus, src.as_ref()),
                    // X=2, Y=4
                    // AND
                    4 => inst::and(self, bus, src.as_ref()),
                    // X=2, Y=5
                    // XOR
                    5 => inst::xor(self, bus, src.as_ref()),
                    // X=2, Y=6
                    // OR
                    6 => inst::or(self, bus, src.as_ref()),
                    // X=2, Y=7
                    // CP
                    7 => inst::cp(self, bus, src.as_ref()),
                    _ => unreachable!()
                }
            },
            // X=3
            3 => {
                match z {
                    0 => {
                        // X=3, Z=0
                        // RET
                        let condition = condition_table(y);
                        if self.condition_met(condition) {
                            inst::ret(self, bus);
                        }
                    },
                    _ => unreachable!()
                }
            },
            _ => unreachable!()
        };
    }
}<|endoftext|>


<commit_before>#![allow(experimental)]

use std::io::{File,fs};
use std::path::posix::Path;
use std::os::{homedir};
use std::rand::random;
use std::finally::Finally;

pub fn new(name: &str) {
    let home_unwrap = homedir().unwrap();
    let home        = home_unwrap.display();

    let muxed_dir = if muxed_dir_exists(&"muxed".to_string()) {
      Path::new(format!("{}/.muxed/", home))
    } else {
      create_muxed_dir(&"muxed".to_string())
    };

    let path = format!("{}/{}", muxed_dir.display(), name);

    File::create(&Path::new(path));
}

fn create_muxed_dir(name: &String) -> Path {
    let home_unwrap = homedir().unwrap();
    let path = &Path::new(format!("{}/.{}", home_unwrap.display(), name));
    fs::mkdir(path, ::std::io::UserRWX);
    path.clone()
}

fn muxed_dir_exists(name: &String) -> bool {
    let home_unwrap = homedir().unwrap();
    let path = &Path::new(format!("{}/.{}", home_unwrap.display(), name));
    path.exists()
}

#[test]
fn muxed_dir_exists_returns_false() {
  let dir = format!("test_dir_{}", random::<f64>());
  assert!(!muxed_dir_exists(&dir));
}

#[test]
fn muxed_dir_exists_returns_true() {
  let dir = format!("test_dir_{}", random::<f64>());
  create_muxed_dir(&dir);
  assert!(muxed_dir_exists(&dir));

  let home_unwrap = homedir().unwrap();
  let muxed_path  = &Path::new(format!("{}/.{}/", home_unwrap.display(), dir.as_slice()));
  fs::rmdir_recursive(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let name        = format!("test_project_{}", random::<f64>());
    let dir         = format!("test_dir_{}", random::<f64>());
    let home_unwrap = homedir().unwrap();
    let muxed_path  = &Path::new(format!("{}/.{}/", home_unwrap.display(), dir.as_slice()));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    fs::rmdir_recursive(muxed_path);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = format!("test_project_{}", random::<f64>());
    let home_unwrap = homedir().unwrap();
    let path = &Path::new(format!("{}/.muxed/{}", home_unwrap.display(), name));
    new(name.as_slice());
    assert!(path.exists());
    fs::unlink(path);
}

#[test]
#[should_fail]
fn new_doesnt_overwrite_existing_file() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    (|| {
        new(name.as_slice());
    }).finally(|| {
        fs::unlink(path);
    })
}
<commit_msg>Use a homdir_path for convenience instead of unwrapping and displaying the path in every method.<commit_after>#![allow(experimental)]

use std::io::{File,fs};
use std::path::posix::Path;
use std::os::{homedir};
use std::rand::random;
use std::finally::Finally;

pub fn new(name: &str) {
    let muxed_dir = if muxed_dir_exists(&"muxed".to_string()) {
      Path::new(format!("{}/.muxed/", homedir_string()))
    } else {
      create_muxed_dir(&"muxed".to_string())
    };

    let path = format!("{}/{}", muxed_dir.display(), name);

    File::create(&Path::new(path));
}

fn create_muxed_dir(name: &String) -> Path {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    fs::mkdir(path, ::std::io::UserRWX);
    path.clone()
}

fn muxed_dir_exists(name: &String) -> bool {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    path.exists()
}

fn homedir_string() -> String {
    let home_unwrap = homedir().unwrap();
    format!("{}", home_unwrap.display())
}

#[test]
fn muxed_dir_exists_returns_false() {
  let dir = format!("test_dir_{}", random::<f64>());
  assert!(!muxed_dir_exists(&dir));
}

#[test]
fn muxed_dir_exists_returns_true() {
  let dir = format!("test_dir_{}", random::<f64>());
  create_muxed_dir(&dir);
  assert!(muxed_dir_exists(&dir));

  let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
  fs::rmdir_recursive(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let name        = format!("test_project_{}", random::<f64>());
    let dir         = format!("test_dir_{}", random::<f64>());
    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    fs::rmdir_recursive(muxed_path);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    assert!(path.exists());
    fs::unlink(path);
}

#[test]
#[should_fail]
fn new_doesnt_overwrite_existing_file() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    (|| {
        new(name.as_slice());
    }).finally(|| {
        fs::unlink(path);
    })
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::HashSet;
use std::os;
use std::io::IoError;
use syntax::ast;

pub struct RPathConfig<'a> {
    pub used_crates: Vec<(ast::CrateNum, Option<Path>)>,
    pub out_filename: Path,
    pub is_like_osx: bool,
    pub has_rpath: bool,
    pub get_install_prefix_lib_path: ||:'a -> Path,
    pub realpath: |&Path|:'a -> Result<Path, IoError>
}

pub fn get_rpath_flags(config: RPathConfig) -> Vec<String> {

    // No rpath on windows
    if !config.has_rpath {
        return Vec::new();
    }

    let mut flags = Vec::new();

    debug!("preparing the RPATH!");

    let libs = config.used_crates.clone();
    let libs = libs.into_iter().filter_map(|(_, l)| {
        l.map(|p| p.clone())
    }).collect::<Vec<_>>();

    let rpaths = get_rpaths(config, libs.as_slice());
    flags.push_all(rpaths_to_flags(rpaths.as_slice()).as_slice());
    flags
}

fn rpaths_to_flags(rpaths: &[String]) -> Vec<String> {
    let mut ret = Vec::new();
    for rpath in rpaths.iter() {
        ret.push(format!("-Wl,-rpath,{}", (*rpath).as_slice()));
    }
    return ret;
}

fn get_rpaths(mut config: RPathConfig,
              libs: &[Path]) -> Vec<String> {
    debug!("output: {}", config.out_filename.display());
    debug!("libs:");
    for libpath in libs.iter() {
        debug!("    {}", libpath.display());
    }

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(&mut config, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = vec!(get_install_prefix_rpath(config));

    fn log_rpaths(desc: &str, rpaths: &[String]) {
        debug!("{} rpaths:", desc);
        for rpath in rpaths.iter() {
            debug!("    {}", *rpath);
        }
    }

    log_rpaths("relative", rel_rpaths.as_slice());
    log_rpaths("fallback", fallback_rpaths.as_slice());

    let mut rpaths = rel_rpaths;
    rpaths.push_all(fallback_rpaths.as_slice());

    // Remove duplicates
    let rpaths = minimize_rpaths(rpaths.as_slice());
    return rpaths;
}

fn get_rpaths_relative_to_output(config: &mut RPathConfig,
                                 libs: &[Path]) -> Vec<String> {
    libs.iter().map(|a| get_rpath_relative_to_output(config, a)).collect()
}

fn get_rpath_relative_to_output(config: &mut RPathConfig,
                                lib: &Path) -> String {
    use std::os;

    // Mac doesn't appear to support $ORIGIN
    let prefix = if config.is_like_osx {
        "@loader_path"
    } else {
        "$ORIGIN"
    };

    let mut lib = (config.realpath)(&os::make_absolute(lib).unwrap()).unwrap();
    lib.pop();
    let mut output = (config.realpath)(&os::make_absolute(&config.out_filename).unwrap()).unwrap();
    output.pop();
    let relative = lib.path_relative_from(&output);
    let relative = relative.expect("could not create rpath relative to output");
    // FIXME (#9639): This needs to handle non-utf8 paths
    format!("{}/{}",
            prefix,
            relative.as_str().expect("non-utf8 component in path"))
}

fn get_install_prefix_rpath(config: RPathConfig) -> String {
    let path = (config.get_install_prefix_lib_path)();
    let path = os::make_absolute(&path).unwrap();
    // FIXME (#9639): This needs to handle non-utf8 paths
    path.as_str().expect("non-utf8 component in rpath").to_string()
}

fn minimize_rpaths(rpaths: &[String]) -> Vec<String> {
    let mut set = HashSet::new();
    let mut minimized = Vec::new();
    for rpath in rpaths.iter() {
        if set.insert(rpath.as_slice()) {
            minimized.push(rpath.clone());
        }
    }
    minimized
}

#[cfg(all(unix, test))]
mod test {
    use super::{RPathConfig};
    use super::{minimize_rpaths, rpaths_to_flags, get_rpath_relative_to_output};
    use syntax::abi;

    #[test]
    fn test_rpaths_to_flags() {
        let flags = rpaths_to_flags(&[
            "path1".to_string(),
            "path2".to_string()
        ]);
        assert_eq!(flags,
                   vec!("-Wl,-rpath,path1".to_string(),
                        "-Wl,-rpath,path2".to_string()));
    }

    #[test]
    fn test_minimize1() {
        let res = minimize_rpaths(&[
            "rpath1".to_string(),
            "rpath2".to_string(),
            "rpath1".to_string()
        ]);
        assert!(res == [
            "rpath1",
            "rpath2",
        ]);
    }

    #[test]
    fn test_minimize2() {
        let res = minimize_rpaths(&[
            "1a".to_string(),
            "2".to_string(),
            "2".to_string(),
            "1a".to_string(),
            "4a".to_string(),
            "1a".to_string(),
            "2".to_string(),
            "3".to_string(),
            "4a".to_string(),
            "3".to_string()
        ]);
        assert!(res == [
            "1a",
            "2",
            "4a",
            "3",
        ]);
    }

    #[test]
    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            has_rpath: true,
            is_like_osx: false,
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "freebsd")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: false,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "dragonfly")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: false,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: true,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "@loader_path/../lib");
    }
}
<commit_msg>librustc_back: remove unnecessary `to_string()` calls<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::HashSet;
use std::os;
use std::io::IoError;
use syntax::ast;

pub struct RPathConfig<'a> {
    pub used_crates: Vec<(ast::CrateNum, Option<Path>)>,
    pub out_filename: Path,
    pub is_like_osx: bool,
    pub has_rpath: bool,
    pub get_install_prefix_lib_path: ||:'a -> Path,
    pub realpath: |&Path|:'a -> Result<Path, IoError>
}

pub fn get_rpath_flags(config: RPathConfig) -> Vec<String> {

    // No rpath on windows
    if !config.has_rpath {
        return Vec::new();
    }

    let mut flags = Vec::new();

    debug!("preparing the RPATH!");

    let libs = config.used_crates.clone();
    let libs = libs.into_iter().filter_map(|(_, l)| {
        l.map(|p| p.clone())
    }).collect::<Vec<_>>();

    let rpaths = get_rpaths(config, libs.as_slice());
    flags.push_all(rpaths_to_flags(rpaths.as_slice()).as_slice());
    flags
}

fn rpaths_to_flags(rpaths: &[String]) -> Vec<String> {
    let mut ret = Vec::new();
    for rpath in rpaths.iter() {
        ret.push(format!("-Wl,-rpath,{}", (*rpath).as_slice()));
    }
    return ret;
}

fn get_rpaths(mut config: RPathConfig,
              libs: &[Path]) -> Vec<String> {
    debug!("output: {}", config.out_filename.display());
    debug!("libs:");
    for libpath in libs.iter() {
        debug!("    {}", libpath.display());
    }

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(&mut config, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = vec!(get_install_prefix_rpath(config));

    fn log_rpaths(desc: &str, rpaths: &[String]) {
        debug!("{} rpaths:", desc);
        for rpath in rpaths.iter() {
            debug!("    {}", *rpath);
        }
    }

    log_rpaths("relative", rel_rpaths.as_slice());
    log_rpaths("fallback", fallback_rpaths.as_slice());

    let mut rpaths = rel_rpaths;
    rpaths.push_all(fallback_rpaths.as_slice());

    // Remove duplicates
    let rpaths = minimize_rpaths(rpaths.as_slice());
    return rpaths;
}

fn get_rpaths_relative_to_output(config: &mut RPathConfig,
                                 libs: &[Path]) -> Vec<String> {
    libs.iter().map(|a| get_rpath_relative_to_output(config, a)).collect()
}

fn get_rpath_relative_to_output(config: &mut RPathConfig,
                                lib: &Path) -> String {
    use std::os;

    // Mac doesn't appear to support $ORIGIN
    let prefix = if config.is_like_osx {
        "@loader_path"
    } else {
        "$ORIGIN"
    };

    let mut lib = (config.realpath)(&os::make_absolute(lib).unwrap()).unwrap();
    lib.pop();
    let mut output = (config.realpath)(&os::make_absolute(&config.out_filename).unwrap()).unwrap();
    output.pop();
    let relative = lib.path_relative_from(&output);
    let relative = relative.expect("could not create rpath relative to output");
    // FIXME (#9639): This needs to handle non-utf8 paths
    format!("{}/{}",
            prefix,
            relative.as_str().expect("non-utf8 component in path"))
}

fn get_install_prefix_rpath(config: RPathConfig) -> String {
    let path = (config.get_install_prefix_lib_path)();
    let path = os::make_absolute(&path).unwrap();
    // FIXME (#9639): This needs to handle non-utf8 paths
    path.as_str().expect("non-utf8 component in rpath").to_string()
}

fn minimize_rpaths(rpaths: &[String]) -> Vec<String> {
    let mut set = HashSet::new();
    let mut minimized = Vec::new();
    for rpath in rpaths.iter() {
        if set.insert(rpath.as_slice()) {
            minimized.push(rpath.clone());
        }
    }
    minimized
}

#[cfg(all(unix, test))]
mod test {
    use super::{RPathConfig};
    use super::{minimize_rpaths, rpaths_to_flags, get_rpath_relative_to_output};
    use syntax::abi;

    #[test]
    fn test_rpaths_to_flags() {
        let flags = rpaths_to_flags(&[
            "path1".to_string(),
            "path2".to_string()
        ]);
        assert_eq!(flags,
                   ["-Wl,-rpath,path1",
                    "-Wl,-rpath,path2"]);
    }

    #[test]
    fn test_minimize1() {
        let res = minimize_rpaths(&[
            "rpath1".to_string(),
            "rpath2".to_string(),
            "rpath1".to_string()
        ]);
        assert!(res == [
            "rpath1",
            "rpath2",
        ]);
    }

    #[test]
    fn test_minimize2() {
        let res = minimize_rpaths(&[
            "1a".to_string(),
            "2".to_string(),
            "2".to_string(),
            "1a".to_string(),
            "4a".to_string(),
            "1a".to_string(),
            "2".to_string(),
            "3".to_string(),
            "4a".to_string(),
            "3".to_string()
        ]);
        assert!(res == [
            "1a",
            "2",
            "4a",
            "3",
        ]);
    }

    #[test]
    #[cfg(any(target_os = "linux", target_os = "android"))]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            has_rpath: true,
            is_like_osx: false,
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "freebsd")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: false,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "dragonfly")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: false,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "$ORIGIN/../lib");
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_rpath_relative() {
        let config = &mut RPathConfig {
            used_crates: Vec::new(),
            has_rpath: true,
            is_like_osx: true,
            out_filename: Path::new("bin/rustc"),
            get_install_prefix_lib_path: || panic!(),
            realpath: |p| Ok(p.clone())
        };
        let res = get_rpath_relative_to_output(config, &Path::new("lib/libstd.so"));
        assert_eq!(res, "@loader_path/../lib");
    }
}
<|endoftext|>


<commit_before>extern crate jsrs_parser;

use jsrs_parser::lalr::parse_Stmt;
use std::io;

fn main() {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();

    println!("{:?}", parse_Stmt(&input).unwrap());
}
<commit_msg>made binary more usable<commit_after>extern crate jsrs_parser;

use jsrs_parser::lalr::parse_Stmt;
use std::io::{self, Write};

fn main() {
    let stdin = io::stdin();
    let mut stdout = io::stdout();

    loop {
        print!(">> ");
        stdout.flush().unwrap();
        let mut input = String::new();
        stdin.read_line(&mut input).unwrap();

        if &input == ".exit\n" {
            return;
        }

        let stmt = match parse_Stmt(&input) {
            Ok(s) => s,
            _ => {
                println!("Error: failed to parse");
                continue
            }
        };

        println!("=> {:?}", stmt);
    }
}
<|endoftext|>


<commit_before>extern crate rav1e;
extern crate y4m;

use rav1e::*;

fn main() {
    let mut files = EncoderConfig::from_cli();
    let mut y4m_dec = y4m::decode(&mut files.input_file).unwrap();
    let width = y4m_dec.get_width();
    let height = y4m_dec.get_height();
    let framerate = y4m_dec.get_framerate();
    let mut y4m_enc = match files.rec_file.as_mut() {
        Some(rec_file) => Some(y4m::encode(width, height, framerate).write_header(rec_file).unwrap()),
        None => None
    };
    let mut fi = FrameInvariants::new(width, height);
    let sequence = Sequence::new();
    write_ivf_header(&mut files.output_file, fi.sb_width*64, fi.sb_height*64, framerate.num, framerate.den);

    let mut last_rec: Option<Frame> = None;
    loop {
        if !process_frame(&sequence, &fi, &mut files.output_file, &mut y4m_dec, y4m_enc.as_mut(), &mut last_rec) {
            break;
        }
        fi.number += 1;
        fi.show_existing_frame = fi.number % 2 == 1;
        if fi.number == files.limit {
            break;
        }
    }
}
<commit_msg>Flush output file on every packet write.<commit_after>extern crate rav1e;
extern crate y4m;

use rav1e::*;

fn main() {
    let mut files = EncoderConfig::from_cli();
    let mut y4m_dec = y4m::decode(&mut files.input_file).unwrap();
    let width = y4m_dec.get_width();
    let height = y4m_dec.get_height();
    let framerate = y4m_dec.get_framerate();
    let mut y4m_enc = match files.rec_file.as_mut() {
        Some(rec_file) => Some(y4m::encode(width, height, framerate).write_header(rec_file).unwrap()),
        None => None
    };
    let mut fi = FrameInvariants::new(width, height);
    let sequence = Sequence::new();
    write_ivf_header(&mut files.output_file, fi.sb_width*64, fi.sb_height*64, framerate.num, framerate.den);

    let mut last_rec: Option<Frame> = None;
    loop {
        if !process_frame(&sequence, &fi, &mut files.output_file, &mut y4m_dec, y4m_enc.as_mut(), &mut last_rec) {
            break;
        }
        fi.number += 1;
        fi.show_existing_frame = fi.number % 2 == 1;
        if fi.number == files.limit {
            break;
        }
        files.output_file.flush().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::fs::File;
use std::f64;

use image::{ DynamicImage, GenericImage, ImageFormat };
use image;

use bitreader::BitReader;

use num_traits::float::Float;

pub struct CoverImage {
    cover_image: DynamicImage
}

impl CoverImage {
    pub fn new(file_path: &str) -> CoverImage {
        CoverImage {
            cover_image: image::open(&Path::new(&file_path)).unwrap()
        }
    }

    pub fn encode_with(&self, message: &str) {
        let bit_vector = get_bit_vec(message);

        encode(&self.cover_image, &bit_vector);

        let ref mut fout = File::create(&Path::new("/home/hugh/Pictures/yurt.png")).unwrap();

        let _ = self.cover_image.save(fout, ImageFormat::PNG).unwrap();
    }
}

// Get a Vector of Bits to Encode the Image with
fn get_bit_vec(message: &str) -> Vec<u8> {
    let mut bit_vector = Vec::new();
    let mut reader = BitReader::new(&message.as_bytes());

    // Multiplied by 8 because it's a Vec of bits not bytes
    for _ in 0..message.len() * 8 {
        bit_vector.push(reader.read_u8(1).unwrap());
    }

    bit_vector
}

// Encode the image with the Vector of Bits
fn encode(c_image: &DynamicImage, bit_vec: &Vec<u8>) {
    let mut tiled_image_vec = tile_image(&c_image);

    //let (width, height) = c_image.dimensions();

    let encoded_image = encode_image(&mut tiled_image_vec, bit_vec);
}

fn tile_image(c_image: &DynamicImage) -> Vec<image::Rgba<u8>> {
    let (width, height) = c_image.dimensions();

    let mut image_blocks = Vec::new();

    for row_index in 0..(height / 8) as u32 {
        for col_index in 0..(width / 8) as u32 {
            for row in 0..8 {
                for column in 0..8 {
                    image_blocks.push(c_image.get_pixel(column + (col_index * 8), row + (row_index * 8)));
                }
            }
        }
    }

    image_blocks
}

fn encode_image(tiled_image: &mut Vec<image::Rgba<u8>>, message: &Vec<u8>) {
    let mut cu = 0.0;
    let mut cv = 0.0;
    let mut total = 0.0;
    let mut inverse_total = 0.0;
    let mut index: usize = 0;
    let mut count = 0;
    let mut rgb_coeffs: Vec<f64> = Vec::new();
    let mut colour_value: u8 = 0;

    let num_of_iterations = if message.len() / 3 == 0 {
                                    (message.len() / 3) as u32
                                } else {
                                    ((message.len() / 3) + 1) as u32
                                };

    for iteration in 0..num_of_iterations {
        for channel in 0..3 {
            for v in 0..8 {
                for u in 0..8 {
                    // Forward transform
                    for y in 0..8 {
                        for x in 0..8 {
                            index = u + (v * 8) + (iteration * 64) as usize;

                            colour_value = tiled_image[index].data[channel];

                            if u == 0 {
                                cu = 1.0 / 2.0.sqrt()
                            } else {
                                cu = 1.0
                            }

                            if v == 0 {
                                cv = 1.0 / 2.0.sqrt()
                            } else {
                                cv = 0.0
                            }

                            total = total + (v as f64 * f64::consts::PI * (2.0 * (y as f64) + 1.0) / 16.0).cos()
                                          * (u as f64 * f64::consts::PI * (2.0 * (x as f64) + 1.0) / 16.0).cos()
                                          * colour_value as f64;
                        }
                    }

                    // 0.25, Cu and Cv are scaling factors
                    total = total * 0.25 * cu * cv;

                    rgb_coeffs.push(total);
                }
            }

            // Encode the message
            let coeff_to_mod = (27 + (iteration * 64)) as usize;

            if rgb_coeffs[coeff_to_mod].trunc() % 2.0 == 0.0 {
                if message[count] == 1 {
                    rgb_coeffs[coeff_to_mod] = rgb_coeffs[coeff_to_mod] + 1.0
                }
            } else {
                if message[count] == 0 {
                    rgb_coeffs[coeff_to_mod] = rgb_coeffs[coeff_to_mod] + 1.0
                }
            }

            // Inverse transform
            for v in 0..8 {
                for u in 0..8 {
                    if u == 0 {
                        cu = 0.0
                    } else {
                        cu = 1.0 / 2.0.sqrt()
                    }

                    if v == 0 {
                        cv = 0.0
                    } else {
                        cv = 1.0 / 2.0.sqrt()
                    }

                    inverse_total = inverse_total + (v as f64 * f64::consts::PI * (2.0 * (y as f64) + 1.0) / 16.0).cos()
                                                  * (u as f64 * f64::consts::PI * (2.0 * (x as f64) + 1.0) / 16.0).cos()
                                                  * colour_value as f64;
                }
            }
        }
    }
}

/**************************************************************************************************
 *                                                                                *
 *                                         Tests                                                  *
 *                                                                                *
 **************************************************************************************************/

#[cfg(test)]
#[test]
fn test_get_bit_vec_len() {
    let bits = get_bit_vec("h");

    assert_eq!(8, bits.len());
}

#[cfg(test)]
#[test]
fn test_get_bit_vec() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let bits = get_bit_vec("h");

    assert_eq!(bits, test_bits);
}

#[cfg(test)]
#[test]
fn test_encode_red_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[0] % 2, bit_vec[0]);
}

#[cfg(test)]
#[test]
fn test_encode_green_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[1] % 2, bit_vec[1]);
}

#[cfg(test)]
#[test]
fn test_encode_blue_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[2] % 2 as u8, bit_vec[2]);
}

#[cfg(test)]
#[test]
fn test_full_byte_is_encoded() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let mut encoded_bit_vec = Vec::new();
    let mut count = 0;

    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    'outer: for x_co_ord in 0..3 {
        let pixel = img.get_pixel(x_co_ord, 0);

        for channel in 0..3 {
            if count >= 8 {
                break 'outer;
            }

            encoded_bit_vec.push(pixel.data[channel] % 2);

            count += 1;
        }
    }

    assert_eq!(encoded_bit_vec, bit_vec);
}

#[cfg(test)]
#[test]
fn test_tile_image_length() {
    let img = image::open(&Path::new("/home/hugh/Pictures/colour.jpg")).unwrap();

    let vec = tile_image(&img);

    assert_eq!(vec.len(), 91200);
}

#[cfg(test)]
#[test]
fn test_tile_image_first_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(0, 0);

    let vec = tile_image(&img);

    assert_eq!(vec[0], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel1() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel2() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(40, 16);

    let vec = tile_image(&img);

    assert_eq!(vec[6720], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel3() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_last_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(399, 399);

    let vec = tile_image(&img);

    assert_eq!(vec[159999], pixel);
}
<commit_msg>More work for idct<commit_after>use std::path::Path;
use std::fs::File;
use std::f64;

use image::{ DynamicImage, GenericImage, ImageFormat };
use image;

use bitreader::BitReader;

use num_traits::float::Float;

pub struct CoverImage {
    cover_image: DynamicImage
}

impl CoverImage {
    pub fn new(file_path: &str) -> CoverImage {
        CoverImage {
            cover_image: image::open(&Path::new(&file_path)).unwrap()
        }
    }

    pub fn encode_with(&self, message: &str) {
        let bit_vector = get_bit_vec(message);

        encode(&self.cover_image, &bit_vector);

        let ref mut fout = File::create(&Path::new("/home/hugh/Pictures/yurt.png")).unwrap();

        let _ = self.cover_image.save(fout, ImageFormat::PNG).unwrap();
    }
}

// Get a Vector of Bits to Encode the Image with
fn get_bit_vec(message: &str) -> Vec<u8> {
    let mut bit_vector = Vec::new();
    let mut reader = BitReader::new(&message.as_bytes());

    // Multiplied by 8 because it's a Vec of bits not bytes
    for _ in 0..message.len() * 8 {
        bit_vector.push(reader.read_u8(1).unwrap());
    }

    bit_vector
}

// Encode the image with the Vector of Bits
fn encode(c_image: &DynamicImage, bit_vec: &Vec<u8>) {
    let mut tiled_image_vec = tile_image(&c_image);

    //let (width, height) = c_image.dimensions();

    let encoded_image = encode_image(&mut tiled_image_vec, bit_vec);
}

fn tile_image(c_image: &DynamicImage) -> Vec<image::Rgba<u8>> {
    let (width, height) = c_image.dimensions();

    let mut image_blocks = Vec::new();

    for row_index in 0..(height / 8) as u32 {
        for col_index in 0..(width / 8) as u32 {
            for row in 0..8 {
                for column in 0..8 {
                    image_blocks.push(c_image.get_pixel(column + (col_index * 8), row + (row_index * 8)));
                }
            }
        }
    }

    image_blocks
}

fn encode_image(tiled_image: &mut Vec<image::Rgba<u8>>, message: &Vec<u8>) {
    let mut cu = 0.0;
    let mut cv = 0.0;
    let mut total = 0.0;
    let mut inverse_total = 0.0;
    let mut index: usize = 0;
    let mut count = 0;
    let mut dct_coeffs: Vec<f64> = Vec::new();
    let mut colour_value: u8 = 0;

    let num_of_iterations = if message.len() / 3 == 0 {
                                    (message.len() / 3) as u32
                                } else {
                                    ((message.len() / 3) + 1) as u32
                                };

    for iteration in 0..num_of_iterations {
        for channel in 0..3 {
            for v in 0..8 {
                for u in 0..8 {
                    // Forward transform
                    for y in 0..8 {
                        for x in 0..8 {
                            index = u + (v * 8) + (iteration * 64) as usize;

                            colour_value = tiled_image[index].data[channel];

                            if u == 0 {
                                cu = 1.0 / 2.0.sqrt()
                            } else {
                                cu = 1.0
                            }

                            if v == 0 {
                                cv = 1.0 / 2.0.sqrt()
                            } else {
                                cv = 0.0
                            }

                            total = total + (v as f64 * f64::consts::PI * (2.0 * (y as f64) + 1.0) / 16.0).cos()
                                          * (u as f64 * f64::consts::PI * (2.0 * (x as f64) + 1.0) / 16.0).cos()
                                          * colour_value as f64;
                        }
                    }

                    // 0.25, Cu and Cv are scaling factors
                    total = total * 0.25 * cu * cv;

                    dct_coeffs.push(total);
                }
            }

            // Encode the message
            let coeff_to_mod = (27 + (iteration * 64)) as usize;

            if dct_coeffs[coeff_to_mod].trunc() % 2.0 == 0.0 {
                if message[count] == 1 {
                    dct_coeffs[coeff_to_mod] = dct_coeffs[coeff_to_mod] + 1.0
                }
            } else {
                if message[count] == 0 {
                    dct_coeffs[coeff_to_mod] = dct_coeffs[coeff_to_mod] + 1.0
                }
            }

            // Inverse transform
            for v in 0..8 {
                for u in 0..8 {
                    for y in 0..8 {
                        for x in 0..8 {
                            index = u + (v * 8) + (iteration * 64) as usize;

                            colour_value = tiled_image[index].data[channel];

                            if u == 0 {
                                cu = 1.0 / 2.0.sqrt()
                            } else {
                                cu = 1.0
                            }

                            if v == 0 {
                                cv = 1.0 / 2.0.sqrt()
                            } else {
                                cv = 0.0
                            }

                            inverse_total = inverse_total * cu * cv + (v as f64 * f64::consts::PI * (2.0 * (y as f64) + 1.0) / 16.0).cos()
                                                                    * (u as f64 * f64::consts::PI * (2.0 * (x as f64) + 1.0) / 16.0).cos()
                                                                    * colour_value as f64;
                        }
                    }
                }
            }

            inverse_total = inverse_total * 0.25;
        }
    }
}

/**************************************************************************************************
 *                                                                                *
 *                                         Tests                                                  *
 *                                                                                *
 **************************************************************************************************/

#[cfg(test)]
#[test]
fn test_get_bit_vec_len() {
    let bits = get_bit_vec("h");

    assert_eq!(8, bits.len());
}

#[cfg(test)]
#[test]
fn test_get_bit_vec() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let bits = get_bit_vec("h");

    assert_eq!(bits, test_bits);
}

#[cfg(test)]
#[test]
fn test_encode_red_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[0] % 2, bit_vec[0]);
}

#[cfg(test)]
#[test]
fn test_encode_green_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[1] % 2, bit_vec[1]);
}

#[cfg(test)]
#[test]
fn test_encode_blue_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[2] % 2 as u8, bit_vec[2]);
}

#[cfg(test)]
#[test]
fn test_full_byte_is_encoded() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let mut encoded_bit_vec = Vec::new();
    let mut count = 0;

    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    'outer: for x_co_ord in 0..3 {
        let pixel = img.get_pixel(x_co_ord, 0);

        for channel in 0..3 {
            if count >= 8 {
                break 'outer;
            }

            encoded_bit_vec.push(pixel.data[channel] % 2);

            count += 1;
        }
    }

    assert_eq!(encoded_bit_vec, bit_vec);
}

#[cfg(test)]
#[test]
fn test_tile_image_length() {
    let img = image::open(&Path::new("/home/hugh/Pictures/colour.jpg")).unwrap();

    let vec = tile_image(&img);

    assert_eq!(vec.len(), 91200);
}

#[cfg(test)]
#[test]
fn test_tile_image_first_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(0, 0);

    let vec = tile_image(&img);

    assert_eq!(vec[0], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel1() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel2() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(40, 16);

    let vec = tile_image(&img);

    assert_eq!(vec[6720], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel3() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_last_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(399, 399);

    let vec = tile_image(&img);

    assert_eq!(vec[159999], pixel);
}
<|endoftext|>


<commit_before>//! Interval-based annotation features.

use std::error::Error;
use std::collections::HashMap;
use std::fmt::{self, Display};

use bio::io::Strand;
use bio::utils::{Interval, IntervalError};

use self::error::FeatureError;

pub mod error {

    use super::*;

    #[derive(Debug, PartialEq)]
    pub enum FeatureError {
        IntervalError,
        StrandCharError,
        ConflictingStrandError,
        UnspecifiedStrandError,
    }

    impl Error for FeatureError {

        fn description(&self) -> &str {
            match *self {
                FeatureError::IntervalError => "interval start coordinate larger than its end coordinate",
                FeatureError::StrandCharError => "strand character is invalid",
                FeatureError::ConflictingStrandError => "conflicting strand inputs specified",
                FeatureError::UnspecifiedStrandError => "strand not specified",
            }
        }

        fn cause(&self) -> Option<&Error> {
            None
        }
    }

    impl Display for FeatureError {

        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "FeatureError: {}", self.description())
        }
    }
}

trait Annotation {
    fn seq_name(&self) -> &str;
    fn interval(&self) -> &Interval<u64>;
    fn attributes(&self) -> &HashMap<String, String>;
    fn attribute(&self, key: &str) -> Option<&str>;
    fn strand(&self) -> &Strand;
}

macro_rules! impl_annotation {
    ($struct_ty:ty) => (

        impl Annotation for $struct_ty {

            fn seq_name(&self) -> &str {
                self.seq_name.as_str()
            }

            fn interval(&self) -> &Interval<u64> {
                &self.interval
            }

            fn attributes(&self) -> &HashMap<String, String> {
                &self.attributes
            }

            fn attribute(&self, key: &str) -> Option<&str> {
                self.attributes.get(key).map(|n| n.as_str())
            }

            fn strand(&self) -> &Strand {
                &self.strand
            }
        }

    );
}

fn char_to_strand(strand_char: char) -> Result<Strand, FeatureError> {
    match strand_char {
        '+' | 'f' | 'F' => Ok(Strand::Forward),
        '-' | 'r' | 'R' => Ok(Strand::Reverse),
        '.' | '?' => Ok(Strand::Unknown),
        _ => Err(FeatureError::StrandCharError),
    }
}

fn resolve_strand_input(strand: Option<Strand>, strand_char: Option<char>) -> Result<Strand, FeatureError> {
    match (strand, strand_char) {
        (None, None) => Err(FeatureError::UnspecifiedStrandError),
        (Some(sv), None) => Ok(sv),
        (None, Some(scv)) => char_to_strand(scv),
        (Some(sv), Some(scv)) => {
            let sv_from_char = char_to_strand(scv)?;
            if sv == sv_from_char {
                Ok(sv)
            } else {
                Err(FeatureError::ConflictingStrandError)
            }
        }
    }
}

fn coords_to_interval(start: u64, end: u64) -> Result<Interval<u64>, FeatureError> {
    Interval::new(start..end)
        .map_err(|err| match err {
            IntervalError::InvalidRange => FeatureError::IntervalError
        })
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TxFeature {
    Exon,
    UTR,
    UTR5,
    UTR3,
    CDS,
    StartCodon,
    StopCodon,
    Any,
}

#[derive(Debug)]
pub struct TranscriptFeature {
    kind: TxFeature,
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>
}

impl_annotation!(TranscriptFeature);

impl TranscriptFeature {

    fn kind(&self) -> &TxFeature {
        &self.kind
    }
}

pub struct TxFeatureBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    attributes: HashMap<String, String>,
    kind: TxFeature,
    strand: Option<Strand>,
    strand_char: Option<char>,
}

impl TxFeatureBuilder {

    fn new<T>(seq_name: T, start: u64, end: u64) -> TxFeatureBuilder
        where T: Into<String>
    {
        TxFeatureBuilder {
            seq_name: seq_name.into(),
            start: start, end: end,
            kind: TxFeature::Any,
            attributes: HashMap::new(),
            strand: None,
            strand_char: None,
        }
    }

    fn kind(mut self, kind: TxFeature) -> TxFeatureBuilder {
        self.kind = kind;
        self
    }

    fn strand(mut self, strand: Strand) -> TxFeatureBuilder {
        self.strand = Some(strand);
        self
    }

    fn strand_char(mut self, strand_char: char) -> TxFeatureBuilder {
        self.strand_char = Some(strand_char);
        self
    }

    fn attribute<K, V>(mut self, key: K, value: V) -> TxFeatureBuilder
        where K: Into<String>, V: Into<String>
    {
        self.attributes.insert(key.into(), value.into());
        self
    }

    fn build(self) -> Result<TranscriptFeature, FeatureError> {
        let interval = coords_to_interval(self.start, self.end)?;
        let strand = resolve_strand_input(self.strand, self.strand_char)?;
        let feature = TranscriptFeature {
            seq_name: self.seq_name, kind: self.kind, interval: interval,
            strand: strand, attributes: self.attributes,
        };
        Ok(feature)
    }
}

#[derive(Debug)]
pub struct Transcript {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>,
    features: Vec<TranscriptFeature>,
}

impl_annotation!(Transcript);

#[derive(Debug)]
pub struct Gene {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>,
    transcripts: HashMap<String, Transcript>,
}

impl_annotation!(Gene);

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

    #[test]
    fn builder() {
        let tfm1 = TxFeatureBuilder::new("chrT", 10, 20)
            .kind(TxFeature::Exon)
            .strand(Strand::Forward)
            .attribute("name", "ex1")
            .build();
        assert!(tfm1.is_ok());
        let tf = tfm1.unwrap();
        assert_eq!(tf.seq_name(), "chrT");
        assert_eq!(tf.kind(), &TxFeature::Exon);
        assert_eq!(tf.strand(), &Strand::Forward);
        assert_eq!(tf.attribute("name"), Some("ex1"));
        assert_eq!(tf.attributes.len(), 1);

        let tfm2 = TxFeatureBuilder::new("chrO", 10, 10)
            .strand_char('-')
            .strand(Strand::Reverse)
            .build();
        assert!(tfm2.is_ok());
    }

    #[test]
    fn builder_interval_invalid() {
        let tfm = TxFeatureBuilder::new("chrE", 20, 10).build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::IntervalError);
    }

    #[test]
    fn builder_strand_unspecified() {
        let tfm = TxFeatureBuilder::new("chrT", 20, 30)
            .build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::UnspecifiedStrandError);
    }

    #[test]
    fn builder_strand_char_unexpected() {
        let tfm = TxFeatureBuilder::new("chrE", 10, 20)
            .strand_char('w')
            .build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::StrandCharError);
    }

    #[test]
    fn builder_strand_char_conflicting() {
        let tfm = TxFeatureBuilder::new("chrE", 10, 20)
            .strand_char('-')
            .strand(Strand::Reverse)
            .build();
        assert!(tfm.is_ok());
        let tf = tfm.unwrap();
        assert_eq!(tf.strand(), &Strand::Reverse);
    }
}
<commit_msg>Temporary implement strand<commit_after>//! Interval-based annotation features.

use std::error::Error;
use std::collections::HashMap;
use std::fmt::{self, Display};

use bio::utils::{Interval, IntervalError};

use self::error::FeatureError;

// TODO: PR for rust-bio so that it implements Copy
//       for now, we use our own
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Strand {
    Forward,
    Reverse,
    Unknown
}

pub mod error {

    use super::*;

    #[derive(Debug, PartialEq)]
    pub enum FeatureError {
        IntervalError,
        StrandCharError,
        ConflictingStrandError,
        UnspecifiedStrandError,
    }

    impl Error for FeatureError {

        fn description(&self) -> &str {
            match *self {
                FeatureError::IntervalError => "interval start coordinate larger than its end coordinate",
                FeatureError::StrandCharError => "strand character is invalid",
                FeatureError::ConflictingStrandError => "conflicting strand inputs specified",
                FeatureError::UnspecifiedStrandError => "strand not specified",
            }
        }

        fn cause(&self) -> Option<&Error> {
            None
        }
    }

    impl Display for FeatureError {

        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "FeatureError: {}", self.description())
        }
    }
}

trait Annotation {
    fn seq_name(&self) -> &str;
    fn interval(&self) -> &Interval<u64>;
    fn attributes(&self) -> &HashMap<String, String>;
    fn attribute(&self, key: &str) -> Option<&str>;
    fn strand(&self) -> &Strand;
}

macro_rules! impl_annotation {
    ($struct_ty:ty) => (

        impl Annotation for $struct_ty {

            fn seq_name(&self) -> &str {
                self.seq_name.as_str()
            }

            fn interval(&self) -> &Interval<u64> {
                &self.interval
            }

            fn attributes(&self) -> &HashMap<String, String> {
                &self.attributes
            }

            fn attribute(&self, key: &str) -> Option<&str> {
                self.attributes.get(key).map(|n| n.as_str())
            }

            fn strand(&self) -> &Strand {
                &self.strand
            }
        }

    );
}

fn char_to_strand(strand_char: char) -> Result<Strand, FeatureError> {
    match strand_char {
        '+' | 'f' | 'F' => Ok(Strand::Forward),
        '-' | 'r' | 'R' => Ok(Strand::Reverse),
        '.' | '?' => Ok(Strand::Unknown),
        _ => Err(FeatureError::StrandCharError),
    }
}

fn resolve_strand_input(strand: Option<Strand>, strand_char: Option<char>) -> Result<Strand, FeatureError> {
    match (strand, strand_char) {
        (None, None) => Err(FeatureError::UnspecifiedStrandError),
        (Some(sv), None) => Ok(sv),
        (None, Some(scv)) => char_to_strand(scv),
        (Some(sv), Some(scv)) => {
            let sv_from_char = char_to_strand(scv)?;
            if sv == sv_from_char {
                Ok(sv)
            } else {
                Err(FeatureError::ConflictingStrandError)
            }
        }
    }
}

fn coords_to_interval(start: u64, end: u64) -> Result<Interval<u64>, FeatureError> {
    Interval::new(start..end)
        .map_err(|err| match err {
            IntervalError::InvalidRange => FeatureError::IntervalError
        })
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TxFeature {
    Exon,
    UTR,
    UTR5,
    UTR3,
    CDS,
    StartCodon,
    StopCodon,
    Any,
}

#[derive(Debug)]
pub struct TranscriptFeature {
    kind: TxFeature,
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>
}

impl_annotation!(TranscriptFeature);

impl TranscriptFeature {

    fn kind(&self) -> &TxFeature {
        &self.kind
    }
}

pub struct TxFeatureBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    attributes: HashMap<String, String>,
    kind: TxFeature,
    strand: Option<Strand>,
    strand_char: Option<char>,
}

impl TxFeatureBuilder {

    fn new<T>(seq_name: T, start: u64, end: u64) -> TxFeatureBuilder
        where T: Into<String>
    {
        TxFeatureBuilder {
            seq_name: seq_name.into(),
            start: start, end: end,
            kind: TxFeature::Any,
            attributes: HashMap::new(),
            strand: None,
            strand_char: None,
        }
    }

    fn kind(mut self, kind: TxFeature) -> TxFeatureBuilder {
        self.kind = kind;
        self
    }

    fn strand(mut self, strand: Strand) -> TxFeatureBuilder {
        self.strand = Some(strand);
        self
    }

    fn strand_char(mut self, strand_char: char) -> TxFeatureBuilder {
        self.strand_char = Some(strand_char);
        self
    }

    fn attribute<K, V>(mut self, key: K, value: V) -> TxFeatureBuilder
        where K: Into<String>, V: Into<String>
    {
        self.attributes.insert(key.into(), value.into());
        self
    }

    fn build(self) -> Result<TranscriptFeature, FeatureError> {
        let interval = coords_to_interval(self.start, self.end)?;
        let strand = resolve_strand_input(self.strand, self.strand_char)?;
        let feature = TranscriptFeature {
            seq_name: self.seq_name, kind: self.kind, interval: interval,
            strand: strand, attributes: self.attributes,
        };
        Ok(feature)
    }
}

#[derive(Debug)]
pub struct Transcript {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>,
    features: Vec<TranscriptFeature>,
}

impl_annotation!(Transcript);

#[derive(Debug)]
pub struct Gene {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    attributes: HashMap<String, String>,
    transcripts: HashMap<String, Transcript>,
}

impl_annotation!(Gene);

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

    #[test]
    fn builder() {
        let tfm1 = TxFeatureBuilder::new("chrT", 10, 20)
            .kind(TxFeature::Exon)
            .strand(Strand::Forward)
            .attribute("name", "ex1")
            .build();
        assert!(tfm1.is_ok());
        let tf = tfm1.unwrap();
        assert_eq!(tf.seq_name(), "chrT");
        assert_eq!(tf.kind(), &TxFeature::Exon);
        assert_eq!(tf.strand(), &Strand::Forward);
        assert_eq!(tf.attribute("name"), Some("ex1"));
        assert_eq!(tf.attributes.len(), 1);

        let tfm2 = TxFeatureBuilder::new("chrO", 10, 10)
            .strand_char('-')
            .strand(Strand::Reverse)
            .build();
        assert!(tfm2.is_ok());
    }

    #[test]
    fn builder_interval_invalid() {
        let tfm = TxFeatureBuilder::new("chrE", 20, 10).build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::IntervalError);
    }

    #[test]
    fn builder_strand_unspecified() {
        let tfm = TxFeatureBuilder::new("chrT", 20, 30)
            .build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::UnspecifiedStrandError);
    }

    #[test]
    fn builder_strand_char_unexpected() {
        let tfm = TxFeatureBuilder::new("chrE", 10, 20)
            .strand_char('w')
            .build();
        assert!(tfm.is_err());
        assert_eq!(tfm.unwrap_err(), FeatureError::StrandCharError);
    }

    #[test]
    fn builder_strand_char_conflicting() {
        let tfm = TxFeatureBuilder::new("chrE", 10, 20)
            .strand_char('-')
            .strand(Strand::Reverse)
            .build();
        assert!(tfm.is_ok());
        let tf = tfm.unwrap();
        assert_eq!(tf.strand(), &Strand::Reverse);
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc, Mutex,
};

use futures::{future::FutureExt, select_biased};
use tbot::{
    connectors::Connector,
    types::parameters::{self, WebPagePreviewState},
    Bot,
};
use tokio::{
    self,
    stream::StreamExt,
    sync::Notify,
    time::{self, delay_for, delay_queue::DelayQueue, Duration, Instant},
};

use crate::client::pull_feed;
use crate::data::{Database, Feed, FeedUpdate};
use crate::messages::{format_large_msg, Escape};

pub fn start(
    bot: Bot<impl Connector>,
    db: Arc<Mutex<Database>>,
    min_interval: u32,
    max_interval: u32,
) {
    let mut queue = FetchQueue::new();
    // TODO: Don't use interval, it can accumulate ticks
    // replace it with delay_until
    let mut interval = time::interval_at(Instant::now(), Duration::from_secs(min_interval as u64));
    let throttle = Throttle::new(min_interval as usize);
    tokio::spawn(async move {
        loop {
            select_biased! {
                feed = queue.next().fuse() => {
                    let feed = feed.expect("unreachable");
                    let bot = bot.clone();
                    let db = db.clone();
                    let opportunity = throttle.acquire();
                    tokio::spawn(async move {
                        opportunity.wait().await;
                        if let Err(e) = fetch_and_push_updates(bot, db, feed).await {
                            dbg!(e);
                        }
                    });
                }
                _ = interval.tick().fuse() => {
                    let feeds = db.lock().unwrap().all_feeds();
                    for feed in feeds {
                        let feed_interval = cmp::min(
                            feed.ttl.map(|ttl| ttl * 60).unwrap_or(min_interval),
                            max_interval,
                        ) as u64 - 1; // after -1, we can stagger with `interval`
                        queue.enqueue(feed, Duration::from_secs(feed_interval));
                    }
                }
            }
        }
    });
}

async fn fetch_and_push_updates(
    bot: Bot<impl Connector>,
    db: Arc<Mutex<Database>>,
    feed: Feed,
) -> Result<(), tbot::errors::MethodCall> {
    let new_feed = match pull_feed(&feed.link).await {
        Ok(feed) => feed,
        Err(e) => {
            let down_time = db.lock().unwrap().get_or_update_down_time(&feed.link);
            // 5 days
            if down_time.as_secs() > 5 * 24 * 60 * 60 {
                let msg = format!(
                    "《{}》\
                     已经连续 5 天拉取出错 ({}),\
                     可能已经关闭, 请取消订阅",
                    Escape(&feed.link),
                    Escape(&feed.title),
                    Escape(&e.to_string())
                );
                push_updates(&bot, &db, feed.subscribers, parameters::Text::html(&msg)).await?;
            }
            return Ok(());
        }
    };

    let updates = db.lock().unwrap().update(&feed.link, new_feed);
    for update in updates {
        match update {
            FeedUpdate::Items(items) => {
                let msgs =
                    format_large_msg(format!("{}", Escape(&feed.title)), &items, |item| {
                        let title = item
                            .title
                            .as_ref()
                            .map(|s| s.as_str())
                            .unwrap_or_else(|| &feed.title);
                        let link = item
                            .link
                            .as_ref()
                            .map(|s| s.as_str())
                            .unwrap_or_else(|| &feed.link);
                        format!("{}", Escape(link), Escape(title))
                    });
                for msg in msgs {
                    push_updates(
                        &bot,
                        &db,
                        feed.subscribers.iter().copied(),
                        parameters::Text::html(&msg),
                    )
                    .await?;
                }
            }
            FeedUpdate::Title(new_title) => {
                let msg = format!(
                    "{} 已更名为 {}",
                    Escape(&feed.link),
                    Escape(&feed.title),
                    Escape(&new_title)
                );
                push_updates(
                    &bot,
                    &db,
                    feed.subscribers.iter().copied(),
                    parameters::Text::html(&msg),
                )
                .await?;
            }
        }
    }
    Ok(())
}

async fn push_updates<I: IntoIterator<Item = i64>>(
    bot: &Bot<impl Connector>,
    db: &Arc<Mutex<Database>>,
    subscribers: I,
    msg: parameters::Text<'_>,
) -> Result<(), tbot::errors::MethodCall> {
    use tbot::errors::MethodCall;
    for mut subscriber in subscribers {
        'retry: for _ in 0..3 {
            match bot
                .send_message(tbot::types::chat::Id(subscriber), msg)
                .web_page_preview(WebPagePreviewState::Disabled)
                .call()
                .await
            {
                Err(MethodCall::RequestError { description, .. })
                    if chat_is_unavailable(&description) =>
                {
                    db.lock().unwrap().delete_subscriber(subscriber);
                }
                Err(MethodCall::RequestError {
                    migrate_to_chat_id: Some(new_chat_id),
                    ..
                }) => {
                    db.lock()
                        .unwrap()
                        .update_subscriber(subscriber, new_chat_id.0);
                    subscriber = new_chat_id.0;
                    continue 'retry;
                }
                Err(MethodCall::RequestError {
                    retry_after: Some(delay),
                    ..
                }) => {
                    time::delay_for(Duration::from_secs(delay)).await;
                    continue 'retry;
                }
                other => {
                    other?;
                }
            }
            break 'retry;
        }
    }
    Ok(())
}

pub fn chat_is_unavailable(s: &str) -> bool {
    s.contains("Forbidden") || s.contains("chat not found")
}

#[derive(Default)]
struct FetchQueue {
    feeds: HashMap<String, Feed>,
    notifies: DelayQueue<String>,
    wakeup: Notify,
}

impl FetchQueue {
    fn new() -> Self {
        Self::default()
    }

    fn enqueue(&mut self, feed: Feed, delay: Duration) -> bool {
        let exists = self.feeds.contains_key(&feed.link);
        if !exists {
            self.notifies.insert(feed.link.clone(), delay);
            self.feeds.insert(feed.link.clone(), feed);
            self.wakeup.notify();
        }
        !exists
    }

    async fn next(&mut self) -> Result<Feed, time::Error> {
        loop {
            if let Some(feed_id) = self.notifies.next().await {
                let feed = self.feeds.remove(feed_id?.get_ref()).unwrap();
                break Ok(feed);
            } else {
                self.wakeup.notified().await;
            }
        }
    }
}

struct Throttle {
    pieces: usize,
    counter: Arc<AtomicUsize>,
}

impl Throttle {
    fn new(pieces: usize) -> Self {
        Throttle {
            pieces,
            counter: Arc::new(AtomicUsize::new(0)),
        }
    }

    fn acquire(&self) -> Opportunity {
        Opportunity {
            n: self.counter.fetch_add(1, Ordering::AcqRel) % self.pieces,
            counter: self.counter.clone(),
        }
    }
}

#[must_use = "Don't lose your opportunity"]
struct Opportunity {
    n: usize,
    counter: Arc<AtomicUsize>,
}

impl Opportunity {
    async fn wait(&self) {
        delay_for(Duration::from_secs(self.n as u64)).await
    }
}

impl Drop for Opportunity {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}
<commit_msg>Don't send fetch error message too frequently<commit_after>use std::cmp;
use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc, Mutex,
};

use futures::{future::FutureExt, select_biased};
use tbot::{
    connectors::Connector,
    types::parameters::{self, WebPagePreviewState},
    Bot,
};
use tokio::{
    self,
    stream::StreamExt,
    sync::Notify,
    time::{self, delay_for, delay_queue::DelayQueue, Duration, Instant},
};

use crate::client::pull_feed;
use crate::data::{Database, Feed, FeedUpdate};
use crate::messages::{format_large_msg, Escape};

pub fn start(
    bot: Bot<impl Connector>,
    db: Arc<Mutex<Database>>,
    min_interval: u32,
    max_interval: u32,
) {
    let mut queue = FetchQueue::new();
    // TODO: Don't use interval, it can accumulate ticks
    // replace it with delay_until
    let mut interval = time::interval_at(Instant::now(), Duration::from_secs(min_interval as u64));
    let throttle = Throttle::new(min_interval as usize);
    tokio::spawn(async move {
        loop {
            select_biased! {
                feed = queue.next().fuse() => {
                    let feed = feed.expect("unreachable");
                    let bot = bot.clone();
                    let db = db.clone();
                    let opportunity = throttle.acquire();
                    tokio::spawn(async move {
                        opportunity.wait().await;
                        if let Err(e) = fetch_and_push_updates(bot, db, feed).await {
                            dbg!(e);
                        }
                    });
                }
                _ = interval.tick().fuse() => {
                    let feeds = db.lock().unwrap().all_feeds();
                    for feed in feeds {
                        let feed_interval = cmp::min(
                            feed.ttl.map(|ttl| ttl * 60).unwrap_or(min_interval),
                            max_interval,
                        ) as u64 - 1; // after -1, we can stagger with `interval`
                        queue.enqueue(feed, Duration::from_secs(feed_interval));
                    }
                }
            }
        }
    });
}

async fn fetch_and_push_updates(
    bot: Bot<impl Connector>,
    db: Arc<Mutex<Database>>,
    feed: Feed,
) -> Result<(), tbot::errors::MethodCall> {
    let new_feed = match pull_feed(&feed.link).await {
        Ok(feed) => feed,
        Err(e) => {
            let down_time = db.lock().unwrap().get_or_update_down_time(&feed.link);
            // 5 days
            if down_time.as_secs() > 5 * 24 * 60 * 60 {
                db.lock().unwrap().reset_down_time(&feed.link);
                let msg = format!(
                    "《{}》\
                     已经连续 5 天拉取出错 ({}),\
                     可能已经关闭, 请取消订阅",
                    Escape(&feed.link),
                    Escape(&feed.title),
                    Escape(&e.to_string())
                );
                push_updates(&bot, &db, feed.subscribers, parameters::Text::html(&msg)).await?;
            }
            return Ok(());
        }
    };

    let updates = db.lock().unwrap().update(&feed.link, new_feed);
    for update in updates {
        match update {
            FeedUpdate::Items(items) => {
                let msgs =
                    format_large_msg(format!("{}", Escape(&feed.title)), &items, |item| {
                        let title = item
                            .title
                            .as_ref()
                            .map(|s| s.as_str())
                            .unwrap_or_else(|| &feed.title);
                        let link = item
                            .link
                            .as_ref()
                            .map(|s| s.as_str())
                            .unwrap_or_else(|| &feed.link);
                        format!("{}", Escape(link), Escape(title))
                    });
                for msg in msgs {
                    push_updates(
                        &bot,
                        &db,
                        feed.subscribers.iter().copied(),
                        parameters::Text::html(&msg),
                    )
                    .await?;
                }
            }
            FeedUpdate::Title(new_title) => {
                let msg = format!(
                    "{} 已更名为 {}",
                    Escape(&feed.link),
                    Escape(&feed.title),
                    Escape(&new_title)
                );
                push_updates(
                    &bot,
                    &db,
                    feed.subscribers.iter().copied(),
                    parameters::Text::html(&msg),
                )
                .await?;
            }
        }
    }
    Ok(())
}

async fn push_updates<I: IntoIterator<Item = i64>>(
    bot: &Bot<impl Connector>,
    db: &Arc<Mutex<Database>>,
    subscribers: I,
    msg: parameters::Text<'_>,
) -> Result<(), tbot::errors::MethodCall> {
    use tbot::errors::MethodCall;
    for mut subscriber in subscribers {
        'retry: for _ in 0..3 {
            match bot
                .send_message(tbot::types::chat::Id(subscriber), msg)
                .web_page_preview(WebPagePreviewState::Disabled)
                .call()
                .await
            {
                Err(MethodCall::RequestError { description, .. })
                    if chat_is_unavailable(&description) =>
                {
                    db.lock().unwrap().delete_subscriber(subscriber);
                }
                Err(MethodCall::RequestError {
                    migrate_to_chat_id: Some(new_chat_id),
                    ..
                }) => {
                    db.lock()
                        .unwrap()
                        .update_subscriber(subscriber, new_chat_id.0);
                    subscriber = new_chat_id.0;
                    continue 'retry;
                }
                Err(MethodCall::RequestError {
                    retry_after: Some(delay),
                    ..
                }) => {
                    time::delay_for(Duration::from_secs(delay)).await;
                    continue 'retry;
                }
                other => {
                    other?;
                }
            }
            break 'retry;
        }
    }
    Ok(())
}

pub fn chat_is_unavailable(s: &str) -> bool {
    s.contains("Forbidden") || s.contains("chat not found")
}

#[derive(Default)]
struct FetchQueue {
    feeds: HashMap<String, Feed>,
    notifies: DelayQueue<String>,
    wakeup: Notify,
}

impl FetchQueue {
    fn new() -> Self {
        Self::default()
    }

    fn enqueue(&mut self, feed: Feed, delay: Duration) -> bool {
        let exists = self.feeds.contains_key(&feed.link);
        if !exists {
            self.notifies.insert(feed.link.clone(), delay);
            self.feeds.insert(feed.link.clone(), feed);
            self.wakeup.notify();
        }
        !exists
    }

    async fn next(&mut self) -> Result<Feed, time::Error> {
        loop {
            if let Some(feed_id) = self.notifies.next().await {
                let feed = self.feeds.remove(feed_id?.get_ref()).unwrap();
                break Ok(feed);
            } else {
                self.wakeup.notified().await;
            }
        }
    }
}

struct Throttle {
    pieces: usize,
    counter: Arc<AtomicUsize>,
}

impl Throttle {
    fn new(pieces: usize) -> Self {
        Throttle {
            pieces,
            counter: Arc::new(AtomicUsize::new(0)),
        }
    }

    fn acquire(&self) -> Opportunity {
        Opportunity {
            n: self.counter.fetch_add(1, Ordering::AcqRel) % self.pieces,
            counter: self.counter.clone(),
        }
    }
}

#[must_use = "Don't lose your opportunity"]
struct Opportunity {
    n: usize,
    counter: Arc<AtomicUsize>,
}

impl Opportunity {
    async fn wait(&self) {
        delay_for(Duration::from_secs(self.n as u64)).await
    }
}

impl Drop for Opportunity {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}
<|endoftext|>