<commit_before>use ::opcodes::{AddressingMode, OpCode};

use cpu::cpu_error::CpuError;
use cpu::flags::StatusFlags;
use cpu::memory_bus::MemoryBus;
use cpu::registers::Registers;
use cpu::stack::Stack;

const DEFAULT_CODE_SEGMENT_START_ADDRESS: u16 = 0xC000;  // Default to a 16KB ROM, leaving 32KB of main memory

const STACK_START: usize = 0x100;
const STACK_END: usize = 0x1FF;

pub enum Operand {
    Immediate(u8),
    Memory(u16),
    Implied,
}

/// A representation of a 6502 microprocessor
pub struct Cpu {
    pub memory: MemoryBus,
    pub registers: Registers,
    pub flags: StatusFlags,
    pub stack: Stack,
}

pub type CpuLoadResult = Result<(), CpuError>;
pub type CpuStepResult = Result<(), CpuError>;

impl Cpu {
    /// Returns a default instance of a Cpu
    pub fn new() -> Cpu {
        Cpu {
            memory: MemoryBus::new(),
            registers: Registers::new(),
            flags: Default::default(),
            stack: Stack::new(),
        }
    }

    /// Loads code into the Cpu main memory at an optional offset. If no
    /// offset is provided, the Cpu will, by default, load the code into
    /// main memory at 0xC000
    pub fn load<T>(&mut self, code: &[u8], addr: T) -> CpuLoadResult
        where T: Into<Option<u16>>
    {
        let addr = addr.into();
        let addr: u16 = if addr.is_some() {
            let addr = addr.unwrap();
            if addr as u32 + code.len() as u32 > u16::max_value() as u32 {
                return Err(CpuError::code_segment_out_of_range(addr));
            } else {
                addr
            }
        } else {
            DEFAULT_CODE_SEGMENT_START_ADDRESS
        };

        for x in 0..code.len() {
            self.memory.write_byte(addr + x as u16, code[x]);
        }

        // Set the Program Counter to point at the
        // start address of the code segment
        self.registers.PC = addr;

        Ok(())
    }

    /// Runs N instructions of code through the Cpu
    pub fn step_n(&mut self, n: u32) -> CpuStepResult {
        for _ in 0..n {
            if self.registers.PC <= (self.memory.len() - 1) as u16 {
                self.step()?;
            } else {
                break;
            }
        }

        Ok(())
    }

    /// Runs a single instruction of code through the Cpu
    pub fn step(&mut self) -> CpuStepResult {
        let byte = self.memory.read_byte(self.registers.PC);

        if let Some(opcode) = OpCode::from_raw_byte(byte) {
            let operand = self.get_operand_from_opcode(&opcode);

            match opcode.mnemonic {
                "ADC" => self.adc(),
                "AND" => self.and(&operand),
                "ASL" => self.asl(&operand),
                "BCC" => self.bcc(&operand),
                "BCS" => self.bcs(&operand),
                "BEQ" => self.beq(&operand),
                "BIT" => self.bit(&operand),
                "BMI" => self.bmi(&operand),
                "BNE" => self.bne(&operand),
                "BPL" => self.bpl(&operand),
                "BRK" => self.brk(),
                "BVC" => self.bvc(&operand),
                "BVS" => self.bvs(&operand),
                "CLC" => self.set_carry_flag(false),
                "CLD" => self.set_decimal_flag(false),
                "CLI" => self.set_interrupt_flag(false),
                "CLV" => self.set_overflow_flag(false),
                "CMP" => self.cmp(&operand),
                "LDA" => self.lda(&operand),
                "SED" => self.set_decimal_flag(true),
                "STA" => self.sta(&operand),
                _ => return Err(CpuError::unknown_opcode(self.registers.PC, opcode.code)),
            }

            self.registers.PC += opcode.length as u16;

            Ok(())
        } else {
            Err(CpuError::unknown_opcode(self.registers.PC, byte))
        }
    }

    fn get_operand_from_opcode(&self, opcode: &OpCode) -> Operand {
        use ::opcodes::AddressingMode::*;

        let operand_start = self.registers.PC + 1;

        match opcode.mode {
            Unknown => unreachable!(),
            Implied => Operand::Implied,
            Immediate => Operand::Immediate(self.read_byte(operand_start)),
            Relative => Operand::Immediate(self.read_byte(operand_start)),
            Accumulator => Operand::Implied,
            ZeroPage => Operand::Memory((self.read_byte(operand_start) as u16) & 0xFF),
            ZeroPageX => {
                Operand::Memory((self.registers.X as u16 + self.read_byte(operand_start) as u16) &
                                0xFF)
            }
            ZeroPageY => {
                Operand::Memory((self.registers.Y as u16 + self.read_byte(operand_start) as u16) &
                                0xFF)
            }
            Absolute => Operand::Memory(self.read_u16(operand_start)),
            AbsoluteX => Operand::Memory(self.registers.X as u16 + self.read_u16(operand_start)),
            AbsoluteY => Operand::Memory(self.registers.Y as u16 + self.read_u16(operand_start)),
            Indirect => Operand::Memory(self.read_u16(self.read_u16(operand_start))),
            IndirectX => {
                Operand::Memory(self.read_u16((self.registers.X as u16 +
                                               self.read_byte(self.registers.PC + 1) as u16) &
                                              0xFF))
            }
            IndirectY => {
                Operand::Memory(self.registers.Y as u16 +
                                self.read_u16(self.read_byte(self.registers.PC + 1) as u16))
            }
        }
    }

    fn unwrap_immediate(&self, operand: &Operand) -> u8 {
        match *operand {
            Operand::Immediate(byte) => byte,
            Operand::Memory(addr) => self.read_byte(addr),
            Operand::Implied => 0,
        }
    }

    fn unwrap_address(&self, operand: &Operand) -> u16 {
        match *operand {
            Operand::Immediate(byte) => byte as u16,
            Operand::Memory(addr) => addr,
            Operand::Implied => 0,
        }
    }

    // ## OpCode handlers ##

    fn adc(&mut self) {
        // This is implemented on the information provided here:
        // http://www.electrical4u.com/bcd-or-binary-coded-decimal-bcd-conversion-addition-subtraction/
        // and here:
        // http://www.6502.org/tutorials/decimal_mode.html,
        // and here:
        // http://www.atariarchives.org/2bml/chapter_10.php,
        // and also here:
        // http://stackoverflow.com/questions/29193303/6502-emulation-proper-way-to-implement-adc-and-sbc

        let carry = if self.flags.carry { 1 } else { 0 };
        let value = self.read_byte(self.registers.PC + 1) as u16;

        let value_signs = self.registers.A & 0x80 == 0x80 && value & 0x80 == 0x80;

        // Do normal binary arithmetic first
        let mut result = self.registers.A as u16 + value as u16 + carry as u16;

        // Handle packed binary coded decimal
        if self.flags.decimal {
            if (self.registers.A as u16 & 0x0F) + (value & 0x0F) + carry > 0x09 {
                result += 0x06;
            }

            if result > 0x99 {
                result += 0x60;
            }

            self.flags.carry = (result & 0x100) == 0x100;
        } else {
            self.flags.carry = result > 0xFF;
        }

        self.flags.zero = result as u8 & 0xFF == 0x00;
        self.flags.sign = result & 0x80 == 0x80;

        if self.flags.sign != value_signs {
            self.flags.overflow = true;
        }

        self.registers.A = result as u8 & 0xFF;
    }

    fn and(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);
        let result = self.registers.A & value;

        self.registers.A = result;

