<commit_before>use jeebie::memory::MMU;
use jeebie::registers::*;

#[derive(Debug)]
pub struct CPU<'a> {
    pub reg: Registers,
    pub mem: &'a mut MMU,
}

impl<'a> CPU<'a> {
    pub fn new(mmu: &'a mut MMU) -> CPU<'a> {
        let r = Registers::new();
        CPU { reg: r, mem: mmu }
    }

    /// executes the instruction
    pub fn dispatch(&mut self, opcode: u8) {}

    /// reads instructions and executes them
    pub fn fetch_and_exec() {
        let mut mmu = MMU::new();
        let mut cpu = CPU::new(&mut mmu);

        loop {
            // fetch
            let opcode = cpu.mem.read_b(cpu.reg.pc);
            // execute
            cpu.dispatch(opcode);
            // increase PC
            cpu.reg.pc.wrapping_add(1);
            // TODO: compute clock timings
        }
    }

    fn combine_as_u16(high: u8, low: u8) -> u16 {
        ((high as u16) << 8) & (low as u16)
    }

    pub fn get8(&mut self, reg: Register8) -> u8 {
        match reg {
            Register8::A => self.reg.a,
            Register8::B => self.reg.b,
            Register8::C => self.reg.c,
            Register8::D => self.reg.d,
            Register8::E => self.reg.e,
            Register8::H => self.reg.h,
            Register8::L => self.reg.l,
            Register8::RegisterAddress(r) => {
                let addr = self.get16(r);
                self.mem.read_b(addr)
            },
            Register8::Address(addr) => self.mem.read_b(addr),
            Register8::Immediate => self.get_immediate8(),
            Register8::Value8(n) => n,
        }
    }

    pub fn set8(&mut self, reg: Register8, value: u8) {
        match reg {
            Register8::A => self.reg.a = value,
            Register8::B => self.reg.b = value,
            Register8::C => self.reg.c = value,
            Register8::D => self.reg.d = value,
            Register8::E => self.reg.e = value,
            Register8::H => self.reg.h = value,
            Register8::L => self.reg.l = value,
            Register8::RegisterAddress(r) => {
                let addr = self.get16(r);
                self.mem.write_b(addr, value);
            },
            Register8::Address(addr) => self.mem.write_b(addr, value),
            _ => {},
        };
    }

    pub fn get16(&mut self, reg: Register16) -> u16 {
        match reg {
            Register16::AF => CPU::combine_as_u16(self.reg.a, self.reg.f),
            Register16::BC => CPU::combine_as_u16(self.reg.b, self.reg.c),
            Register16::DE => CPU::combine_as_u16(self.reg.d, self.reg.e),
            Register16::HL => CPU::combine_as_u16(self.reg.h, self.reg.l),
            Register16::SP => self.reg.sp,
            Register16::PC => self.reg.pc,
            Register16::Immediate16 => self.get_immediate16(),
            Register16::Value16(n) => n,
        }
    }

    pub fn set16(&mut self, reg: Register16, value: u16) {
        match reg {
            Register16::AF => { self.reg.a = (value >> 8) as u8 ; self.reg.f = value as u8; },
            Register16::BC => { self.reg.b = (value >> 8) as u8 ; self.reg.c = value as u8; },
            Register16::DE => { self.reg.d = (value >> 8) as u8 ; self.reg.e = value as u8; },
            Register16::HL => { self.reg.h = (value >> 8) as u8 ; self.reg.l = value as u8; },
            Register16::SP => { self.reg.sp = value },
            Register16::PC => { self.reg.pc = value },
            _ => {},
        };
    }

    // Swaps low and high nibble of an 8 bit value and sets flags.
    // Returns the result of the swap operation.
    pub fn compute_swap(&mut self, reg: Register8) {
        let value = self.get8(reg);
        let result = (value << 4) | (value >> 4);

        self.reg.clear_all_flags();
        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(reg, result);
    }

    pub fn load_rr(&mut self, reg1: Register8, reg2: Register8) {
        let value = self.get8(reg2);
        self.set8(reg1, value);
    }

    pub fn load_rr16(&mut self, reg1: Register16, reg2: Register16) {
        let value = self.get16(reg2);
        self.set16(reg1, value);
    }

    // Computes the flags and result for a 16-bit ADD instruction.
    // The result is put in the specified `reg1`.
    pub fn compute_add16(&mut self, reg1: Register16, reg2: Register16) {
        let lhs = self.get16(reg1);
        let rhs = self.get16(reg2);
        let result = lhs.wrapping_add(rhs);

        // flag zero is not affected, sub is reset
        self.reg.clear_flag(Flags::Sub);

        // checked add returns None if add overflows (carry from bit 15)
        if let None = lhs.checked_add(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        // halfcarry is set if bit 11 has carry
        // sum the 12 LSb and check bit 12
        let low_result = (rhs & 0xFFF) + (lhs & 0xFFF);
        if (low_result >> 12) == 1 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set16(reg1, result);
    }

    // Computes the flags and result for an 8-bit ADD instruction.
    pub fn compute_add(&mut self, reg1: Register8, reg2: Register8) {
        let lhs = self.get8(reg1);
        let rhs = self.get8(reg2);
        let result = lhs.wrapping_add(rhs);

        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // checked add returns None if add overflows
        if let None = lhs.checked_add(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        let low_result = (rhs & 0x0F) + (lhs & 0x0F);
        // check if bit 4 is set
        if (low_result >> 4) == 1 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(Register8::A, result);
    }

    pub fn compute_adc(&mut self, reg1: Register8, reg2: Register8) {
        let carry = if self.reg.is_set(Flags::Carry) { 1 } else { 0 };
        let reg2 = Register8::Value8(self.get8(reg2).wrapping_add(carry));
        self.compute_add(reg1, reg2);
    }

    // Computes the flags and result for a SUB instruction.
    // Left hand operator is always register A
    pub fn compute_sub(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs.wrapping_sub(rhs);

        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.reg.set_flag(Flags::Sub);

        // checked sub returns None if sub borrows (carry flag)
        if let None = lhs.checked_sub(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        // same as carry flag but values limited to their 4 low bits (half carry flag)
        if let None = (lhs & 0xF).checked_sub(rhs & 0xF) {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(Register8::A, result);
    }

    pub fn compute_sbc(&mut self, reg: Register8) {
        let carry = if self.reg.is_set(Flags::Carry) { 1 } else { 0 };
        let reg = Register8::Value8(self.get8(reg).wrapping_sub(carry));
        self.compute_sub(reg);
    }

    // Computes the flags for a CP instruction.
    // this has the same effect on flags as a SUB, but the result is discarded.
    pub fn compute_cp(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);

        self.reg.clear_all_flags();
        self.reg.set_flag(Flags::Sub);

        // Same as A - n
        if lhs == rhs {
            self.reg.set_flag(Flags::Zero);
        }

        // if A < n
        if lhs < rhs {
            self.reg.set_flag(Flags::Carry);
        }

        // if 4 low bits of A < 4 low bits of n
        if (lhs & 0xF) < (rhs & 0xF) {
            self.reg.set_flag(Flags::HalfCarry);
        }
    }

    // Computes the flags and result for an AND instruction.
    // lhs is always the register A
    pub fn compute_and(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs & rhs;

        // HC flag is always set, other flags always cleared except zero.
        self.reg.clear_all_flags();
        self.reg.set_flag(Flags::HalfCarry);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    // Computes the flags and result for an OR instruction.
    // lhs is always the register A
    pub fn compute_or(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs | rhs;

        // all flags are cleared except zero.
        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    // Computes the flags and result for an XOR instruction.
    // lhs is always the register A
    pub fn compute_xor(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs ^ rhs;

        // all flags are cleared except zero.
        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    /// Computes an INC instruction.
    /// Carry flag is left untouched, so they are not cleared.
    pub fn compute_inc(&mut self, reg: Register8) {
        let result = self.get8(reg).wrapping_add(1);
        self.reg.clear_flag(Flags::Sub);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // set HC if bit 0-3 were 1 before adding
        if (result.wrapping_sub(1) & 0x0F) == 0x0F {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(reg, result);
    }

    /// Computes an INC instruction on a 16 bit register.
    /// Flags are not affected.
    pub fn compute_inc16(&mut self, reg: Register16) {
        let value = self.get16(reg).wrapping_add(1);
        self.set16(reg, value);
    }

    /// Computes an INC instruction on a 16 bit register.
    /// Flags are not affected.
    pub fn compute_dec16(&mut self, reg: Register16) {
        let value = self.get16(reg).wrapping_sub(1);
        self.set16(reg, value);
    }

    /// Computes flags after a DEC instruction based on the final value.
    /// Carry flag is left untouched, so no clearing all of them.
    pub fn compute_dec(&mut self, reg: Register8) {
        let result = self.get8(reg).wrapping_sub(1);
        self.reg.set_flag(Flags::Sub);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // HC is set if bit 4 was borrowed
        // this happens only when bit 0-3 are all 0 before decrementing.
        if (result.wrapping_add(1) & 0x0F) == 0x00 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(reg, result);
    }

    /// Pushes on the stack a 16-bit register value,
    /// it decrements SP and pushes the LSB before the MSB.
    pub fn push_stack(&mut self, reg: Register16) {
        let value = self.get16(reg);

        self.reg.sp.wrapping_sub(1);
        let addr = self.reg.sp;
        let low = (value & 0x00FF) as u8;
        self.mem.write_b(addr, low);

        self.reg.sp.wrapping_sub(1);
        let addr = self.reg.sp;
        let high = ((value >> 8) & 0x00FF) as u8;
        self.mem.write_b(addr, high);
    }

    /// Pops a 16-bit value from the stack, MSB first, returning the u16 value.
    pub fn pop_stack(&mut self, dest: Register16) {
        let high = self.mem.read_b(self.reg.sp);
        self.reg.sp.wrapping_add(1);

        let low = self.mem.read_b(self.reg.sp);
        self.reg.sp.wrapping_add(1);

        let result = ((high as u16) << 8) & (low as u16);
        self.set16(dest, result);
    }

    /// Retrieves an immediate 8-bit value.
    /// Immediates are retrieved by reading at the address in the PC register.
    pub fn get_immediate8(&mut self) -> u8 {
        let value = self.mem.read_b(self.reg.pc);
        self.reg.pc.wrapping_add(1);

        value
    }

    /// Retrieves an immediate 16-bit value.
    /// 16-bit immediates are read as two 8-bit immediates, the first being the LSB.
    pub fn get_immediate16(&mut self) -> u16 {
        let low = self.get_immediate8();
        let high = self.get_immediate8();

        (high as u16) & (low as u16)
    }
}
<commit_msg>Added instruction execution and timing counter<commit_after>use jeebie::memory::MMU;
use jeebie::registers::*;

use jeebie::opcodes::{ CB_OPCODE_TABLE, OPCODE_TABLE };
use jeebie::timings::{ CB_TIMING_TABLE, TIMING_TABLE };

#[derive(Debug)]
pub struct CPU<'a> {
    pub reg: Registers,
    pub mem: &'a mut MMU,

    // amount of machine cycles (as reported in timing tables) elapsed.
    cycles: u32,

}

impl<'a> CPU<'a> {
    pub fn new(mmu: &'a mut MMU) -> CPU<'a> {
        let r = Registers::new();
        CPU { reg: r, mem: mmu, cycles: 0}
    }

    /// executes one instruction, updating cycles and PC register accordingly
    pub fn exec(&mut self) {
        // fetch
        let opcode = self.mem.read_b(self.reg.pc);
        self.reg.pc.wrapping_add(1);

        let instr_timing = match opcode {
            0xCB => {
                // 2-byte opcodes are prefixed with 0xCB
                let second_byte = self.mem.read_b(self.reg.pc);
                self.reg.pc.wrapping_add(1);

                CB_OPCODE_TABLE[second_byte as usize](self);
                CB_TIMING_TABLE[second_byte as usize]
            },
            _ => {
                OPCODE_TABLE[opcode as usize](self);
                TIMING_TABLE[opcode as usize]
            }
        };

        self.cycles.wrapping_add(instr_timing as u32);
    }

    pub fn exec_one_frame(&mut self) {

    }

    fn combine_as_u16(high: u8, low: u8) -> u16 {
        ((high as u16) << 8) & (low as u16)
    }

    pub fn get8(&mut self, reg: Register8) -> u8 {
        match reg {
            Register8::A => self.reg.a,
            Register8::B => self.reg.b,
            Register8::C => self.reg.c,
            Register8::D => self.reg.d,
            Register8::E => self.reg.e,
            Register8::H => self.reg.h,
            Register8::L => self.reg.l,
            Register8::RegisterAddress(r) => {
                let addr = self.get16(r);
                self.mem.read_b(addr)
            },
            Register8::Address(addr) => self.mem.read_b(addr),
            Register8::Immediate => self.get_immediate8(),
            Register8::Value8(n) => n,
        }
    }

    pub fn set8(&mut self, reg: Register8, value: u8) {
        match reg {
            Register8::A => self.reg.a = value,
            Register8::B => self.reg.b = value,
            Register8::C => self.reg.c = value,
            Register8::D => self.reg.d = value,
            Register8::E => self.reg.e = value,
            Register8::H => self.reg.h = value,
            Register8::L => self.reg.l = value,
            Register8::RegisterAddress(r) => {
                let addr = self.get16(r);
                self.mem.write_b(addr, value);
            },
            Register8::Address(addr) => self.mem.write_b(addr, value),
            _ => {},
        };
    }

    pub fn get16(&mut self, reg: Register16) -> u16 {
        match reg {
            Register16::AF => CPU::combine_as_u16(self.reg.a, self.reg.f),
            Register16::BC => CPU::combine_as_u16(self.reg.b, self.reg.c),
            Register16::DE => CPU::combine_as_u16(self.reg.d, self.reg.e),
            Register16::HL => CPU::combine_as_u16(self.reg.h, self.reg.l),
            Register16::SP => self.reg.sp,
            Register16::PC => self.reg.pc,
            Register16::Immediate16 => self.get_immediate16(),
            Register16::Value16(n) => n,
        }
    }

    pub fn set16(&mut self, reg: Register16, value: u16) {
        match reg {
            Register16::AF => { self.reg.a = (value >> 8) as u8 ; self.reg.f = value as u8; },
            Register16::BC => { self.reg.b = (value >> 8) as u8 ; self.reg.c = value as u8; },
            Register16::DE => { self.reg.d = (value >> 8) as u8 ; self.reg.e = value as u8; },
            Register16::HL => { self.reg.h = (value >> 8) as u8 ; self.reg.l = value as u8; },
            Register16::SP => { self.reg.sp = value },
            Register16::PC => { self.reg.pc = value },
            _ => {},
        };
    }

    // Swaps low and high nibble of an 8 bit value and sets flags.
    // Returns the result of the swap operation.
    pub fn compute_swap(&mut self, reg: Register8) {
        let value = self.get8(reg);
        let result = (value << 4) | (value >> 4);

        self.reg.clear_all_flags();
        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(reg, result);
    }

    pub fn load_rr(&mut self, reg1: Register8, reg2: Register8) {
        let value = self.get8(reg2);
        self.set8(reg1, value);
    }

    pub fn load_rr16(&mut self, reg1: Register16, reg2: Register16) {
        let value = self.get16(reg2);
        self.set16(reg1, value);
    }

    // Computes the flags and result for a 16-bit ADD instruction.
    // The result is put in the specified `reg1`.
    pub fn compute_add16(&mut self, reg1: Register16, reg2: Register16) {
        let lhs = self.get16(reg1);
        let rhs = self.get16(reg2);
        let result = lhs.wrapping_add(rhs);

        // flag zero is not affected, sub is reset
        self.reg.clear_flag(Flags::Sub);

        // checked add returns None if add overflows (carry from bit 15)
        if let None = lhs.checked_add(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        // halfcarry is set if bit 11 has carry
        // sum the 12 LSb and check bit 12
        let low_result = (rhs & 0xFFF) + (lhs & 0xFFF);
        if (low_result >> 12) == 1 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set16(reg1, result);
    }

    // Computes the flags and result for an 8-bit ADD instruction.
    pub fn compute_add(&mut self, reg1: Register8, reg2: Register8) {
        let lhs = self.get8(reg1);
        let rhs = self.get8(reg2);
        let result = lhs.wrapping_add(rhs);

        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // checked add returns None if add overflows
        if let None = lhs.checked_add(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        let low_result = (rhs & 0x0F) + (lhs & 0x0F);
        // check if bit 4 is set
        if (low_result >> 4) == 1 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(Register8::A, result);
    }

    pub fn compute_adc(&mut self, reg1: Register8, reg2: Register8) {
        let carry = if self.reg.is_set(Flags::Carry) { 1 } else { 0 };
        let reg2 = Register8::Value8(self.get8(reg2).wrapping_add(carry));
        self.compute_add(reg1, reg2);
    }

    // Computes the flags and result for a SUB instruction.
    // Left hand operator is always register A
    pub fn compute_sub(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs.wrapping_sub(rhs);

        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.reg.set_flag(Flags::Sub);

        // checked sub returns None if sub borrows (carry flag)
        if let None = lhs.checked_sub(rhs) {
            self.reg.set_flag(Flags::Carry);
        }

        // same as carry flag but values limited to their 4 low bits (half carry flag)
        if let None = (lhs & 0xF).checked_sub(rhs & 0xF) {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(Register8::A, result);
    }

    pub fn compute_sbc(&mut self, reg: Register8) {
        let carry = if self.reg.is_set(Flags::Carry) { 1 } else { 0 };
        let reg = Register8::Value8(self.get8(reg).wrapping_sub(carry));
        self.compute_sub(reg);
    }

    // Computes the flags for a CP instruction.
    // this has the same effect on flags as a SUB, but the result is discarded.
    pub fn compute_cp(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);

        self.reg.clear_all_flags();
        self.reg.set_flag(Flags::Sub);

        // Same as A - n
        if lhs == rhs {
            self.reg.set_flag(Flags::Zero);
        }

        // if A < n
        if lhs < rhs {
            self.reg.set_flag(Flags::Carry);
        }

        // if 4 low bits of A < 4 low bits of n
        if (lhs & 0xF) < (rhs & 0xF) {
            self.reg.set_flag(Flags::HalfCarry);
        }
    }

    // Computes the flags and result for an AND instruction.
    // lhs is always the register A
    pub fn compute_and(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs & rhs;

        // HC flag is always set, other flags always cleared except zero.
        self.reg.clear_all_flags();
        self.reg.set_flag(Flags::HalfCarry);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    // Computes the flags and result for an OR instruction.
    // lhs is always the register A
    pub fn compute_or(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs | rhs;

        // all flags are cleared except zero.
        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    // Computes the flags and result for an XOR instruction.
    // lhs is always the register A
    pub fn compute_xor(&mut self, reg: Register8) {
        let lhs = self.get8(Register8::A);
        let rhs = self.get8(reg);
        let result = lhs ^ rhs;

        // all flags are cleared except zero.
        self.reg.clear_all_flags();

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        self.set8(Register8::A, result);
    }

    /// Computes an INC instruction.
    /// Carry flag is left untouched, so they are not cleared.
    pub fn compute_inc(&mut self, reg: Register8) {
        let result = self.get8(reg).wrapping_add(1);
        self.reg.clear_flag(Flags::Sub);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // set HC if bit 0-3 were 1 before adding
        if (result.wrapping_sub(1) & 0x0F) == 0x0F {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(reg, result);
    }

    /// Computes an INC instruction on a 16 bit register.
    /// Flags are not affected.
    pub fn compute_inc16(&mut self, reg: Register16) {
        let value = self.get16(reg).wrapping_add(1);
        self.set16(reg, value);
    }

    /// Computes an INC instruction on a 16 bit register.
    /// Flags are not affected.
    pub fn compute_dec16(&mut self, reg: Register16) {
        let value = self.get16(reg).wrapping_sub(1);
        self.set16(reg, value);
    }

    /// Computes flags after a DEC instruction based on the final value.
    /// Carry flag is left untouched, so no clearing all of them.
    pub fn compute_dec(&mut self, reg: Register8) {
        let result = self.get8(reg).wrapping_sub(1);
        self.reg.set_flag(Flags::Sub);

        if result == 0 {
            self.reg.set_flag(Flags::Zero);
        }

        // HC is set if bit 4 was borrowed
        // this happens only when bit 0-3 are all 0 before decrementing.
        if (result.wrapping_add(1) & 0x0F) == 0x00 {
            self.reg.set_flag(Flags::HalfCarry);
        }

        self.set8(reg, result);
    }

    /// Pushes on the stack a 16-bit register value,
    /// it decrements SP and pushes the LSB before the MSB.
    pub fn push_stack(&mut self, reg: Register16) {
        let value = self.get16(reg);

        self.reg.sp.wrapping_sub(1);
        let addr = self.reg.sp;
        let low = (value & 0x00FF) as u8;
        self.mem.write_b(addr, low);

        self.reg.sp.wrapping_sub(1);
        let addr = self.reg.sp;
        let high = ((value >> 8) & 0x00FF) as u8;
        self.mem.write_b(addr, high);
    }

    /// Pops a 16-bit value from the stack, MSB first, returning the u16 value.
    pub fn pop_stack(&mut self, dest: Register16) {
        let high = self.mem.read_b(self.reg.sp);
        self.reg.sp.wrapping_add(1);

        let low = self.mem.read_b(self.reg.sp);
        self.reg.sp.wrapping_add(1);

        let result = ((high as u16) << 8) & (low as u16);
        self.set16(dest, result);
    }

    /// Retrieves an immediate 8-bit value.
    /// Immediates are retrieved by reading at the address in the PC register.
    pub fn get_immediate8(&mut self) -> u8 {
        let value = self.mem.read_b(self.reg.pc);
        self.reg.pc.wrapping_add(1);

        value
    }

    /// Retrieves an immediate 16-bit value.
    /// 16-bit immediates are read as two 8-bit immediates, the first being the LSB.
    pub fn get_immediate16(&mut self) -> u16 {
        let low = self.get_immediate8();
        let high = self.get_immediate8();

        (high as u16) & (low as u16)
    }
}
<|endoftext|>


<commit_before>mod page;
pub use self::page::Page;

use std::env;
use std::io::Read;
use std::time::Duration;
#[allow(unused_imports)]
use std::ascii::AsciiExt;
use std::default::Default;
use std::collections::BTreeMap;

use url::Url;
use mime;
use reqwest::Client as HttpClient;
use reqwest::header::{UserAgent, ContentType};
use reqwest::StatusCode;
use reqwest::Response as HttpResponse;
use robotparser::RobotFileParser;
use html5ever::tokenizer::Tokenizer;
use html5ever::tokenizer::BufferQueue;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_8;
use url_serde::Serde;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Serde<Url>>,
    pub unvisited_urls: Vec<Serde<Url>>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(redis_pool: RedisPool,
               base_url: Url,
               limit: isize,
               mime_types: Vec<mime::Mime>)
               -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts { namespace: Some(maman_env.to_string()) };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url: base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq: sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit: limit,
            mime_types: mime_types,
        }
    }

    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser
            .can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for h in response.headers().iter() {
                headers.insert(h.name().to_ascii_lowercase(), h.value_string());
            }
        }
        let mut document = vec![];
        let _ = response.read_to_end(&mut document);
        match UTF_8.decode(&document, DecoderTrap::Ignore) {
            Ok(doc) => {
                let page = Page::new(page_url.clone(),
                                     doc.to_string(),
                                     headers.clone(),
                                     response.status().to_string());
                let read = Spider::read_page(page, &doc);
                Some(read.sink)
            }
            Err(_) => None,
        }
    }

    fn load_url(url: &str, mime_types: &Vec<mime::Mime>) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build().expect("HttpClient failed to construct");
        match client.get(url).header(UserAgent::new(maman_user_agent!().to_owned())).send() {
            Err(_) => None,
            Ok(response) => {
                match response.status() {
                    StatusCode::Ok | StatusCode::NotModified => {
                        if mime_types.is_empty() {
                            Some(response)
                        } else if mime_types.contains(response
                                                          .headers()
                                                          .get::<ContentType>()
                                                          .unwrap()) {
                            Some(response)
                        } else {
                            None
                        }
                    }
                    _ => None,
                }
            }
        }
    }
}
<commit_msg>Fix clippy warnings<commit_after>mod page;
pub use self::page::Page;

use std::env;
use std::io::Read;
use std::time::Duration;
#[allow(unused_imports)]
use std::ascii::AsciiExt;
use std::default::Default;
use std::collections::BTreeMap;

use url::Url;
use mime;
use reqwest::Client as HttpClient;
use reqwest::header::{UserAgent, ContentType};
use reqwest::StatusCode;
use reqwest::Response as HttpResponse;
use robotparser::RobotFileParser;
use html5ever::tokenizer::Tokenizer;
use html5ever::tokenizer::BufferQueue;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_8;
use url_serde::Serde;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Serde<Url>>,
    pub unvisited_urls: Vec<Serde<Url>>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(redis_pool: RedisPool,
               base_url: Url,
               limit: isize,
               mime_types: Vec<mime::Mime>)
               -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts { namespace: Some(maman_env.to_string()) };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url: base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq: sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit: limit,
            mime_types: mime_types,
        }
    }

    #[allow(unknown_lints, needless_pass_by_value)]
    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser
            .can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for h in response.headers().iter() {
                headers.insert(h.name().to_ascii_lowercase(), h.value_string());
            }
        }
        let mut document = vec![];
        let _ = response.read_to_end(&mut document);
        match UTF_8.decode(&document, DecoderTrap::Ignore) {
            Ok(doc) => {
                let page = Page::new(page_url.clone(),
                                     doc.to_string(),
                                     headers.clone(),
                                     response.status().to_string());
                let read = Spider::read_page(page, &doc);
                Some(read.sink)
            }
            Err(_) => None,
        }
    }

    fn load_url(url: &str, mime_types: &[mime::Mime]) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build().expect("HttpClient failed to construct");
        match client.get(url).header(UserAgent::new(maman_user_agent!().to_owned())).send() {
            Err(_) => None,
            Ok(response) => {
                match response.status() {
                    StatusCode::Ok | StatusCode::NotModified => {
                        if mime_types.is_empty() || mime_types.contains(response
                                                                        .headers()
                                                                        .get::<ContentType>()
                                                                        .unwrap()) {
                            Some(response)
                        } else {
                            None
                        }
                    }
                    _ => None,
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Defines a series of convenience modifiers for editing Responses
//!
//! Modifiers can be used to edit Responses through the owning
//! method `set` or the mutating `set_mut`, both of which are
//! defined in the `Set` trait.
//!
//! Instead of having a combinatorial explosion of Response methods
//! and constructors, this provides a series of modifiers that
//! can be used through the `Set` trait.
//!
//! For instance, instead of `Response::redirect` constructing a
//! redirect response, we provide a `Redirect` modifier, so you
//! can just do:
//!
//! ```rust,ignore
//! Response::new()
//!     .set(status)
//!     .set(Redirect(url));
//! ```
//!
//! This is more extensible as it allows you to combine
//! arbitrary modifiers without having a massive number of
//! Response constructors.
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::io::{File, MemReader};
use std::path::Path;

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Response, Url};

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<Reader + Send> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(MemReader::new(self)) as Box<Reader + Send>);
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_string().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    #[inline]
    fn modify(self, res: &mut Response) {
        // Also set the content type.
        // self.path().extension_str()
        //     .and_then(get_content_type)
        //     .and_then(|ct| { res.set_mut(ct) });
        res.body = Some(Box::new(self) as Box<Reader + Send>);
    }
}

impl Modifier<Response> for Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .ok().expect(format!("No such file: {}", self.display()).as_slice())
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(format!("{:?}", url)));
    }
}

<commit_msg>(fix) Fixed bug in redirect code<commit_after>//! Defines a series of convenience modifiers for editing Responses
//!
//! Modifiers can be used to edit Responses through the owning
//! method `set` or the mutating `set_mut`, both of which are
//! defined in the `Set` trait.
//!
//! Instead of having a combinatorial explosion of Response methods
//! and constructors, this provides a series of modifiers that
//! can be used through the `Set` trait.
//!
//! For instance, instead of `Response::redirect` constructing a
//! redirect response, we provide a `Redirect` modifier, so you
//! can just do:
//!
//! ```rust,ignore
//! Response::new()
//!     .set(status)
//!     .set(Redirect(url));
//! ```
//!
//! This is more extensible as it allows you to combine
//! arbitrary modifiers without having a massive number of
//! Response constructors.
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::io::{File, MemReader};
use std::path::Path;

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Response, Url};

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<Reader + Send> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(MemReader::new(self)) as Box<Reader + Send>);
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_string().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    #[inline]
    fn modify(self, res: &mut Response) {
        // Also set the content type.
        // self.path().extension_str()
        //     .and_then(get_content_type)
        //     .and_then(|ct| { res.set_mut(ct) });
        res.body = Some(Box::new(self) as Box<Reader + Send>);
    }
}