        self.flags.zero = result as u8 & 0xFF == 0;
        self.flags.sign = result & 0x80 == 0x80;
    }

    fn asl(&mut self, operand: &Operand) {
        let mut value = if let &Operand::Implied = operand {
            // Implied ASL uses the A register
            self.registers.A
        } else {
            self.unwrap_immediate(&operand)
        };

        // Test the seventh bit - if its set, shift it
        // into the carry flag
        self.flags.carry = (value & 0x80) == 0x80;

        // Shift the value left
        value = value << 0x01;
        self.flags.sign = value & 0x80 == 0x80;
        self.flags.zero = value as u8 & 0xFF == 0;

        if let &Operand::Implied = operand {
            self.registers.A = value;
        } else {
            let addr = self.unwrap_address(&operand);
            self.write_byte(addr, value);
        }
    }

    fn bcc(&mut self, operand: &Operand) {
        // Branch if the carry flag is not set
        if !self.flags.carry {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bcs(&mut self, operand: &Operand) {
        // Branch if the carry flag is set
        if self.flags.carry {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn beq(&mut self, operand: &Operand) {
        // Branch if the zero flag is set
        if self.flags.zero {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bit(&mut self, operand: &Operand) {
        let a = self.registers.A;
        let value = self.unwrap_immediate(&operand);
        let result = value & a;

        self.flags.zero = result == 0x00;
        self.flags.overflow = value & 0x40 == 0x40; // "The V flag and the N flag receive copies of the sixth and seventh bits of the tested number"
        self.flags.sign = value & 0x80 == 0x80;
    }

    fn bmi(&mut self, operand: &Operand) {
        // Branch if the sign flag is set
        if self.flags.sign {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bne(&mut self, operand: &Operand) {
        // Branch if the zero flag is not set
        if !self.flags.zero {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bpl(&mut self, operand: &Operand) {
        // Branch if the sign flag is not set
        if !self.flags.sign {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn brk(&mut self) {
        let mut mem = &mut self.memory[STACK_START..STACK_END];

        // Return address is BRK + 0x02, but we do + 0x01 here
        // because after the cpu step we add another 0x01
        self.stack.push_u16(mem, self.registers.PC + 0x01);
        self.stack.push(mem, self.flags.to_u8());

        self.flags.interrupt_disabled = true;
    }

    fn bvc(&mut self, operand: &Operand) {
        // Branch if the overflow flag is not set
        if !self.flags.overflow {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bvs(&mut self, operand: &Operand) {
        // Branch if the overflow flag is set
        if self.flags.overflow {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn set_carry_flag(&mut self, value: bool) {
        self.flags.carry = value;
    }

    fn set_decimal_flag(&mut self, value: bool) {
        self.flags.decimal = value;
    }

    fn set_interrupt_flag(&mut self, value: bool) {
        self.flags.interrupt_disabled = value;
    }

    fn set_overflow_flag(&mut self, value: bool) {
        self.flags.overflow = value;
    }

    fn cmp(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);
        let result: i16 = self.registers.A as i16 - value as i16;

        self.flags.carry = (result as u16) < 0x100;
        println!("A: {}, value: {}, result: {:08X}, result < 0x100: {}",
                 self.registers.A,
                 value,
                 result,
                 (result as u16) < 0x100);
        println!("Carry: {}", self.flags.carry);
        self.flags.zero = result & 0xFF == 0x00;
        self.flags.sign = result & 0x80 == 0x80;
    }

    fn lda(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);

        self.registers.A = value;
        self.flags.sign = value & 0x80 == 0x80;
        self.flags.zero = value & 0xFF == 0x00;
    }

    fn sta(&mut self, operand: &Operand) {
        let addr = self.unwrap_address(&operand);
        let value = self.registers.A;

        self.write_byte(addr, value);
    }

    fn relative_jump(&mut self, offset: u8) {
        // If the sign bit is there, negate the PC by the difference
        // between 256 and the offset
        if offset & 0x80 == 0x80 {
            self.registers.PC -= 0x100 - offset as u16;
        } else {
            self.registers.PC += offset as u16;
        }
    }

    /// Convenience wrapper for accessing a byte
    /// in memory
    fn read_byte(&self, addr: u16) -> u8 {
        self.memory.read_byte(addr)
    }

    /// Convenience wrapper for writing a byte
    /// to memory
    fn write_byte(&mut self, addr: u16, byte: u8) {
        self.memory.write_byte(addr, byte);
    }

    /// Convenience wrapper for accessing a word
    /// in memory
    fn read_u16(&self, addr: u16) -> u16 {
        self.memory.read_u16(addr)
    }
}
<commit_msg>Remove debug statements<commit_after>use ::opcodes::{AddressingMode, OpCode};

use cpu::cpu_error::CpuError;
use cpu::flags::StatusFlags;
use cpu::memory_bus::MemoryBus;
use cpu::registers::Registers;
use cpu::stack::Stack;

const DEFAULT_CODE_SEGMENT_START_ADDRESS: u16 = 0xC000;  // Default to a 16KB ROM, leaving 32KB of main memory

const STACK_START: usize = 0x100;
const STACK_END: usize = 0x1FF;

pub enum Operand {
    Immediate(u8),
    Memory(u16),
    Implied,
}

/// A representation of a 6502 microprocessor
pub struct Cpu {
    pub memory: MemoryBus,
    pub registers: Registers,
    pub flags: StatusFlags,
    pub stack: Stack,
}

pub type CpuLoadResult = Result<(), CpuError>;
pub type CpuStepResult = Result<(), CpuError>;

impl Cpu {
    /// Returns a default instance of a Cpu
    pub fn new() -> Cpu {
        Cpu {
            memory: MemoryBus::new(),
            registers: Registers::new(),
            flags: Default::default(),
            stack: Stack::new(),
        }
    }

    /// Loads code into the Cpu main memory at an optional offset. If no
    /// offset is provided, the Cpu will, by default, load the code into
    /// main memory at 0xC000
    pub fn load<T>(&mut self, code: &[u8], addr: T) -> CpuLoadResult
        where T: Into<Option<u16>>
    {
        let addr = addr.into();
        let addr: u16 = if addr.is_some() {
            let addr = addr.unwrap();
            if addr as u32 + code.len() as u32 > u16::max_value() as u32 {
                return Err(CpuError::code_segment_out_of_range(addr));
            } else {
                addr
            }
        } else {
            DEFAULT_CODE_SEGMENT_START_ADDRESS
        };

        for x in 0..code.len() {
            self.memory.write_byte(addr + x as u16, code[x]);
        }

        // Set the Program Counter to point at the
        // start address of the code segment
        self.registers.PC = addr;

        Ok(())
    }

    /// Runs N instructions of code through the Cpu
    pub fn step_n(&mut self, n: u32) -> CpuStepResult {
        for _ in 0..n {
            if self.registers.PC <= (self.memory.len() - 1) as u16 {
                self.step()?;
            } else {
                break;
            }
        }

        Ok(())
    }

    /// Runs a single instruction of code through the Cpu
    pub fn step(&mut self) -> CpuStepResult {
        let byte = self.memory.read_byte(self.registers.PC);

        if let Some(opcode) = OpCode::from_raw_byte(byte) {
            let operand = self.get_operand_from_opcode(&opcode);

            match opcode.mnemonic {
                "ADC" => self.adc(),
                "AND" => self.and(&operand),
                "ASL" => self.asl(&operand),
                "BCC" => self.bcc(&operand),
                "BCS" => self.bcs(&operand),
                "BEQ" => self.beq(&operand),
                "BIT" => self.bit(&operand),
                "BMI" => self.bmi(&operand),
                "BNE" => self.bne(&operand),
                "BPL" => self.bpl(&operand),
                "BRK" => self.brk(),
                "BVC" => self.bvc(&operand),
                "BVS" => self.bvs(&operand),
                "CLC" => self.set_carry_flag(false),
                "CLD" => self.set_decimal_flag(false),
                "CLI" => self.set_interrupt_flag(false),
                "CLV" => self.set_overflow_flag(false),
                "CMP" => self.cmp(&operand),
                "LDA" => self.lda(&operand),
                "SED" => self.set_decimal_flag(true),
                "STA" => self.sta(&operand),
                _ => return Err(CpuError::unknown_opcode(self.registers.PC, opcode.code)),
            }

            self.registers.PC += opcode.length as u16;

            Ok(())
        } else {
            Err(CpuError::unknown_opcode(self.registers.PC, byte))
        }
    }

    fn get_operand_from_opcode(&self, opcode: &OpCode) -> Operand {
        use ::opcodes::AddressingMode::*;

        let operand_start = self.registers.PC + 1;

        match opcode.mode {
            Unknown => unreachable!(),
            Implied => Operand::Implied,
            Immediate => Operand::Immediate(self.read_byte(operand_start)),
            Relative => Operand::Immediate(self.read_byte(operand_start)),
            Accumulator => Operand::Implied,
            ZeroPage => Operand::Memory((self.read_byte(operand_start) as u16) & 0xFF),
            ZeroPageX => {
                Operand::Memory((self.registers.X as u16 + self.read_byte(operand_start) as u16) &
                                0xFF)
            }
            ZeroPageY => {
                Operand::Memory((self.registers.Y as u16 + self.read_byte(operand_start) as u16) &
                                0xFF)
            }
            Absolute => Operand::Memory(self.read_u16(operand_start)),
            AbsoluteX => Operand::Memory(self.registers.X as u16 + self.read_u16(operand_start)),
            AbsoluteY => Operand::Memory(self.registers.Y as u16 + self.read_u16(operand_start)),
            Indirect => Operand::Memory(self.read_u16(self.read_u16(operand_start))),
            IndirectX => {
                Operand::Memory(self.read_u16((self.registers.X as u16 +
                                               self.read_byte(self.registers.PC + 1) as u16) &
                                              0xFF))
            }
            IndirectY => {
                Operand::Memory(self.registers.Y as u16 +
                                self.read_u16(self.read_byte(self.registers.PC + 1) as u16))
            }
        }
    }

    fn unwrap_immediate(&self, operand: &Operand) -> u8 {
        match *operand {
            Operand::Immediate(byte) => byte,
            Operand::Memory(addr) => self.read_byte(addr),
            Operand::Implied => 0,
        }
    }

    fn unwrap_address(&self, operand: &Operand) -> u16 {
        match *operand {
            Operand::Immediate(byte) => byte as u16,
            Operand::Memory(addr) => addr,
            Operand::Implied => 0,
        }
    }

    // ## OpCode handlers ##

    fn adc(&mut self) {
        // This is implemented on the information provided here:
        // http://www.electrical4u.com/bcd-or-binary-coded-decimal-bcd-conversion-addition-subtraction/
        // and here:
        // http://www.6502.org/tutorials/decimal_mode.html,
        // and here:
        // http://www.atariarchives.org/2bml/chapter_10.php,
        // and also here:
        // http://stackoverflow.com/questions/29193303/6502-emulation-proper-way-to-implement-adc-and-sbc

        let carry = if self.flags.carry { 1 } else { 0 };
        let value = self.read_byte(self.registers.PC + 1) as u16;

        let value_signs = self.registers.A & 0x80 == 0x80 && value & 0x80 == 0x80;

        // Do normal binary arithmetic first
        let mut result = self.registers.A as u16 + value as u16 + carry as u16;

        // Handle packed binary coded decimal
        if self.flags.decimal {
            if (self.registers.A as u16 & 0x0F) + (value & 0x0F) + carry > 0x09 {
                result += 0x06;
            }

            if result > 0x99 {
                result += 0x60;
            }

            self.flags.carry = (result & 0x100) == 0x100;
        } else {
            self.flags.carry = result > 0xFF;
        }

        self.flags.zero = result as u8 & 0xFF == 0x00;
        self.flags.sign = result & 0x80 == 0x80;

        if self.flags.sign != value_signs {
            self.flags.overflow = true;
        }

        self.registers.A = result as u8 & 0xFF;
    }

    fn and(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);
        let result = self.registers.A & value;

        self.registers.A = result;

        self.flags.zero = result as u8 & 0xFF == 0;
        self.flags.sign = result & 0x80 == 0x80;
    }

    fn asl(&mut self, operand: &Operand) {
        let mut value = if let &Operand::Implied = operand {
            // Implied ASL uses the A register
            self.registers.A
        } else {
            self.unwrap_immediate(&operand)
        };

        // Test the seventh bit - if its set, shift it
        // into the carry flag
        self.flags.carry = (value & 0x80) == 0x80;

        // Shift the value left
        value = value << 0x01;
        self.flags.sign = value & 0x80 == 0x80;
        self.flags.zero = value as u8 & 0xFF == 0;

        if let &Operand::Implied = operand {
            self.registers.A = value;
        } else {
            let addr = self.unwrap_address(&operand);
            self.write_byte(addr, value);
        }
    }

    fn bcc(&mut self, operand: &Operand) {
        // Branch if the carry flag is not set
        if !self.flags.carry {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bcs(&mut self, operand: &Operand) {
        // Branch if the carry flag is set
        if self.flags.carry {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn beq(&mut self, operand: &Operand) {
        // Branch if the zero flag is set
        if self.flags.zero {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bit(&mut self, operand: &Operand) {
        let a = self.registers.A;
        let value = self.unwrap_immediate(&operand);
        let result = value & a;

        self.flags.zero = result == 0x00;
        self.flags.overflow = value & 0x40 == 0x40; // "The V flag and the N flag receive copies of the sixth and seventh bits of the tested number"
        self.flags.sign = value & 0x80 == 0x80;
    }

    fn bmi(&mut self, operand: &Operand) {
        // Branch if the sign flag is set
        if self.flags.sign {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bne(&mut self, operand: &Operand) {
        // Branch if the zero flag is not set
        if !self.flags.zero {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bpl(&mut self, operand: &Operand) {
        // Branch if the sign flag is not set
        if !self.flags.sign {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn brk(&mut self) {
        let mut mem = &mut self.memory[STACK_START..STACK_END];

        // Return address is BRK + 0x02, but we do + 0x01 here
        // because after the cpu step we add another 0x01
        self.stack.push_u16(mem, self.registers.PC + 0x01);
        self.stack.push(mem, self.flags.to_u8());

        self.flags.interrupt_disabled = true;
    }

    fn bvc(&mut self, operand: &Operand) {
        // Branch if the overflow flag is not set
        if !self.flags.overflow {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn bvs(&mut self, operand: &Operand) {
        // Branch if the overflow flag is set
        if self.flags.overflow {
            let offset = self.unwrap_immediate(&operand);
            self.relative_jump(offset);
        }
    }

    fn set_carry_flag(&mut self, value: bool) {
        self.flags.carry = value;
    }

    fn set_decimal_flag(&mut self, value: bool) {
        self.flags.decimal = value;
    }

    fn set_interrupt_flag(&mut self, value: bool) {
        self.flags.interrupt_disabled = value;
    }

    fn set_overflow_flag(&mut self, value: bool) {
        self.flags.overflow = value;
    }

    fn cmp(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);
        let result: i16 = self.registers.A as i16 - value as i16;

        self.flags.carry = (result as u16) < 0x100;
        self.flags.zero = result & 0xFF == 0x00;
        self.flags.sign = result & 0x80 == 0x80;
    }

    fn lda(&mut self, operand: &Operand) {
        let value = self.unwrap_immediate(&operand);

        self.registers.A = value;
        self.flags.sign = value & 0x80 == 0x80;
        self.flags.zero = value & 0xFF == 0x00;
    }

    fn sta(&mut self, operand: &Operand) {
        let addr = self.unwrap_address(&operand);
        let value = self.registers.A;

        self.write_byte(addr, value);
    }

    fn relative_jump(&mut self, offset: u8) {
        // If the sign bit is there, negate the PC by the difference
        // between 256 and the offset
        if offset & 0x80 == 0x80 {
            self.registers.PC -= 0x100 - offset as u16;
        } else {
            self.registers.PC += offset as u16;
        }
    }

    /// Convenience wrapper for accessing a byte
    /// in memory
    fn read_byte(&self, addr: u16) -> u8 {
        self.memory.read_byte(addr)
    }

    /// Convenience wrapper for writing a byte
    /// to memory
    fn write_byte(&mut self, addr: u16, byte: u8) {
        self.memory.write_byte(addr, byte);
    }

    /// Convenience wrapper for accessing a word
    /// in memory
    fn read_u16(&self, addr: u16) -> u16 {
        self.memory.read_u16(addr)
    }
}
<|endoftext|>


<commit_before>use heck::KebabCase;
use quote::{ToTokens, Tokens, Ident};

use ast::*;

impl ToTokens for Resource {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let resource_ty = Ident::new(&self.header.ty[..]);
        let resource_name = self.header.ty.to_kebab_case();
        let endpoint = self.header.endpoint.as_ref().unwrap_or(&resource_name);
        let rel_links = self.rel_links();
        let relationships = self.relationships();
        let rel_ids = self.rel_ids();
        let rel_ids_ty = Ident::new(&rel_ids.ty[..]);
        tokens.append(quote! {
            impl ::cargonauts::routing::ResourceEndpoint for #resource_ty {
                const ENDPOINT: &'static str = #endpoint;
                const RESOURCE: &'static str = #resource_name;
                const REL_LINKS: &'static [::cargonauts::routing::RelationshipLink] = &#rel_links;
                type RelIds = #rel_ids_ty;
            }

            #(#relationships)*

            #rel_ids;
        })
    }
}

impl Resource {
    fn rel_links(&self) -> Vec<RelLink> {
        self.members.iter().filter_map(|m| m.as_relation()).map(|rel| RelLink {
            resource: &self.header.ty,
            relation: &rel.rel,
        }).collect()
    }

    fn relationships(&self) -> Vec<Relationship> {
        self.members.iter().filter_map(|m| m.as_relation()).map(|rel| Relationship {
            resource: &self.header.ty,
            relation: &rel.rel,
            endpoint: rel.endpoint.as_ref(),
            kind: rel.kind,
        }).collect()
    }

    fn rel_ids(&self) -> RelIds {
        RelIds {
            ty: self.header.ty.clone() + "RelIds",
            resource: &self.header.ty,
            rels: self.members.iter().filter_map(|m| m.as_relation()).map(|rel| (&rel.rel[..], rel.kind)).collect(),
        }
    }
}

struct RelLink<'a> {
    resource: &'a str,
    relation: &'a str,
}

impl<'a> ToTokens for RelLink<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let relation_ty = Ident::new(self.relation);
        let resource_ty = Ident::new(self.resource);
        tokens.append(quote! {
            ::cargonauts::routing::RelationshipLink {
                endpoint: <#resource_ty as ::cargonauts::routing::RelationEndpoint<#relation_ty>>::REL_ENDPOINT,
                relation: <#resource_ty as ::cargonauts::routing::RelationEndpoint<#relation_ty>>::RELATION,
            }
        })
    }
}

struct Relationship<'a> {
    resource: &'a str,
    relation: &'a str,
    endpoint: Option<&'a String>,
    kind: RelationKind,
}

impl<'a> ToTokens for Relationship<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let relation_ty = Ident::new(self.relation);
        let resource_ty = Ident::new(self.resource);
        let relation = self.relation.to_kebab_case();
        let endpoint = self.endpoint.unwrap_or(&relation);
        tokens.append(quote! {
            impl ::cargonauts::routing::RelationEndpoint<#relation_ty> for #resource_ty {
                const REL_ENDPOINT: &'static str = #endpoint;
                const RELATION: &'static str = #relation;
            }
        });

        tokens.append(match self.kind {
            RelationKind::Single => quote! {
                impl ::cargonauts::routing::HasOneEndpoint<#relation_ty> for #resource_ty {
                }
            },
            RelationKind::Many => quote! {
                impl ::cargonauts::routing::HasManyEndpoint<#relation_ty> for #resource_ty {
                }
            },
        });
    }
}

struct RelIds<'a> {
    ty: String,
    resource: &'a str,
    rels: Vec<(&'a str, RelationKind)>,
}

impl<'a> ToTokens for RelIds<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let resource = Ident::new(self.resource);
        let ty = Ident::new(&self.ty[..]);
        let rels = self.rels.iter().map(|&(rel, kind)| {
            let rel = Ident::new(rel);
            match kind {
                RelationKind::Single    => quote!(#rel: Option<String>,),
                RelationKind::Many      => quote!(#rel: Vec<String>,),
            }
        });

        let single_try_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => {
                    self.#rel = Some(id);
                    true
                }))
            } else { None }
        });

        let many_try_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => {
                    self.#rel = ids;
                    true
                }))
            } else { None }
        });

        let single_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel = Some(id),))
            } else { None }
        });

        let single_get_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel.as_ref().map(|s| &s[..]),))
            } else { None }
        });

        let many_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Many = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel = ids,))
            } else { None }
        });

        let many_get_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Many = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => &self.#rel[..],))
            } else { None }
        });

        tokens.append(quote! {
            #[allow(non_snake_case)]
            #[derive(Default)]
            pub struct #ty {
                #(#rels)*
            }

            #[allow(unused_variables)]
            impl ::cargonauts::routing::RelIds<#resource> for #ty {
                fn try_set_rel_id(&mut self, rel: &str, id: String) -> bool {
                    match rel {
                        #(#single_try_set_arms)*
                        _ => false,
                    }
                }

                fn try_set_rel_ids(&mut self, rel: &str, ids: Vec<String>) -> bool {
                    match rel {
                        #(#many_try_set_arms)*
                        _ => false,
                    }
                }

                fn set_rel_id<R>(&mut self, id: String)
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasOneEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#single_set_arms)*
                        _ => (),
                    }
                }

                fn rel_id<R>(&self) -> Option<&str>
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasOneEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#single_get_arms)*
                        _ => None,
                    }
                }

                fn set_rel_ids<R>(&mut self, id: Vec<String>)
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasManyEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#many_set_arms)*
                        _ => (),
                    }
                }

                fn rel_ids<R>(&self) -> &[String]
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasManyEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#many_get_arms)*
                        _ => &[],
                    }
                }
            }
        })
    }
}
<commit_msg>Fix bug in relid generation.<commit_after>use heck::KebabCase;
use quote::{ToTokens, Tokens, Ident};