impl Modifier<Response> for Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .ok().expect(format!("No such file: {}", self.display()).as_slice())
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(url.to_string()));
    }
}

<|endoftext|>


<commit_before>//! This module contains a trait representing the (real-valued) objective associated with a
//! discrete optimization problem. The greedy algorithm contained in this crate is implemented for
//! any struct implementing `Objective`.
use std::fmt::Debug;
use std::hash::Hash;
use errors::*;
use fnv::FnvHashSet;
use setlike::Setlike;

pub type Set<E> = FnvHashSet<E>;

pub type ElementIterator<'a, O> = Box<Iterator<Item = <O as Objective>::Element> + 'a>;

pub mod curvature {
    //! Marker traits for Curvature bounds.
    //! These allow compile-time verification that an algorithm called on an `Objective` has its
    //! curvature constraints satisfied.
    //!
    //! As a concrete example, the `greedy::lazier_greedy` method requires a `Submodular`
    //! objective. Previously, this was done by checking an enum and panicking if it failed. Marker
    //! traits move the error from runtime to compilation.

    /// An objective with fixed bounds that are not the (0, 1) of `Submodular` or (1, ∞) of
    /// `Supermodular`.
    pub trait Bounded: super::Objective {
        fn bounds() -> (Option<f64>, Option<f64>);
    }

    /// A submodular objective, which satisfies `∀ S ⊆ T: f(S ∪ {e}) - f(S) ≥ f(T ∪ {e}) - f(T)` or (equivalently) has curvature bounded above by `1`.
    ///
    /// The `Bounded` implementation should return a fixed `(None, Some(1.0))`.
    pub trait Submodular: Bounded {}

    /// A supermodular objective, which satisfies `∀ S ⊆ T: f(S ∪ {e}) - f(S) ≤ f(T ∪ {e}) - f(T)` or (equivalently) has curvature bounded below by `1`.
    ///
    /// The `Bounded` implementation should return a fixed `(Some(1.0), None)`.
    pub trait Supermodular: Bounded {}

    /// A modular objective. The `Bounded` implementation should return a fixed `(Some(1.0),
    /// Some(1.0))`.
    ///
    /// This trait is the least common. The semantics of this type of objective are best understood
    /// by the equation `f(S) = Σf(e)`. That is, the objective is exactly the sum of the weights of
    /// the elements of the set `S`.
    pub trait Modular: Bounded {}
    impl<T: Modular> Submodular for T {}
    impl<T: Modular> Supermodular for T {}
}

/// An objective to be optimized.
///
/// The functions are assumed to be 'internally-stateless'. That is, they do not modify the problem
/// instance, and instead track any relevant information on the associated `State` type. This
/// allows snapshotting of the state and allows safe parallel computation.
pub trait Objective: Sized {
    type Element: Eq + Hash + Clone + Copy + Debug;
    /// Holds any state that is tracked either to assist in correctness or to provide e.g.
    /// cacheing.
    type State: Default + Clone;

    /// The domain of the solution.
    fn elements(&self) -> ElementIterator<Self>;

    /// The implementation of the function `f`. This is the eponymous `Objective`.
    fn benefit<S: Setlike<Self::Element>>(&self, s: &S, state: &Self::State) -> Result<f64>;

    /// The marginal gain $\delta(u \mid S) = f(S \cup u) - f(S)$
    ///
    /// Should return `0` for $u \in S$.
    ///
    /// The default implementation is incredibly inefficient (two calls to `benefit`), so
    /// overriding this is *strongly* encouraged.
    fn delta<S: Setlike<Self::Element>>(&self, u: Self::Element, s: &S, state: &Self::State) -> Result<f64>;

    /// The primal curvature $\nabla(u, v\mid S) = \frac{\delta(u \mid S \cup v)}{\delta(u \mid
    /// S)}$.
    ///
    /// Should return `1` for $u \in S$ or when `depends(u, S)` does not contain `v`.
    ///
    /// The default implementation panics with `unimplemented!()`
    #[allow(unused_variables)]
    fn nabla<S: Setlike<Self::Element>>(&self,
             u: Self::Element,
             v: Self::Element,
             s: &S,
             state: &Self::State)
        -> Result<f64> {
            unimplemented!()
        }