use ast::*;

impl ToTokens for Resource {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let resource_ty = Ident::new(&self.header.ty[..]);
        let resource_name = self.header.ty.to_kebab_case();
        let endpoint = self.header.endpoint.as_ref().unwrap_or(&resource_name);
        let rel_links = self.rel_links();
        let relationships = self.relationships();
        let rel_ids = self.rel_ids();
        let rel_ids_ty = Ident::new(&rel_ids.ty[..]);
        tokens.append(quote! {
            impl ::cargonauts::routing::ResourceEndpoint for #resource_ty {
                const ENDPOINT: &'static str = #endpoint;
                const RESOURCE: &'static str = #resource_name;
                const REL_LINKS: &'static [::cargonauts::routing::RelationshipLink] = &#rel_links;
                type RelIds = #rel_ids_ty;
            }

            #(#relationships)*

            #rel_ids;
        })
    }
}

impl Resource {
    fn rel_links(&self) -> Vec<RelLink> {
        self.members.iter().filter_map(|m| m.as_relation()).map(|rel| RelLink {
            resource: &self.header.ty,
            relation: &rel.rel,
        }).collect()
    }

    fn relationships(&self) -> Vec<Relationship> {
        self.members.iter().filter_map(|m| m.as_relation()).map(|rel| Relationship {
            resource: &self.header.ty,
            relation: &rel.rel,
            endpoint: rel.endpoint.as_ref(),
            kind: rel.kind,
        }).collect()
    }