    /// The elements on which the marginal gain of `u` depends. The technical definition I have
    /// been using is based on the primal curvature (`nabla`): if `u` depends on `v`, then
    /// `nabla(u, v, sol, state) != 1`.
    fn depends(&self, u: Self::Element, state: &Self::State) -> Result<ElementIterator<Self>>;
    /// Same as `depends`, except that `nabla(u, v, sol, state) > 1`.
    ///
    /// Not all objectives need to implement this. If it is not implemented, calling it panics (via
    /// `unimplemented!()`)
    #[allow(unused_variables)]
    fn supermodular_depends(&self,
                            u: Self::Element,
                            state: &Self::State)
                            -> Result<ElementIterator<Self>> {
        unimplemented!()
    }

    /// Update the state tracked in `s` to reflect the insertion of `u` into the solution.
    fn insert_mut(&self, u: Self::Element, s: &mut Self::State) -> Result<()>;

    /// Produce a new state reflecting the insertion of `u` into the solution.
    fn insert(&self, u: Self::Element, s: &Self::State) -> Result<Self::State> {
        let mut state = s.clone();
        self.insert_mut(u, &mut state)?;
        Ok(state)
    }
}

/// Lazy extension of `Objective`. While the same state type is used, state updates are handled
/// independently. If you need two different kinds of state for each kind of update, consider using
/// an enum-struct. For example:
///
/// ```
/// #[derive(Clone, Debug)]
/// pub enum State {
///     Empty,
///     Strict {
///         foo: usize,
///         bar: Vec<usize>,
///     },
///     Lazy {
///         baz: f64,
///     }
/// }
///
/// impl Default for State {
///     fn default() -> Self {
///         State::Empty
///     }
/// }
/// ```
pub trait LazyObjective: Objective {
    /// Updates the marginal gain of `element`, returning it value after updating the state to
    /// reflect the insertion of `previous` elements.
    ///
    /// Returns `Ok(None)` if the marginal gain does not need to be updated.
    fn update_lazy_mut<S: Setlike<Self::Element>>(&self,
                       element: Self::Element,
                       previous: &S,
                       state: &mut Self::State)
                       -> Result<Option<f64>>;

    /// Update the state to preserve invariants upon insertion of `element`.
    ///
    /// This should do the least amount of work possible without making `update_lazy_mut` too
    /// complex.
    fn insert_lazy_mut(&self, element: Self::Element, state: &mut Self::State) -> Result<()>;
}

/// An objective with further constraints than mere cardinality.
pub trait ConstrainedObjective: Objective {
    /// Returns true if adding `el` to `sol` would create a valid set.
    fn valid_addition(&self,
                      el: Self::Element,
                      sol: &Vec<Self::Element>,
                      state: &Self::State)
                      -> bool;
}
<commit_msg>Added IntoIterator constraint on delta/benefit.<commit_after>//! This module contains a trait representing the (real-valued) objective associated with a
//! discrete optimization problem. The greedy algorithm contained in this crate is implemented for
//! any struct implementing `Objective`.
use std::fmt::Debug;
use std::hash::Hash;
use errors::*;
use fnv::FnvHashSet;
use setlike::Setlike;

pub type Set<E> = FnvHashSet<E>;

pub type ElementIterator<'a, O> = Box<Iterator<Item = <O as Objective>::Element> + 'a>;

pub mod curvature {
    //! Marker traits for Curvature bounds.
    //! These allow compile-time verification that an algorithm called on an `Objective` has its
    //! curvature constraints satisfied.
    //!
    //! As a concrete example, the `greedy::lazier_greedy` method requires a `Submodular`
    //! objective. Previously, this was done by checking an enum and panicking if it failed. Marker
    //! traits move the error from runtime to compilation.

    /// An objective with fixed bounds that are not the (0, 1) of `Submodular` or (1, ∞) of
    /// `Supermodular`.
    pub trait Bounded: super::Objective {
        fn bounds() -> (Option<f64>, Option<f64>);
    }

    /// A submodular objective, which satisfies `∀ S ⊆ T: f(S ∪ {e}) - f(S) ≥ f(T ∪ {e}) - f(T)` or (equivalently) has curvature bounded above by `1`.
    ///
    /// The `Bounded` implementation should return a fixed `(None, Some(1.0))`.
    pub trait Submodular: Bounded {}

    /// A supermodular objective, which satisfies `∀ S ⊆ T: f(S ∪ {e}) - f(S) ≤ f(T ∪ {e}) - f(T)` or (equivalently) has curvature bounded below by `1`.
    ///
    /// The `Bounded` implementation should return a fixed `(Some(1.0), None)`.
    pub trait Supermodular: Bounded {}

    /// A modular objective. The `Bounded` implementation should return a fixed `(Some(1.0),
    /// Some(1.0))`.
    ///
    /// This trait is the least common. The semantics of this type of objective are best understood
    /// by the equation `f(S) = Σf(e)`. That is, the objective is exactly the sum of the weights of
    /// the elements of the set `S`.
    pub trait Modular: Bounded {}
    impl<T: Modular> Submodular for T {}
    impl<T: Modular> Supermodular for T {}
}

/// An objective to be optimized.
///
/// The functions are assumed to be 'internally-stateless'. That is, they do not modify the problem
/// instance, and instead track any relevant information on the associated `State` type. This
/// allows snapshotting of the state and allows safe parallel computation.
pub trait Objective: Sized {
    type Element: Eq + Hash + Clone + Copy + Debug;
    /// Holds any state that is tracked either to assist in correctness or to provide e.g.
    /// cacheing.
    type State: Default + Clone;

    /// The domain of the solution.
    fn elements(&self) -> ElementIterator<Self>;

    /// The implementation of the function `f`. This is the eponymous `Objective`.
    fn benefit<S>(&self, s: &S, state: &Self::State) -> Result<f64>
        where S: Setlike<Self::Element> + IntoIterator<Item = Self::Element>;

    /// The marginal gain $\delta(u \mid S) = f(S \cup u) - f(S)$
    ///
    /// Should return `0` for $u \in S$.
    ///
    /// The default implementation is incredibly inefficient (two calls to `benefit`), so
    /// overriding this is *strongly* encouraged.
    fn delta<S>(&self, u: Self::Element, s: &S, state: &Self::State) -> Result<f64>
        where S: Setlike<Self::Element> + IntoIterator<Item = Self::Element>;

    /// The primal curvature $\nabla(u, v\mid S) = \frac{\delta(u \mid S \cup v)}{\delta(u \mid
    /// S)}$.
    ///
    /// Should return `1` for $u \in S$ or when `depends(u, S)` does not contain `v`.
    ///
    /// The default implementation panics with `unimplemented!()`
    #[allow(unused_variables)]
    fn nabla<S: Setlike<Self::Element>>(&self,
                                        u: Self::Element,
                                        v: Self::Element,
                                        s: &S,
                                        state: &Self::State)
                                        -> Result<f64> {
        unimplemented!()
    }

    /// The elements on which the marginal gain of `u` depends. The technical definition I have
    /// been using is based on the primal curvature (`nabla`): if `u` depends on `v`, then
    /// `nabla(u, v, sol, state) != 1`.
    fn depends(&self, u: Self::Element, state: &Self::State) -> Result<ElementIterator<Self>>;
    /// Same as `depends`, except that `nabla(u, v, sol, state) > 1`.
    ///
    /// Not all objectives need to implement this. If it is not implemented, calling it panics (via
    /// `unimplemented!()`)
    #[allow(unused_variables)]
    fn supermodular_depends(&self,
                            u: Self::Element,
                            state: &Self::State)
                            -> Result<ElementIterator<Self>> {
        unimplemented!()
    }

    /// Update the state tracked in `s` to reflect the insertion of `u` into the solution.
    fn insert_mut(&self, u: Self::Element, s: &mut Self::State) -> Result<()>;

    /// Produce a new state reflecting the insertion of `u` into the solution.
    fn insert(&self, u: Self::Element, s: &Self::State) -> Result<Self::State> {
        let mut state = s.clone();
        self.insert_mut(u, &mut state)?;
        Ok(state)
    }
}

/// Lazy extension of `Objective`. While the same state type is used, state updates are handled
/// independently. If you need two different kinds of state for each kind of update, consider using
/// an enum-struct. For example:
///
/// ```
/// #[derive(Clone, Debug)]
/// pub enum State {
///     Empty,
///     Strict {
///         foo: usize,
///         bar: Vec<usize>,
///     },
///     Lazy {
///         baz: f64,
///     }
/// }
///
/// impl Default for State {
///     fn default() -> Self {
///         State::Empty
///     }
/// }
/// ```
pub trait LazyObjective: Objective {
    /// Updates the marginal gain of `element`, returning it value after updating the state to
    /// reflect the insertion of `previous` elements.
    ///
    /// Returns `Ok(None)` if the marginal gain does not need to be updated.
    fn update_lazy_mut<S: Setlike<Self::Element>>(&self,
                                                  element: Self::Element,
                                                  previous: &S,
                                                  state: &mut Self::State)
                                                  -> Result<Option<f64>>;

    /// Update the state to preserve invariants upon insertion of `element`.
    ///
    /// This should do the least amount of work possible without making `update_lazy_mut` too
    /// complex.
    fn insert_lazy_mut(&self, element: Self::Element, state: &mut Self::State) -> Result<()>;
}

/// An objective with further constraints than mere cardinality.
pub trait ConstrainedObjective: Objective {
    /// Returns true if adding `el` to `sol` would create a valid set.
    fn valid_addition(&self,
                      el: Self::Element,
                      sol: &Vec<Self::Element>,
                      state: &Self::State)
                      -> bool;
}
<|endoftext|>


<commit_before>#![warn(rust_2018_idioms)]

#[macro_use]
extern crate slog;

use slog::Drain;
use smithay::{
    backend::drm::{
        common::Error,
        device_bind,
        legacy::{LegacyDrmDevice, LegacyDrmSurface},
        Device, DeviceHandler, RawSurface, Surface,
    },
    reexports::{
        calloop::EventLoop,
        drm::{
            buffer::format::PixelFormat,
            control::{
                connector::State as ConnectorState, crtc, dumbbuffer::DumbBuffer, framebuffer,
                Device as ControlDevice,
            },
        },
    },
};
use std::{
    fs::{File, OpenOptions},
    io::Error as IoError,
    rc::Rc,
    sync::Mutex,
};

fn main() {
    let log = slog::Logger::root(Mutex::new(slog_term::term_full().fuse()).fuse(), o!());

    /*
     * Initialize the drm backend
     */

    // "Find" a suitable drm device
    let mut options = OpenOptions::new();
    options.read(true);
    options.write(true);
    let mut device = LegacyDrmDevice::new(options.open("/dev/dri/card0").unwrap(), log.clone()).unwrap();

    // Get a set of all modesetting resource handles (excluding planes):
    let res_handles = Device::resource_handles(&device).unwrap();

    // Use first connected connector
    let connector_info = res_handles
        .connectors()
        .iter()
        .map(|conn| device.get_connector_info(*conn).unwrap())
        .find(|conn| conn.state() == ConnectorState::Connected)
        .unwrap();

    // Use the first encoder
    let encoder = connector_info
        .encoders()
        .iter()
        .filter_map(|&e| e)
        .next()
        .unwrap();
    let encoder_info = device.get_encoder_info(encoder).unwrap();

    // use the connected crtc if any
    let crtc = encoder_info
        .crtc()
        // or use the first one that is compatible with the encoder
        .unwrap_or_else(|| {
            *res_handles
                .filter_crtcs(encoder_info.possible_crtcs())
                .iter()
                .next()
                .unwrap()
        });

    // Assuming we found a good connector and loaded the info into `connector_info`
    let mode = connector_info.modes()[0]; // Use first mode (usually highest resoltion, but in reality you should filter and sort and check and match with other connectors, if you use more then one.)

    // Initialize the hardware backend
    let surface = Rc::new(device.create_surface(crtc).unwrap());

    surface.use_mode(Some(mode)).unwrap();
    for conn in surface.current_connectors().into_iter() {
        if conn != connector_info.handle() {
            surface.remove_connector(conn).unwrap();
        }
    }
    surface.add_connector(connector_info.handle()).unwrap();

    /*
     * Lets create buffers and framebuffers.
     * We use drm-rs DumbBuffers, because they always work and require little to no setup.
     * But they are very slow, this is just for demonstration purposes.
     */
    let (w, h) = mode.size();
    let front_buffer = device
        .create_dumb_buffer((w as u32, h as u32), PixelFormat::XRGB8888)
        .unwrap();
    let front_framebuffer = device.add_framebuffer(&front_buffer).unwrap();
    let back_buffer = device
        .create_dumb_buffer((w as u32, h as u32), PixelFormat::XRGB8888)
        .unwrap();
    let back_framebuffer = device.add_framebuffer(&back_buffer).unwrap();

    device.set_handler(DrmHandlerImpl {
        current: front_framebuffer,
        front: (front_buffer, front_framebuffer),
        back: (back_buffer, back_framebuffer),
        surface: surface.clone(),
    });

    /*
     * Register the DrmDevice on the EventLoop
     */
    let mut event_loop = EventLoop::<()>::new().unwrap();
    let _source = device_bind(&event_loop.handle(), device)
        .map_err(|err| -> IoError { err.into() })
        .unwrap();

    // Start rendering
    if surface.commit_pending() {
        surface.commit(front_framebuffer).unwrap();
    }

    // Run
    event_loop.run(None, &mut (), |_| {}).unwrap();
}

pub struct DrmHandlerImpl {
    front: (DumbBuffer, framebuffer::Handle),
    back: (DumbBuffer, framebuffer::Handle),
    current: framebuffer::Handle,
    surface: Rc<LegacyDrmSurface<File>>,
}

impl DeviceHandler for DrmHandlerImpl {
    type Device = LegacyDrmDevice<File>;

    fn vblank(&mut self, _crtc: crtc::Handle) {
        {
            // Swap and map buffer
            let mut mapping = if self.current == self.front.1 {
                self.current = self.back.1;
                self.surface.map_dumb_buffer(&mut self.back.0).unwrap()
            } else {
                self.current = self.front.1;
                self.surface.map_dumb_buffer(&mut self.front.0).unwrap()
            };

            // now we could render to the mapping via software rendering.
            // this example just sets some grey color

            for x in mapping.as_mut() {
                *x = 128;
            }
        }
        RawSurface::page_flip(&*self.surface, self.current).unwrap();
    }

    fn error(&mut self, error: Error) {
        panic!("{:?}", error);
    }
}
<commit_msg>example: use set_connectors in raw_legacy_drm<commit_after>#![warn(rust_2018_idioms)]

#[macro_use]
extern crate slog;

use slog::Drain;
use smithay::{
    backend::drm::{
        common::Error,
        device_bind,
        legacy::{LegacyDrmDevice, LegacyDrmSurface},
        Device, DeviceHandler, RawSurface, Surface,
    },
    reexports::{
        calloop::EventLoop,
        drm::{
            buffer::format::PixelFormat,
            control::{
                connector::State as ConnectorState, crtc, dumbbuffer::DumbBuffer, framebuffer,
                Device as ControlDevice,
            },
        },
    },
};
use std::{
    fs::{File, OpenOptions},
    io::Error as IoError,
    rc::Rc,
    sync::Mutex,
};