    fn rel_ids(&self) -> RelIds {
        RelIds {
            ty: self.header.ty.clone() + "RelIds",
            resource: &self.header.ty,
            rels: self.members.iter().filter_map(|m| m.as_relation()).map(|rel| (&rel.rel[..], rel.kind)).collect(),
        }
    }
}

struct RelLink<'a> {
    resource: &'a str,
    relation: &'a str,
}

impl<'a> ToTokens for RelLink<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let relation_ty = Ident::new(self.relation);
        let resource_ty = Ident::new(self.resource);
        tokens.append(quote! {
            ::cargonauts::routing::RelationshipLink {
                endpoint: <#resource_ty as ::cargonauts::routing::RelationEndpoint<#relation_ty>>::REL_ENDPOINT,
                relation: <#resource_ty as ::cargonauts::routing::RelationEndpoint<#relation_ty>>::RELATION,
            }
        })
    }
}

struct Relationship<'a> {
    resource: &'a str,
    relation: &'a str,
    endpoint: Option<&'a String>,
    kind: RelationKind,
}

impl<'a> ToTokens for Relationship<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let relation_ty = Ident::new(self.relation);
        let resource_ty = Ident::new(self.resource);
        let relation = self.relation.to_kebab_case();
        let endpoint = self.endpoint.unwrap_or(&relation);
        tokens.append(quote! {
            impl ::cargonauts::routing::RelationEndpoint<#relation_ty> for #resource_ty {
                const REL_ENDPOINT: &'static str = #endpoint;
                const RELATION: &'static str = #relation;
            }
        });

        tokens.append(match self.kind {
            RelationKind::Single => quote! {
                impl ::cargonauts::routing::HasOneEndpoint<#relation_ty> for #resource_ty {
                }
            },
            RelationKind::Many => quote! {
                impl ::cargonauts::routing::HasManyEndpoint<#relation_ty> for #resource_ty {
                }
            },
        });
    }
}