fn main() {
    let log = slog::Logger::root(Mutex::new(slog_term::term_full().fuse()).fuse(), o!());

    /*
     * Initialize the drm backend
     */

    // "Find" a suitable drm device
    let mut options = OpenOptions::new();
    options.read(true);
    options.write(true);
    let mut device = LegacyDrmDevice::new(options.open("/dev/dri/card0").unwrap(), log.clone()).unwrap();

    // Get a set of all modesetting resource handles (excluding planes):
    let res_handles = Device::resource_handles(&device).unwrap();

    // Use first connected connector
    let connector_info = res_handles
        .connectors()
        .iter()
        .map(|conn| device.get_connector_info(*conn).unwrap())
        .find(|conn| conn.state() == ConnectorState::Connected)
        .unwrap();

    // Use the first encoder
    let encoder = connector_info
        .encoders()
        .iter()
        .filter_map(|&e| e)
        .next()
        .unwrap();
    let encoder_info = device.get_encoder_info(encoder).unwrap();

    // use the connected crtc if any
    let crtc = encoder_info
        .crtc()
        // or use the first one that is compatible with the encoder
        .unwrap_or_else(|| {
            *res_handles
                .filter_crtcs(encoder_info.possible_crtcs())
                .iter()
                .next()
                .unwrap()
        });

    // Assuming we found a good connector and loaded the info into `connector_info`
    let mode = connector_info.modes()[0]; // Use first mode (usually highest resoltion, but in reality you should filter and sort and check and match with other connectors, if you use more then one.)

    // Initialize the hardware backend
    let surface = Rc::new(device.create_surface(crtc).unwrap());

    surface.use_mode(Some(mode)).unwrap();
    surface.set_connectors(&[connector_info.handle()]).unwrap();
    /*
     * Lets create buffers and framebuffers.
     * We use drm-rs DumbBuffers, because they always work and require little to no setup.
     * But they are very slow, this is just for demonstration purposes.
     */
    let (w, h) = mode.size();
    let front_buffer = device
        .create_dumb_buffer((w as u32, h as u32), PixelFormat::XRGB8888)
        .unwrap();
    let front_framebuffer = device.add_framebuffer(&front_buffer).unwrap();
    let back_buffer = device
        .create_dumb_buffer((w as u32, h as u32), PixelFormat::XRGB8888)
        .unwrap();
    let back_framebuffer = device.add_framebuffer(&back_buffer).unwrap();

    device.set_handler(DrmHandlerImpl {
        current: front_framebuffer,
        front: (front_buffer, front_framebuffer),
        back: (back_buffer, back_framebuffer),
        surface: surface.clone(),
    });

    /*
     * Register the DrmDevice on the EventLoop
     */
    let mut event_loop = EventLoop::<()>::new().unwrap();
    let _source = device_bind(&event_loop.handle(), device)
        .map_err(|err| -> IoError { err.into() })
        .unwrap();

    // Start rendering
    if surface.commit_pending() {
        surface.commit(front_framebuffer).unwrap();
    }

    // Run
    event_loop.run(None, &mut (), |_| {}).unwrap();
}

pub struct DrmHandlerImpl {
    front: (DumbBuffer, framebuffer::Handle),
    back: (DumbBuffer, framebuffer::Handle),
    current: framebuffer::Handle,
    surface: Rc<LegacyDrmSurface<File>>,
}

impl DeviceHandler for DrmHandlerImpl {
    type Device = LegacyDrmDevice<File>;

    fn vblank(&mut self, _crtc: crtc::Handle) {
        {
            // Swap and map buffer
            let mut mapping = if self.current == self.front.1 {
                self.current = self.back.1;
                self.surface.map_dumb_buffer(&mut self.back.0).unwrap()
            } else {
                self.current = self.front.1;
                self.surface.map_dumb_buffer(&mut self.front.0).unwrap()
            };

            // now we could render to the mapping via software rendering.
            // this example just sets some grey color

            for x in mapping.as_mut() {
                *x = 128;
            }
        }
        RawSurface::page_flip(&*self.surface, self.current).unwrap();
    }

    fn error(&mut self, error: Error) {
        panic!("{:?}", error);
    }
}
<|endoftext|>


<commit_before>#[cfg(not(feature = "rustc-dep-of-std"))]
extern crate alloc;

use crate::{
	block_current_task, get_priority, getpid, set_priority, wakeup_task, yield_now, Priority, Tid,
	NO_PRIORITIES,
};
use alloc::collections::vec_deque::VecDeque;
use core::cell::UnsafeCell;
use core::hint;
use core::ops::{Deref, DerefMut, Drop};
use core::sync::atomic::{AtomicUsize, Ordering};

/// This type provides a lock based on busy waiting to realize mutual exclusion
///
/// # Description
///
/// This structure behaves a lot like a common mutex. There are some differences:
///
/// - By using busy waiting, it can be used outside the runtime.
/// - It is a so called ticket lock and is completely fair.
#[cfg_attr(target_arch = "x86_64", repr(align(128)))]
#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))]
struct Spinlock<T: ?Sized> {
	queue: AtomicUsize,
	dequeue: AtomicUsize,
	data: UnsafeCell<T>,
}

unsafe impl<T: ?Sized + Send> Sync for Spinlock<T> {}
unsafe impl<T: ?Sized + Send> Send for Spinlock<T> {}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
struct SpinlockGuard<'a, T: ?Sized> {
	dequeue: &'a AtomicUsize,
	data: &'a mut T,
}

impl<T> Spinlock<T> {
	pub const fn new(user_data: T) -> Spinlock<T> {
		Spinlock {
			queue: AtomicUsize::new(0),
			dequeue: AtomicUsize::new(1),
			data: UnsafeCell::new(user_data),
		}
	}

	#[inline]
	fn obtain_lock(&self) {
		let ticket = self.queue.fetch_add(1, Ordering::SeqCst) + 1;
		let mut counter: u16 = 0;
		while self.dequeue.load(Ordering::SeqCst) != ticket {
			counter += 1;
			if counter < 100 {
				hint::spin_loop();
			} else {
				counter = 0;
				unsafe {
					yield_now();
				}
			}
		}
	}

	#[inline]
	pub unsafe fn lock(&self) -> SpinlockGuard<'_, T> {
		self.obtain_lock();
		SpinlockGuard {
			dequeue: &self.dequeue,
			data: &mut *self.data.get(),
		}
	}
}

impl<T: ?Sized + Default> Default for Spinlock<T> {
	fn default() -> Spinlock<T> {
		Spinlock::new(Default::default())
	}
}

impl<'a, T: ?Sized> Deref for SpinlockGuard<'a, T> {
	type Target = T;
	fn deref(&self) -> &T {
		&*self.data
	}
}

impl<'a, T: ?Sized> DerefMut for SpinlockGuard<'a, T> {
	fn deref_mut(&mut self) -> &mut T {
		&mut *self.data
	}
}

impl<'a, T: ?Sized> Drop for SpinlockGuard<'a, T> {
	/// The dropping of the SpinlockGuard will release the lock it was created from.
	fn drop(&mut self) {
		self.dequeue.fetch_add(1, Ordering::SeqCst);
	}
}

/// Realize a priority queue for tasks
struct PriorityQueue {
	queues: [Option<VecDeque<Tid>>; NO_PRIORITIES],
	prio_bitmap: u64,
}

impl PriorityQueue {
	pub const fn new() -> PriorityQueue {
		PriorityQueue {
			queues: [
				None, None, None, None, None, None, None, None, None, None, None, None, None, None,
				None, None, None, None, None, None, None, None, None, None, None, None, None, None,
				None, None, None,
			],
			prio_bitmap: 0,
		}
	}

	/// Add a task id by its priority to the queue
	pub fn push(&mut self, prio: Priority, id: Tid) {
		let i: usize = prio.into().into();
		self.prio_bitmap |= (1 << i) as u64;
		if let Some(queue) = &mut self.queues[i] {
			queue.push_back(id);
		} else {
			let mut queue = VecDeque::new();
			queue.push_back(id);
			self.queues[i] = Some(queue);
		}
	}

	fn pop_from_queue(&mut self, queue_index: usize) -> Option<Tid> {
		if let Some(queue) = &mut self.queues[queue_index] {
			let id = queue.pop_front();

			if queue.is_empty() {
				self.prio_bitmap &= !(1 << queue_index as u64);
			}

			id
		} else {
			None
		}
	}

	/// Pop the task handle with the highest priority from the queue
	pub fn pop(&mut self) -> Option<Tid> {
		for i in 0..NO_PRIORITIES {
			if self.prio_bitmap & (1 << i) != 0 {
				return self.pop_from_queue(i);
			}
		}

		None
	}
}

struct MutexInner {
	current_prio: Priority,
	base_prio: Priority,
	id: Tid,
	blocked_task: PriorityQueue,
}

impl MutexInner {
	pub const fn new() -> MutexInner {
		MutexInner {
			current_prio: Priority::from(0),
			base_prio: Priority::from(0),
			id: 0,
			blocked_task: PriorityQueue::new(),
		}
	}
}

pub struct Mutex {
	inner: Spinlock<MutexInner>,
}

unsafe impl Send for Mutex {}
unsafe impl Sync for Mutex {}

impl Mutex {
	pub const fn new() -> Mutex {
		Mutex {
			inner: Spinlock::new(MutexInner::new()),
		}
	}

	#[inline]
	pub unsafe fn init(&mut self) {
		self.inner = Spinlock::new(MutexInner::new());
	}

	#[inline]
	pub unsafe fn lock(&self) {
		loop {
			let mut guard = self.inner.lock();
			if guard.id == 0 {
				guard.current_prio = get_priority();
				guard.base_prio = guard.current_prio;
				guard.id = getpid();
				return;
			} else {
				let prio = get_priority();

				if guard.current_prio < prio {
					set_priority(guard.id, prio);
					guard.current_prio = prio;
				}
				guard.blocked_task.push(prio, getpid());
				block_current_task();
				drop(guard);
				yield_now();
			}
		}
	}

	#[inline]
	pub unsafe fn unlock(&self) {
		let mut guard = self.inner.lock();

		if guard.base_prio != guard.current_prio {
			set_priority(guard.id, guard.base_prio);
		}

		// reset data
		guard.current_prio = Priority::from(0);
		guard.base_prio = Priority::from(0);
		guard.id = 0;

		if let Some(tid) = guard.blocked_task.pop() {
			wakeup_task(tid);
		}
	}

	#[inline]
	pub unsafe fn try_lock(&self) -> bool {
		let mut guard = self.inner.lock();
		let locked = guard.id == 0;
		if locked {
			guard.current_prio = get_priority();
			guard.base_prio = guard.current_prio;
			guard.id = getpid();
		}
		locked
	}

	#[inline]
	pub unsafe fn destroy(&self) {}
}
<commit_msg>using `Option<Tid>` to identify if the mutex is hold by a task<commit_after>#[cfg(not(feature = "rustc-dep-of-std"))]
extern crate alloc;

use crate::{
	block_current_task, get_priority, getpid, set_priority, wakeup_task, yield_now, Priority, Tid,
	NO_PRIORITIES,
};
use alloc::collections::vec_deque::VecDeque;
use core::cell::UnsafeCell;
use core::hint;
use core::ops::{Deref, DerefMut, Drop};
use core::sync::atomic::{AtomicUsize, Ordering};

/// This type provides a lock based on busy waiting to realize mutual exclusion
///
/// # Description
///
/// This structure behaves a lot like a common mutex. There are some differences:
///
/// - By using busy waiting, it can be used outside the runtime.
/// - It is a so called ticket lock and is completely fair.
#[cfg_attr(target_arch = "x86_64", repr(align(128)))]
#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))]
struct Spinlock<T: ?Sized> {
	queue: AtomicUsize,
	dequeue: AtomicUsize,
	data: UnsafeCell<T>,
}

unsafe impl<T: ?Sized + Send> Sync for Spinlock<T> {}
unsafe impl<T: ?Sized + Send> Send for Spinlock<T> {}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
struct SpinlockGuard<'a, T: ?Sized> {
	dequeue: &'a AtomicUsize,
	data: &'a mut T,
}

impl<T> Spinlock<T> {
	pub const fn new(user_data: T) -> Spinlock<T> {
		Spinlock {
			queue: AtomicUsize::new(0),
			dequeue: AtomicUsize::new(1),
			data: UnsafeCell::new(user_data),
		}
	}

	#[inline]
	fn obtain_lock(&self) {
		let ticket = self.queue.fetch_add(1, Ordering::SeqCst) + 1;
		let mut counter: u16 = 0;
		while self.dequeue.load(Ordering::SeqCst) != ticket {
			counter += 1;
			if counter < 100 {
				hint::spin_loop();
			} else {
				counter = 0;
				unsafe {
					yield_now();
				}
			}
		}
	}

	#[inline]
	pub unsafe fn lock(&self) -> SpinlockGuard<'_, T> {
		self.obtain_lock();
		SpinlockGuard {
			dequeue: &self.dequeue,
			data: &mut *self.data.get(),
		}
	}
}

impl<T: ?Sized + Default> Default for Spinlock<T> {
	fn default() -> Spinlock<T> {
		Spinlock::new(Default::default())
	}
}

impl<'a, T: ?Sized> Deref for SpinlockGuard<'a, T> {
	type Target = T;
	fn deref(&self) -> &T {
		&*self.data
	}
}

impl<'a, T: ?Sized> DerefMut for SpinlockGuard<'a, T> {
	fn deref_mut(&mut self) -> &mut T {
		&mut *self.data
	}
}

impl<'a, T: ?Sized> Drop for SpinlockGuard<'a, T> {
	/// The dropping of the SpinlockGuard will release the lock it was created from.
	fn drop(&mut self) {
		self.dequeue.fetch_add(1, Ordering::SeqCst);
	}
}

/// Realize a priority queue for tasks
struct PriorityQueue {
	queues: [Option<VecDeque<Tid>>; NO_PRIORITIES],
	prio_bitmap: u64,
}

impl PriorityQueue {
	pub const fn new() -> PriorityQueue {
		PriorityQueue {
			queues: [
				None, None, None, None, None, None, None, None, None, None, None, None, None, None,
				None, None, None, None, None, None, None, None, None, None, None, None, None, None,
				None, None, None,
			],
			prio_bitmap: 0,
		}
	}

	/// Add a task id by its priority to the queue
	pub fn push(&mut self, prio: Priority, id: Tid) {
		let i: usize = prio.into().into();
		self.prio_bitmap |= (1 << i) as u64;
		if let Some(queue) = &mut self.queues[i] {
			queue.push_back(id);
		} else {
			let mut queue = VecDeque::new();
			queue.push_back(id);
			self.queues[i] = Some(queue);
		}
	}

	fn pop_from_queue(&mut self, queue_index: usize) -> Option<Tid> {
		if let Some(queue) = &mut self.queues[queue_index] {
			let id = queue.pop_front();

			if queue.is_empty() {
				self.prio_bitmap &= !(1 << queue_index as u64);
			}

			id
		} else {
			None
		}
	}

	/// Pop the task handle with the highest priority from the queue
	pub fn pop(&mut self) -> Option<Tid> {
		for i in 0..NO_PRIORITIES {
			if self.prio_bitmap & (1 << i) != 0 {
				return self.pop_from_queue(i);
			}
		}

		None
	}
}

struct MutexInner {
	/// The lock is free, if `id` is none. Otherwise `id` identifies the task.
	id: Option<Tid>,
	/// Current priority of the task, which holds the lock
	current_prio: Priority,
	/// Original priority of the task, which holds the lock
	base_prio: Priority,
	/// Priority queue of blocked tasks
	blocked_task: PriorityQueue,
}

impl MutexInner {
	pub const fn new() -> MutexInner {
		MutexInner {
			id: None,
			current_prio: Priority::from(0),
			base_prio: Priority::from(0),
			blocked_task: PriorityQueue::new(),
		}
	}
}

pub struct Mutex {
	inner: Spinlock<MutexInner>,
}

unsafe impl Send for Mutex {}
unsafe impl Sync for Mutex {}

impl Mutex {
	pub const fn new() -> Mutex {
		Mutex {
			inner: Spinlock::new(MutexInner::new()),
		}
	}

	#[inline]
	pub unsafe fn init(&mut self) {
		self.inner = Spinlock::new(MutexInner::new());
	}

	#[inline]
	pub unsafe fn lock(&self) {
		loop {
			let mut guard = self.inner.lock();
			match guard.id {
				None => {
					guard.current_prio = get_priority();
					guard.base_prio = guard.current_prio;
					guard.id = Some(getpid());
					return;
				},
				Some(id) => {
					let prio = get_priority();

					if guard.current_prio < prio {
						set_priority(id, prio);
						guard.current_prio = prio;
					}
					guard.blocked_task.push(prio, getpid());
					block_current_task();
					drop(guard);
					yield_now();
				},
			}
		}
	}

	#[inline]
	pub unsafe fn unlock(&self) {
		let mut guard = self.inner.lock();

		if guard.base_prio != guard.current_prio {
			if let Some(id) = guard.id {
				set_priority(id, guard.base_prio);
			}
		}

		// reset data
		guard.current_prio = Priority::from(0);
		guard.base_prio = Priority::from(0);
		guard.id = None;

		if let Some(tid) = guard.blocked_task.pop() {
			wakeup_task(tid);
		}
	}

	#[inline]
	pub unsafe fn try_lock(&self) -> bool {
		let mut guard = self.inner.lock();
		if guard.id.is_none() {
			guard.current_prio = get_priority();
			guard.base_prio = guard.current_prio;
			guard.id = Some(getpid());

			true
		} else {
			false
		}
	}

	#[inline]
	pub unsafe fn destroy(&self) {}
}
<|endoftext|>