struct RelIds<'a> {
    ty: String,
    resource: &'a str,
    rels: Vec<(&'a str, RelationKind)>,
}

impl<'a> ToTokens for RelIds<'a> {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let resource = Ident::new(self.resource);
        let ty = Ident::new(&self.ty[..]);
        let rels = self.rels.iter().map(|&(rel, kind)| {
            let rel = Ident::new(rel);
            match kind {
                RelationKind::Single    => quote!(#rel: Option<String>,),
                RelationKind::Many      => quote!(#rel: Vec<String>,),
            }
        });

        let single_try_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => {
                    self.#rel = Some(id);
                    true
                }))
            } else { None }
        });

        let many_try_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Many = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => {
                    self.#rel = ids;
                    true
                }))
            } else { None }
        });

        let single_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel = Some(id),))
            } else { None }
        });

        let single_get_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Single = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel.as_ref().map(|s| &s[..]),))
            } else { None }
        });

        let many_set_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Many = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => self.#rel = ids,))
            } else { None }
        });

        let many_get_arms = self.rels.iter().filter_map(|&(rel, kind)| {
            if let RelationKind::Many = kind {
                let rel = Ident::new(rel);
                Some(quote!(<#resource as ::cargonauts::routing::RelationEndpoint<#rel>>::RELATION => &self.#rel[..],))
            } else { None }
        });

        tokens.append(quote! {
            #[allow(non_snake_case)]
            #[derive(Default)]
            pub struct #ty {
                #(#rels)*
            }

            #[allow(unused_variables)]
            impl ::cargonauts::routing::RelIds<#resource> for #ty {
                fn try_set_rel_id(&mut self, rel: &str, id: String) -> bool {
                    match rel {
                        #(#single_try_set_arms)*
                        _ => false,
                    }
                }

                fn try_set_rel_ids(&mut self, rel: &str, ids: Vec<String>) -> bool {
                    match rel {
                        #(#many_try_set_arms)*
                        _ => false,
                    }
                }

                fn set_rel_id<R>(&mut self, id: String)
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasOneEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#single_set_arms)*
                        _ => (),
                    }
                }

                fn rel_id<R>(&self) -> Option<&str>
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasOneEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#single_get_arms)*
                        _ => None,
                    }
                }

                fn set_rel_ids<R>(&mut self, id: Vec<String>)
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasManyEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#many_set_arms)*
                        _ => (),
                    }
                }

                fn rel_ids<R>(&self) -> &[String]
                where
                    R: ::cargonauts::api::Relationship,
                    #resource: ::cargonauts::routing::HasManyEndpoint<R>,
                    R::Related: ::cargonauts::routing::ResourceEndpoint,
                {
                    match <#resource as ::cargonauts::routing::RelationEndpoint<R>>::RELATION {
                        #(#many_get_arms)*
                        _ => &[],
                    }
                }
            }
        })
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "cargo-clippy", allow(redundant_field_names))]

use std::io::{self, Write};
use std::cell::RefCell;
use std::fmt::Write as FmtWrite;

use time::{self, Duration};
use bytes::{BytesMut, BufMut};
use futures::{Async, Poll};
use tokio_io::AsyncWrite;
use http::{Version, HttpTryFrom};
use http::header::{HeaderValue, DATE,
                   CONNECTION, CONTENT_ENCODING, CONTENT_LENGTH, TRANSFER_ENCODING};
use flate2::Compression;
use flate2::write::{GzEncoder, DeflateEncoder};
use brotli2::write::BrotliEncoder;

use body::{Body, Binary};
use headers::ContentEncoding;
use server::WriterState;
use server::shared::SharedBytes;
use server::encoding::{ContentEncoder, TransferEncoding};

use client::ClientRequest;

const AVERAGE_HEADER_SIZE: usize = 30;

bitflags! {
    struct Flags: u8 {
        const STARTED = 0b0000_0001;
        const UPGRADE = 0b0000_0010;
        const KEEPALIVE = 0b0000_0100;
        const DISCONNECTED = 0b0000_1000;
    }
}

pub(crate) struct HttpClientWriter {
    flags: Flags,
    written: u64,
    headers_size: u32,
    buffer: SharedBytes,
    buffer_capacity: usize,
    encoder: ContentEncoder,
}

impl HttpClientWriter {

    pub fn new(buffer: SharedBytes) -> HttpClientWriter {
        let encoder = ContentEncoder::Identity(TransferEncoding::eof(buffer.clone()));
        HttpClientWriter {
            flags: Flags::empty(),
            written: 0,
            headers_size: 0,
            buffer_capacity: 0,
            buffer,
            encoder,
        }
    }

    pub fn disconnected(&mut self) {
        self.buffer.take();
    }

    // pub fn keepalive(&self) -> bool {
    //    self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
    // }

    fn write_to_stream<T: AsyncWrite>(&mut self, stream: &mut T) -> io::Result<WriterState> {
        while !self.buffer.is_empty() {
            match stream.write(self.buffer.as_ref()) {
                Ok(0) => {
                    self.disconnected();
                    return Ok(WriterState::Done);
                },
                Ok(n) => {
                    let _ = self.buffer.split_to(n);
                },
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    if self.buffer.len() > self.buffer_capacity {
                        return Ok(WriterState::Pause)
                    } else {
                        return Ok(WriterState::Done)
                    }
                }
                Err(err) => return Err(err),
            }
        }
        Ok(WriterState::Done)
    }
}

impl HttpClientWriter {

    pub fn start(&mut self, msg: &mut ClientRequest) -> io::Result<()> {
        // prepare task
        self.flags.insert(Flags::STARTED);
        self.encoder = content_encoder(self.buffer.clone(), msg);

        // render message
        {
            let mut buffer = self.buffer.get_mut();
            if let Body::Binary(ref bytes) = *msg.body() {
                buffer.reserve(256 + msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
            } else {
                buffer.reserve(256 + msg.headers().len() * AVERAGE_HEADER_SIZE);
            }

            if msg.upgrade() {
                self.flags.insert(Flags::UPGRADE);
            }

            let path = msg.uri().path_and_query().map(|u| u.as_str()).unwrap_or("");
            // status line
            let _ = write!(buffer, "{} {} {:?}\r\n",
                           msg.method(), path, msg.version());

            // write headers
            for (key, value) in msg.headers() {
                let v = value.as_ref();
                let k = key.as_str().as_bytes();
                buffer.reserve(k.len() + v.len() + 4);
                buffer.put_slice(k);
                buffer.put_slice(b": ");
                buffer.put_slice(v);
                buffer.put_slice(b"\r\n");
            }

            // set date header
            if !msg.headers().contains_key(DATE) {
                buffer.extend_from_slice(b"date: ");
                set_date(&mut buffer);
                buffer.extend_from_slice(b"\r\n\r\n");
            } else {
                buffer.extend_from_slice(b"\r\n");
            }
            self.headers_size = buffer.len() as u32;

            if msg.body().is_binary() {
                if let Body::Binary(bytes) = msg.replace_body(Body::Empty) {
                    self.written += bytes.len() as u64;
                    self.encoder.write(bytes)?;
                }
            } else {
                self.buffer_capacity = msg.write_buffer_capacity();
            }
        }
        Ok(())
    }

    pub fn write(&mut self, payload: Binary) -> io::Result<WriterState> {
        self.written += payload.len() as u64;
        if !self.flags.contains(Flags::DISCONNECTED) {
            if self.flags.contains(Flags::UPGRADE) {
                self.buffer.extend(payload);
            } else {
                self.encoder.write(payload)?;
            }
        }

        if self.buffer.len() > self.buffer_capacity {
            Ok(WriterState::Pause)
        } else {
            Ok(WriterState::Done)
        }
    }

    pub fn write_eof(&mut self) -> io::Result<()> {
        self.encoder.write_eof()?;

        if self.encoder.is_eof() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other,
                               "Last payload item, but eof is not reached"))
        }
    }

    #[inline]
    pub fn poll_completed<T: AsyncWrite>(&mut self, stream: &mut T, shutdown: bool)
                                         -> Poll<(), io::Error>
    {
        match self.write_to_stream(stream) {
            Ok(WriterState::Done) => {
                if shutdown {
                    stream.shutdown()
                } else {
                    Ok(Async::Ready(()))
                }
            },
            Ok(WriterState::Pause) => Ok(Async::NotReady),
            Err(err) => Err(err)
        }
    }
}

fn content_encoder(buf: SharedBytes, req: &mut ClientRequest) -> ContentEncoder {
    let version = req.version();
    let mut body = req.replace_body(Body::Empty);
    let mut encoding = req.content_encoding();

    let transfer = match body {
        Body::Empty => {
            req.headers_mut().remove(CONTENT_LENGTH);
            TransferEncoding::length(0, buf)
        },
        Body::Binary(ref mut bytes) => {
            if encoding.is_compression() {
                let tmp = SharedBytes::default();
                let transfer = TransferEncoding::eof(tmp.clone());
                let mut enc = match encoding {
                    ContentEncoding::Deflate => ContentEncoder::Deflate(
                        DeflateEncoder::new(transfer, Compression::default())),
                    ContentEncoding::Gzip => ContentEncoder::Gzip(
                        GzEncoder::new(transfer, Compression::default())),
                    ContentEncoding::Br => ContentEncoder::Br(
                        BrotliEncoder::new(transfer, 5)),
                    ContentEncoding::Identity => ContentEncoder::Identity(transfer),
                    ContentEncoding::Auto => unreachable!()
                };
                // TODO return error!
                let _ = enc.write(bytes.clone());
                let _ = enc.write_eof();
                *bytes = Binary::from(tmp.take());

                req.headers_mut().insert(
                    CONTENT_ENCODING, HeaderValue::from_static(encoding.as_str()));
                encoding = ContentEncoding::Identity;
            }
            let mut b = BytesMut::new();
            let _ = write!(b, "{}", bytes.len());
            req.headers_mut().insert(
                CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
            TransferEncoding::eof(buf)
        },
        Body::Streaming(_) | Body::Actor(_) => {
            if req.upgrade() {
                if version == Version::HTTP_2 {
                    error!("Connection upgrade is forbidden for HTTP/2");
                } else {
                    req.headers_mut().insert(CONNECTION, HeaderValue::from_static("upgrade"));
                }
                if encoding != ContentEncoding::Identity {
                    encoding = ContentEncoding::Identity;
                    req.headers_mut().remove(CONTENT_ENCODING);
                }
                TransferEncoding::eof(buf)
            } else {
                streaming_encoding(buf, version, req)
            }
        }
    };

    if encoding.is_compression() {
        req.headers_mut().insert(
            CONTENT_ENCODING, HeaderValue::from_static(encoding.as_str()));
    }

    req.replace_body(body);
    match encoding {
        ContentEncoding::Deflate => ContentEncoder::Deflate(
            DeflateEncoder::new(transfer, Compression::default())),
        ContentEncoding::Gzip => ContentEncoder::Gzip(
            GzEncoder::new(transfer, Compression::default())),
        ContentEncoding::Br => ContentEncoder::Br(
            BrotliEncoder::new(transfer, 5)),
        ContentEncoding::Identity | ContentEncoding::Auto => ContentEncoder::Identity(transfer),
    }
}

fn streaming_encoding(buf: SharedBytes, version: Version, req: &mut ClientRequest)
                      -> TransferEncoding {
    if req.chunked() {
        // Enable transfer encoding
        req.headers_mut().remove(CONTENT_LENGTH);
        if version == Version::HTTP_2 {
            req.headers_mut().remove(TRANSFER_ENCODING);
            TransferEncoding::eof(buf)
        } else {
            req.headers_mut().insert(
                TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
            TransferEncoding::chunked(buf)
        }
    } else {
        // if Content-Length is specified, then use it as length hint
        let (len, chunked) =
            if let Some(len) = req.headers().get(CONTENT_LENGTH) {
                // Content-Length
                if let Ok(s) = len.to_str() {
                    if let Ok(len) = s.parse::<u64>() {
                        (Some(len), false)
                    } else {
                        error!("illegal Content-Length: {:?}", len);
                        (None, false)
                    }
                } else {
                    error!("illegal Content-Length: {:?}", len);
                    (None, false)
                }
            } else {
                (None, true)
            };

        if !chunked {
            if let Some(len) = len {
                TransferEncoding::length(len, buf)
            } else {
                TransferEncoding::eof(buf)
            }
        } else {
            // Enable transfer encoding
            match version {
                Version::HTTP_11 => {
                    req.headers_mut().insert(
                        TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
                    TransferEncoding::chunked(buf)
                },
                _ => {
                    req.headers_mut().remove(TRANSFER_ENCODING);
                    TransferEncoding::eof(buf)
                }
            }
        }
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
pub const DATE_VALUE_LENGTH: usize = 29;

fn set_date(dst: &mut BytesMut) {
    CACHED.with(|cache| {
        let mut cache = cache.borrow_mut();
        let now = time::get_time();
        if now > cache.next_update {
            cache.update(now);
        }
        dst.extend_from_slice(cache.buffer());
    })
}

struct CachedDate {
    bytes: [u8; DATE_VALUE_LENGTH],
    next_update: time::Timespec,
}

thread_local!(static CACHED: RefCell<CachedDate> = RefCell::new(CachedDate {
    bytes: [0; DATE_VALUE_LENGTH],
    next_update: time::Timespec::new(0, 0),
}));

impl CachedDate {
    fn buffer(&self) -> &[u8] {
        &self.bytes[..]
    }

    fn update(&mut self, now: time::Timespec) {
        write!(&mut self.bytes[..], "{}", time::at_utc(now).rfc822()).unwrap();
        self.next_update = now + Duration::seconds(1);
        self.next_update.nsec = 0;
    }
}
<commit_msg>Move path call inline into write<commit_after>#![cfg_attr(feature = "cargo-clippy", allow(redundant_field_names))]

use std::io::{self, Write};
use std::cell::RefCell;
use std::fmt::Write as FmtWrite;

use time::{self, Duration};
use bytes::{BytesMut, BufMut};
use futures::{Async, Poll};
use tokio_io::AsyncWrite;
use http::{Version, HttpTryFrom};
use http::header::{HeaderValue, DATE,
                   CONNECTION, CONTENT_ENCODING, CONTENT_LENGTH, TRANSFER_ENCODING};
use flate2::Compression;
use flate2::write::{GzEncoder, DeflateEncoder};
use brotli2::write::BrotliEncoder;

use body::{Body, Binary};
use headers::ContentEncoding;
use server::WriterState;
use server::shared::SharedBytes;
use server::encoding::{ContentEncoder, TransferEncoding};

use client::ClientRequest;

const AVERAGE_HEADER_SIZE: usize = 30;

bitflags! {
    struct Flags: u8 {
        const STARTED = 0b0000_0001;
        const UPGRADE = 0b0000_0010;
        const KEEPALIVE = 0b0000_0100;
        const DISCONNECTED = 0b0000_1000;
    }
}

pub(crate) struct HttpClientWriter {
    flags: Flags,
    written: u64,
    headers_size: u32,
    buffer: SharedBytes,
    buffer_capacity: usize,
    encoder: ContentEncoder,
}

impl HttpClientWriter {

    pub fn new(buffer: SharedBytes) -> HttpClientWriter {
        let encoder = ContentEncoder::Identity(TransferEncoding::eof(buffer.clone()));
        HttpClientWriter {
            flags: Flags::empty(),
            written: 0,
            headers_size: 0,
            buffer_capacity: 0,
            buffer,
            encoder,
        }
    }

    pub fn disconnected(&mut self) {
        self.buffer.take();
    }

    // pub fn keepalive(&self) -> bool {
    //    self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
    // }

    fn write_to_stream<T: AsyncWrite>(&mut self, stream: &mut T) -> io::Result<WriterState> {
        while !self.buffer.is_empty() {
            match stream.write(self.buffer.as_ref()) {
                Ok(0) => {
                    self.disconnected();
                    return Ok(WriterState::Done);
                },
                Ok(n) => {
                    let _ = self.buffer.split_to(n);
                },
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    if self.buffer.len() > self.buffer_capacity {
                        return Ok(WriterState::Pause)
                    } else {
                        return Ok(WriterState::Done)
                    }
                }
                Err(err) => return Err(err),
            }
        }
        Ok(WriterState::Done)
    }
}

impl HttpClientWriter {

    pub fn start(&mut self, msg: &mut ClientRequest) -> io::Result<()> {
        // prepare task
        self.flags.insert(Flags::STARTED);
        self.encoder = content_encoder(self.buffer.clone(), msg);

        // render message
        {
            let mut buffer = self.buffer.get_mut();
            if let Body::Binary(ref bytes) = *msg.body() {
                buffer.reserve(256 + msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
            } else {
                buffer.reserve(256 + msg.headers().len() * AVERAGE_HEADER_SIZE);
            }

            if msg.upgrade() {
                self.flags.insert(Flags::UPGRADE);
            }

            // status line
            let _ = write!(buffer, "{} {} {:?}\r\n",
                           msg.method(),
                           msg.uri().path_and_query().map(|u| u.as_str()).unwrap_or(""),
                           msg.version());

            // write headers
            for (key, value) in msg.headers() {
                let v = value.as_ref();
                let k = key.as_str().as_bytes();
                buffer.reserve(k.len() + v.len() + 4);
                buffer.put_slice(k);
                buffer.put_slice(b": ");
                buffer.put_slice(v);
                buffer.put_slice(b"\r\n");
            }

            // set date header
            if !msg.headers().contains_key(DATE) {
                buffer.extend_from_slice(b"date: ");
                set_date(&mut buffer);
                buffer.extend_from_slice(b"\r\n\r\n");
            } else {
                buffer.extend_from_slice(b"\r\n");
            }
            self.headers_size = buffer.len() as u32;

            if msg.body().is_binary() {
                if let Body::Binary(bytes) = msg.replace_body(Body::Empty) {
                    self.written += bytes.len() as u64;
                    self.encoder.write(bytes)?;
                }
            } else {
                self.buffer_capacity = msg.write_buffer_capacity();
            }
        }
        Ok(())
    }

    pub fn write(&mut self, payload: Binary) -> io::Result<WriterState> {
        self.written += payload.len() as u64;
        if !self.flags.contains(Flags::DISCONNECTED) {
            if self.flags.contains(Flags::UPGRADE) {
                self.buffer.extend(payload);
            } else {
                self.encoder.write(payload)?;
            }
        }

        if self.buffer.len() > self.buffer_capacity {
            Ok(WriterState::Pause)
        } else {
            Ok(WriterState::Done)
        }
    }

    pub fn write_eof(&mut self) -> io::Result<()> {
        self.encoder.write_eof()?;

        if self.encoder.is_eof() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other,
                               "Last payload item, but eof is not reached"))
        }
    }

    #[inline]
    pub fn poll_completed<T: AsyncWrite>(&mut self, stream: &mut T, shutdown: bool)
                                         -> Poll<(), io::Error>
    {
        match self.write_to_stream(stream) {
            Ok(WriterState::Done) => {
                if shutdown {
                    stream.shutdown()
                } else {
                    Ok(Async::Ready(()))
                }
            },
            Ok(WriterState::Pause) => Ok(Async::NotReady),
            Err(err) => Err(err)
        }
    }
}

fn content_encoder(buf: SharedBytes, req: &mut ClientRequest) -> ContentEncoder {
    let version = req.version();
    let mut body = req.replace_body(Body::Empty);
    let mut encoding = req.content_encoding();

    let transfer = match body {
        Body::Empty => {
            req.headers_mut().remove(CONTENT_LENGTH);
            TransferEncoding::length(0, buf)
        },
        Body::Binary(ref mut bytes) => {
            if encoding.is_compression() {
                let tmp = SharedBytes::default();
                let transfer = TransferEncoding::eof(tmp.clone());
                let mut enc = match encoding {
                    ContentEncoding::Deflate => ContentEncoder::Deflate(
                        DeflateEncoder::new(transfer, Compression::default())),
                    ContentEncoding::Gzip => ContentEncoder::Gzip(
                        GzEncoder::new(transfer, Compression::default())),
                    ContentEncoding::Br => ContentEncoder::Br(
                        BrotliEncoder::new(transfer, 5)),
                    ContentEncoding::Identity => ContentEncoder::Identity(transfer),
                    ContentEncoding::Auto => unreachable!()
                };
                // TODO return error!
                let _ = enc.write(bytes.clone());
                let _ = enc.write_eof();
                *bytes = Binary::from(tmp.take());

                req.headers_mut().insert(
                    CONTENT_ENCODING, HeaderValue::from_static(encoding.as_str()));
                encoding = ContentEncoding::Identity;
            }
            let mut b = BytesMut::new();
            let _ = write!(b, "{}", bytes.len());
            req.headers_mut().insert(
                CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
            TransferEncoding::eof(buf)
        },
        Body::Streaming(_) | Body::Actor(_) => {
            if req.upgrade() {
                if version == Version::HTTP_2 {
                    error!("Connection upgrade is forbidden for HTTP/2");
                } else {
                    req.headers_mut().insert(CONNECTION, HeaderValue::from_static("upgrade"));
                }
                if encoding != ContentEncoding::Identity {
                    encoding = ContentEncoding::Identity;
                    req.headers_mut().remove(CONTENT_ENCODING);
                }
                TransferEncoding::eof(buf)
            } else {
                streaming_encoding(buf, version, req)
            }
        }
    };

    if encoding.is_compression() {
        req.headers_mut().insert(
            CONTENT_ENCODING, HeaderValue::from_static(encoding.as_str()));
    }

    req.replace_body(body);
    match encoding {
        ContentEncoding::Deflate => ContentEncoder::Deflate(
            DeflateEncoder::new(transfer, Compression::default())),
        ContentEncoding::Gzip => ContentEncoder::Gzip(
            GzEncoder::new(transfer, Compression::default())),
        ContentEncoding::Br => ContentEncoder::Br(
            BrotliEncoder::new(transfer, 5)),
        ContentEncoding::Identity | ContentEncoding::Auto => ContentEncoder::Identity(transfer),
    }
}

fn streaming_encoding(buf: SharedBytes, version: Version, req: &mut ClientRequest)
                      -> TransferEncoding {
    if req.chunked() {
        // Enable transfer encoding
        req.headers_mut().remove(CONTENT_LENGTH);
        if version == Version::HTTP_2 {
            req.headers_mut().remove(TRANSFER_ENCODING);
            TransferEncoding::eof(buf)
        } else {
            req.headers_mut().insert(
                TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
            TransferEncoding::chunked(buf)
        }
    } else {
        // if Content-Length is specified, then use it as length hint
        let (len, chunked) =
            if let Some(len) = req.headers().get(CONTENT_LENGTH) {
                // Content-Length
                if let Ok(s) = len.to_str() {
                    if let Ok(len) = s.parse::<u64>() {
                        (Some(len), false)
                    } else {
                        error!("illegal Content-Length: {:?}", len);
                        (None, false)
                    }
                } else {
                    error!("illegal Content-Length: {:?}", len);
                    (None, false)
                }
            } else {
                (None, true)
            };

        if !chunked {
            if let Some(len) = len {
                TransferEncoding::length(len, buf)
            } else {
                TransferEncoding::eof(buf)
            }
        } else {
            // Enable transfer encoding
            match version {
                Version::HTTP_11 => {
                    req.headers_mut().insert(
                        TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
                    TransferEncoding::chunked(buf)
                },
                _ => {
                    req.headers_mut().remove(TRANSFER_ENCODING);
                    TransferEncoding::eof(buf)
                }
            }
        }
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
pub const DATE_VALUE_LENGTH: usize = 29;

fn set_date(dst: &mut BytesMut) {
    CACHED.with(|cache| {
        let mut cache = cache.borrow_mut();
        let now = time::get_time();
        if now > cache.next_update {
            cache.update(now);
        }
        dst.extend_from_slice(cache.buffer());
    })
}

struct CachedDate {
    bytes: [u8; DATE_VALUE_LENGTH],
    next_update: time::Timespec,
}

thread_local!(static CACHED: RefCell<CachedDate> = RefCell::new(CachedDate {
    bytes: [0; DATE_VALUE_LENGTH],
    next_update: time::Timespec::new(0, 0),
}));

impl CachedDate {
    fn buffer(&self) -> &[u8] {
        &self.bytes[..]
    }

    fn update(&mut self, now: time::Timespec) {
        write!(&mut self.bytes[..], "{}", time::at_utc(now).rfc822()).unwrap();
        self.next_update = now + Duration::seconds(1);
        self.next_update.nsec = 0;
    }
}
<|endoftext|>


<commit_before>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_format::stream::Stream;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational32;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement};
use nom::{self, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub clusters: Vec<Cluster>,
    pub queue: VecDeque<Event>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            clusters: Vec::new(),
            queue: VecDeque::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return IResult::Done(input, ());
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            IResult::Done(i, _) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| {
                    d as u64
                });
                //self.tracks.as_ref().map(|t| t.tracks.
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            IResult::Incomplete(_) => Err(ErrorKind::MoreDataNeeded.into()),
            e => Err(ErrorKind::InvalidData.into()),
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        match segment_element(buf.data()) {
            IResult::Done(i, element) => {
                let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                match element {
                    SegmentElement::Cluster(c) => {
                        self.clusters.push(c);
                    }
                    el => {
                        println!("got element: {:#?}", el);
                    }
                }

                Ok((seek, Event::MoreDataNeeded))
            }
            IResult::Incomplete(_) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
            e => {
                println!("parsing issue: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }
}

/*
pub fn track_stream(t: &TrackEntry) -> Stream {

  Stream {
    id: t.track_uid as usize,
    index: t.track_number as usize,
    start: None,
    duration: t.default_duration,
    timebase : Rational32::from_integer(1),
    params : CodecParams {
      extradata: Vec<u8>,
      tag: u32,
      bit_rate: usize,
      bits_per_coded_sample: usize,
      profile: usize,
      level: usize,
      codec_id: CodecID,
      kind: MediaKind::

    },
  }
}
*/

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::{IResult, Offset};
    use av_format::demuxer::context::DemuxerContext;

    const mkv: &'static [u8] = include_bytes!("../assets/single_stream.mkv");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(mkv);
        println!("got parsing res: {:?}", res);
        match res {
            IResult::Done(i, _) => {
                println!("offset: {} bytes", mkv.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(mkv)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<commit_msg>draft: Populate the Stream struct from the Track information<commit_after>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational32;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement};
use nom::{self, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub clusters: Vec<Cluster>,
    pub queue: VecDeque<Event>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            clusters: Vec::new(),
            queue: VecDeque::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return IResult::Done(input, ());
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            IResult::Done(i, _) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| {
                    d as u64
                });
                //self.tracks.as_ref().map(|t| t.tracks.
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            IResult::Incomplete(_) => Err(ErrorKind::MoreDataNeeded.into()),
            e => Err(ErrorKind::InvalidData.into()),
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        match segment_element(buf.data()) {
            IResult::Done(i, element) => {
                let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                match element {
                    SegmentElement::Cluster(c) => {
                        self.clusters.push(c);
                    }
                    el => {
                        println!("got element: {:#?}", el);
                    }
                }

                Ok((seek, Event::MoreDataNeeded))
            }
            IResult::Incomplete(_) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
            e => {
                println!("parsing issue: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9"  => Some(CodecID::VP9),
        _ => None
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width : video.pixel_width as usize,
            height : video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format : None
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None
    }
}

pub fn track_stream(t: &TrackEntry) -> Stream {
    Stream {
    id: t.track_uid as usize,
    index: t.track_number as usize,
    start: None,
    duration: t.default_duration,
    timebase : Rational32::from_integer(1),
    // TODO: Extend CodecParams and fill it with the remaining information
    params : CodecParams {
        extradata: t.codec_private.clone(),
        bit_rate: 0,
        codec_id: track_entry_codec_id(t),
        kind: track_entry_media_kind(t)
    },
  }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::{IResult, Offset};
    use av_format::demuxer::context::DemuxerContext;

    const mkv: &'static [u8] = include_bytes!("../assets/single_stream.mkv");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(mkv);
        println!("got parsing res: {:?}", res);
        match res {
            IResult::Done(i, _) => {
                println!("offset: {} bytes", mkv.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(mkv)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<|endoftext|>


<commit_before>
use {Field, Borrowed, Value};
use vecmath::{rotate_radians, multiply, translate, scale, shear};
use {Transform2d, Matrix2d, Color};
use {RectangleColorContext};

/// A context with color information.
pub struct ColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> Transform2d<'a> for ColorContext<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let trans = translate(x, y);
                Value(multiply(&trans, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let rot = rotate_radians(angle);
                Value(multiply(&rot, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let scale = scale(sx, sy);
                Value(multiply(&scale, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let shear = shear(sx, sy);
                Value(multiply(&shear, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> ColorContext<'a> {
    /// Creates a RectangleColorContext.
    #[inline(always)]
    pub fn rect(&'a self, x: f64, y: f64, w: f64, h: f64) -> RectangleColorContext<'a> {
        RectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Borrowed(self.color.get()),
            rect: Value([x, y, w, h]),
        }
    }
}

#[test]
fn test_rect() {
    use {Context, AddColor};

    let c = Context::new();
    let color = c.rgba(1.0, 0.0, 0.0, 1.0);
    let rect_color = color.rect(0.0, 0.0, 100.0, 100.0);
    assert_eq!(rect_color.rect.get()[2], 100.0);
}

<commit_msg>Made ColorContext implement AddRectangle<commit_after>
use {Field, Borrowed, Value};
use vecmath::{rotate_radians, multiply, translate, scale, shear};
use {Transform2d, Matrix2d, Color};
use {RectangleColorContext};
use {AddRectangle};

/// A context with color information.
pub struct ColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> Transform2d<'a> for ColorContext<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let trans = translate(x, y);
                Value(multiply(&trans, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let rot = rotate_radians(angle);
                Value(multiply(&rot, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let scale = scale(sx, sy);
                Value(multiply(&scale, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let shear = shear(sx, sy);
                Value(multiply(&shear, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> AddRectangle<'a, RectangleColorContext<'a>> for ColorContext<'a> {
    #[inline(always)]
    fn rect(&'a self, x: f64, y: f64, w: f64, h: f64) -> RectangleColorContext<'a> {
        RectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Borrowed(self.color.get()),
            rect: Value([x, y, w, h]),
        }
    }
}

#[test]
fn test_rect() {
    use {Context, AddColor};

    let c = Context::new();
    let color = c.rgba(1.0, 0.0, 0.0, 1.0);
    let rect_color = color.rect(0.0, 0.0, 100.0, 100.0);
    assert_eq!(rect_color.rect.get()[2], 100.0);
}

<|endoftext|>