<commit_before>use std::io::{BufWriter, Reader};
use std::num::Float;
use error::Ch8Error;
use ops::{Op, Instruction};
use std::slice::Chunks;

const RAM_SIZE: usize = 4096;
const PROGRAM_START: usize = 0x200;
const CLOCK_SPEED_HZ: u32 = 512 * 1024; // 512KHz
const STEP_TIME: f32 = 1.0f32 / CLOCK_SPEED_HZ as f32;

pub struct Vm {
    reg: [u8; 16],
    i: usize,
    pc: usize,
    sp: usize,
    stack: [usize; 256],
    ram: [u8; RAM_SIZE],

    timer: u8,
    t_tick: f32,

    sound_timer: u8,
    st_tick: f32,

    screen: [u8; 64 * 32],
    keys: [u8; 16],
}

impl Vm {
    pub fn new() -> Vm {
        Vm {
            reg: [0; 16],
            i: 0,
            pc: PROGRAM_START,
            sp: 0,
            stack: [0; 256],
            ram: [0; RAM_SIZE],

            timer: 0,
            t_tick: 0.0,

            sound_timer: 0,
            st_tick: 0.0,

            screen: [0; 64 * 32],
            keys: [0; 16],
        }
    }

    pub fn load_rom(&mut self, reader: &mut Reader) -> Result<usize, Ch8Error> {
        let rom = try!(reader.read_to_end());
        if rom.len() > (RAM_SIZE - PROGRAM_START) {
           return Err(Ch8Error::Io("Rom was larger than available RAM", None))
        }
        let mut ram = BufWriter::new(&mut self.ram[PROGRAM_START..RAM_SIZE]);
        try!(ram.write(rom.as_slice()));
        return Ok(rom.len());
    }

    pub fn dump_ram(&self, writer: &mut Writer) {
        writer.write(&self.ram).unwrap();
    }

    pub fn set_key(&mut self, idx: u8) {
        self.keys[idx as usize] = 1;
    }

    pub fn unset_key(&mut self, idx: u8) {
        self.keys[idx as usize] = 0;
    }

    fn exec(&mut self, op: &Op) -> bool {
        use ops::Instruction::*;
        let ins = Instruction::from_op(op);
//        if self.timer == 0 {
//            println!("Executing instruction: 0x{:X} {:?}", self.pc, ins);
//            println!("v: {:?}", self.reg);
//            println!("timers: t: {} s: {}", self.timer, self.sound_timer);
//            self.print_screen();
//        }
        match ins {
            // Sys(addr) intentionally left unimplemented.
            Clear => {
                for b in self.screen.iter_mut() {
                    *b = 0;
                }
            },
            Return => {
                self.pc = self.stack[self.sp];
                self.sp-=1;
            },
            Jump(addr) => {
                let idle = self.pc-2 == addr as usize;
                self.pc = addr as usize;
                if idle { return true; }
            }
            Call(addr) => {
                self.sp+=1;
                self.stack[self.sp] = self.pc;
                self.pc = addr as usize;
            },
            SkipEqualK(vx, k) => {
                if self.reg[vx as usize] == k {
                    self.pc += 2;
                }
            },
            SkipNotEqualK(vx, k) => {
                if self.reg[vx as usize] != k {
                    self.pc += 2;
                }
            },
            SkipEqual(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];
                if x == y {
                    self.pc += 2;
                }
            },
            SetK(vx, byte) => {
                self.reg[vx as usize] = byte;
            },
            AddK(vx, byte) => {
                self.reg[vx as usize] += byte;
            },
            Set(vx, vy) => self.reg[vx as usize] = self.reg[vy as usize],
            Or(vx, vy)  => self.reg[vx as usize] |= self.reg[vy as usize],
            And(vx, vy) => self.reg[vx as usize] &= self.reg[vy as usize],
            XOr(vx, vy) => self.reg[vx as usize] ^= self.reg[vy as usize],
            Add(vx, vy) => {
                let x = self.reg[vx as usize] as u16;
                let y = self.reg[vy as usize] as u16;
                let res = x + y;

                // VF is carryover
                self.reg[15] = (res > 255) as u8;

                self.reg[vx as usize] = res as u8;
            },
            Sub(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];

                // VF is Not Borrow i.e. x > y
                self.reg[15] = (x > y) as u8;

                self.reg[vx as usize] = x - y;
            },
            ShiftRight(vx, vy) => {
                let y = self.reg[vy as usize];

                // VF is lsb before shift
                self.reg[15] = 0x1 & y;

                self.reg[vx as usize] = y >> 1;
            },
            SubInv(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];

                // VF is Not Borrow i.e. y > x
                self.reg[15] = (y > x) as u8;

                self.reg[vx as usize] = y - x;
            },
            ShiftLeft(vx, vy) => {
                let y = self.reg[vy as usize];

                // VF is msb before shift
                self.reg[15] =  0x80 & y;

                self.reg[vx as usize] = y << 1;
            }
            SkipNotEqual(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];
                if x != y {
                    self.pc += 2;
                }
            },
            LoadI(addr) => {
                self.i = addr as usize;
            },
            LongJump(addr) => {
                self.pc = (self.reg[0] as u16 + addr) as usize;
            },
            //TODO Rand(byte) => {}
            Draw(vx, vy, n) => {
                let x = self.reg[vx as usize] as usize;
                let y = self.reg[vy as usize] as usize;
                let i = self.i as usize;
                let n = n as usize;

                let sprite = &self.ram[i..i+n];

                self.reg[15] = 0;
                for (sy, byte) in sprite.iter().enumerate() {
                    let dy = (y + sy) % 32;
                    for sx in 0us..8 {
                        let px = (*byte >> (7 - sx)) & 0b00000001;
                        let dx = (x + sx) % 64;
                        let idx = dy * 64 + dx;
                        self.screen[idx] ^= px;

                        // Vf is if there was a collision
                        self.reg[15] |= (self.screen[idx] == 0 && px == 1) as u8;
                    }
                }
            },
            //TODO SkipPressed(vx) => {}
            //TODO SkipNotPressed(vx) => {}
            SkipPressed(vx) => {
                let idx = self.reg[vx as usize];
                if self.keys[idx as usize] == 1 {
                    self.pc += 2;
                }
            }
            SkipNotPressed(vx) => {
                let idx = self.reg[vx as usize];
                if self.keys[idx  as usize] != 1 {
                    self.pc += 2;
                }
            }
            GetTimer(vx) => {
                self.reg[vx as usize] = self.timer;
            },
            //TODO WaitKey(vx) => {}
            SetTimer(vx) => {
                self.timer = self.reg[vx as usize];
                self.t_tick = 1.0 / 60.0;
            },
            SetSoundTimer(vx) => {
                self.sound_timer = self.reg[vx as usize];
                self.st_tick = 1.0 / 60.0;
            },
            AddToI(vx) => {
                self.i += self.reg[vx as usize] as usize;
            },
            //TODO LoadHexGlyph(vx) => {}
            //TODO StoreBCD(vx) => {}
            StoreRegisters(vx) => {
                let vx = vx as usize;
                let i = self.i as usize;

                let mut dst = &mut self.ram[i..i+vx+1];
                for (x,b) in dst.iter_mut().enumerate() {
                  //println!("store reg {} to {}", x, i);
                    *b = self.reg[x];
                }
            },
            LoadRegisters(vx) => {
                let vx = vx as usize;
                let i = self.i as usize;

                let src = &self.ram[i..i+vx+1];
                for (x,b) in src.iter().enumerate() {
                  //println!("read reg {} from {}", x, i);
                    self.reg[x] = *b;
                }
            },
            other => {
                println!("Instruction not implemented {:?} skipping...", other)
            }
        }
        return false;
    }

    fn time_step(&mut self) {
        if self.timer > 0 {
            self.t_tick -= STEP_TIME;
            if self.t_tick <= 0.0 {
                self.timer -= 1;
                self.t_tick = 1.0 / 60.0;
            }
        }

        if self.sound_timer > 0 {
            self.st_tick -= STEP_TIME;
            if self.st_tick <= 0.0 {
                self.sound_timer -= 1;
                self.st_tick = 1.0 / 60.0;
            }
        }
    }

    // dt: Time in seconds since last step
    pub fn step(&mut self, dt:f32) -> bool {
        let steps = (CLOCK_SPEED_HZ as f32 * dt).round() as u32;
        let mut idle = false;

        for _ in 0u32..steps {
            let raw = {
                let codes = &self.ram[self.pc..self.pc+2];
                ((codes[0] as u16) << 8) | codes[1] as u16
            };
            let op = Op::new(raw);
            self.pc += 2;
            idle = self.exec(&op);
            self.time_step();
            if idle { break; }
        }
        return idle;
    }

    pub fn screen_rows<'a>(&'a self) -> Chunks<'a, u8> {
        self.screen.chunks(64)
    }

    pub fn print_screen(&self) {
        for row in self.screen.chunks(64) {
            println!("");
            for byte in row.iter() {
                match *byte {
                    0x0 => print!("░"),
                    _ => print!("▓")
                }
            }
        }
    }

    pub fn print_disassembly(&self) {
        for i in self.ram.chunks(2) {
            match i {
                [0, 0] => continue,
                [h, l] => {
                    let op = Op::new(((h as u16) << 8) | l as u16);
                    println!("raw: 0x{:X}", op.raw());
                    println!("instruction: {:?}", Instruction::from_op(&op));
                    println!("addr: 0x{:X}", op.addr());
                    println!("x: 0x{:X}", op.x());
                    println!("y: 0x{:X}", op.y());
                    println!("n_high: 0x{:X}", op.n_high());
                    println!("n_low: 0x{:X}", op.n_low());
                    println!("k: 0x{:X}\n", op.k());
                },
                _ => continue
            }
        }
    }
}
<commit_msg>Add font handling<commit_after>use std::io::{BufWriter, Reader};
use std::num::Float;
use error::Ch8Error;
use ops::{Op, Instruction};
use std::slice::Chunks;

use std::rand::Rng;
use std::rand;

const RAM_SIZE: usize = 4096;
const PROGRAM_START: usize = 0x200;
const CLOCK_SPEED_HZ: u32 = 512 * 1024; // 512KHz
const STEP_TIME: f32 = 1.0f32 / CLOCK_SPEED_HZ as f32;

const FONT_ADDR: usize = 0;
const FONT_HEIGHT: usize = 5;
const FONT_BYTES: usize = FONT_HEIGHT * 16;
const FONT: [u8; FONT_BYTES] = [
	0xF0, 0x90, 0x90, 0x90, 0xF0, // 0
	0x20, 0x60, 0x20, 0x20, 0x70, // 1
	0xF0, 0x10, 0xF0, 0x80, 0xF0, // 2
	0xF0, 0x10, 0xF0, 0x10, 0xF0, // 3
	0x90, 0x90, 0xF0, 0x10, 0x10, // 4
	0xF0, 0x80, 0xF0, 0x10, 0xF0, // 5
	0xF0, 0x80, 0xF0, 0x90, 0xF0, // 6
	0xF0, 0x10, 0x20, 0x40, 0x40, // 7
	0xF0, 0x90, 0xF0, 0x90, 0xF0, // 8
	0xF0, 0x90, 0xF0, 0x10, 0xF0, // 9
	0xF0, 0x90, 0xF0, 0x90, 0x90, // A
	0xE0, 0x90, 0xE0, 0x90, 0xE0, // B
	0xF0, 0x80, 0x80, 0x80, 0xF0, // C
	0xE0, 0x90, 0x90, 0x90, 0xE0, // D
	0xF0, 0x80, 0xF0, 0x80, 0xF0, // E
	0xF0, 0x80, 0xF0, 0x80, 0x80, // F
];

pub struct Vm {
    reg: [u8; 16],
    i: usize,
    pc: usize,
    sp: usize,
    stack: [usize; 256],
    ram: [u8; RAM_SIZE],

    timer: u8,
    t_tick: f32,

    sound_timer: u8,
    st_tick: f32,

    screen: [u8; 64 * 32],
    keys: [u8; 16],
}

impl Vm {
    pub fn new() -> Vm {
        let mut vm = Vm {
            reg: [0; 16],
            i: 0,
            pc: PROGRAM_START,
            sp: 0,
            stack: [0; 256],
            ram: [0; RAM_SIZE],

            timer: 0,
            t_tick: 0.0,

            sound_timer: 0,
            st_tick: 0.0,

            screen: [0; 64 * 32],
            keys: [0; 16],
        };
        {
            let mut ram = BufWriter::new(&mut vm.ram[FONT_ADDR..(FONT_ADDR + FONT_BYTES)]);
            ram.write(FONT.as_slice());
        }
        vm
    }

    pub fn load_rom(&mut self, reader: &mut Reader) -> Result<usize, Ch8Error> {
        let rom = try!(reader.read_to_end());
        if rom.len() > (RAM_SIZE - PROGRAM_START) {
           return Err(Ch8Error::Io("Rom was larger than available RAM", None))
        }
        let mut ram = BufWriter::new(&mut self.ram[PROGRAM_START..RAM_SIZE]);
        try!(ram.write(rom.as_slice()));
        return Ok(rom.len());
    }

    pub fn dump_ram(&self, writer: &mut Writer) {
        writer.write(&self.ram).unwrap();
    }

    pub fn set_key(&mut self, idx: u8) {
        self.keys[idx as usize] = 1;
    }

    pub fn unset_key(&mut self, idx: u8) {
        self.keys[idx as usize] = 0;
    }

    fn exec(&mut self, op: &Op) -> bool {
        use ops::Instruction::*;
        let ins = Instruction::from_op(op);
//        if self.timer == 0 {
//            println!("Executing instruction: 0x{:X} {:?}", self.pc, ins);
//            println!("v: {:?}", self.reg);
//            println!("timers: t: {} s: {}", self.timer, self.sound_timer);
//            self.print_screen();
//        }
        match ins {
            // Sys(addr) intentionally left unimplemented.
            Clear => {
                for b in self.screen.iter_mut() {
                    *b = 0;
                }
            },
            Return => {
                self.pc = self.stack[self.sp];
                self.sp-=1;
            },
            Jump(addr) => {
                let idle = self.pc-2 == addr as usize;
                self.pc = addr as usize;
                if idle { return true; }
            }
            Call(addr) => {
                self.sp+=1;
                self.stack[self.sp] = self.pc;
                self.pc = addr as usize;
            },
            SkipEqualK(vx, k) => {
                if self.reg[vx as usize] == k {
                    self.pc += 2;
                }
            },
            SkipNotEqualK(vx, k) => {
                if self.reg[vx as usize] != k {
                    self.pc += 2;
                }
            },
            SkipEqual(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];
                if x == y {
                    self.pc += 2;
                }
            },
            SetK(vx, byte) => {
                self.reg[vx as usize] = byte;
            },
            AddK(vx, byte) => {
                self.reg[vx as usize] += byte;
            },
            Set(vx, vy) => self.reg[vx as usize] = self.reg[vy as usize],
            Or(vx, vy)  => self.reg[vx as usize] |= self.reg[vy as usize],
            And(vx, vy) => self.reg[vx as usize] &= self.reg[vy as usize],
            XOr(vx, vy) => self.reg[vx as usize] ^= self.reg[vy as usize],
            Add(vx, vy) => {
                let x = self.reg[vx as usize] as u16;
                let y = self.reg[vy as usize] as u16;
                let res = x + y;

                // VF is carryover
                self.reg[15] = (res > 255) as u8;

                self.reg[vx as usize] = res as u8;
            },
            Sub(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];

                // VF is Not Borrow i.e. x > y
                self.reg[15] = (x > y) as u8;

                self.reg[vx as usize] = x - y;
            },
            ShiftRight(vx, vy) => {
                let y = self.reg[vy as usize];

                // VF is lsb before shift
                self.reg[15] = 0x1 & y;

                self.reg[vx as usize] = y >> 1;
            },
            SubInv(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];

                // VF is Not Borrow i.e. y > x
                self.reg[15] = (y > x) as u8;

                self.reg[vx as usize] = y - x;
            },
            ShiftLeft(vx, vy) => {
                let y = self.reg[vy as usize];

                // VF is msb before shift
                self.reg[15] =  0x80 & y;

                self.reg[vx as usize] = y << 1;
            }
            SkipNotEqual(vx, vy) => {
                let x = self.reg[vx as usize];
                let y = self.reg[vy as usize];
                if x != y {
                    self.pc += 2;
                }
            },
            LoadI(addr) => {
                self.i = addr as usize;
            },
            LongJump(addr) => {
                self.pc = (self.reg[0] as u16 + addr) as usize;
            },
            Rand(vx, byte) => {
                self.reg[vx as usize] = rand::thread_rng().gen::<u8>() & byte;
            }
            Draw(vx, vy, n) => {
                let x = self.reg[vx as usize] as usize;
                let y = self.reg[vy as usize] as usize;
                let i = self.i as usize;
                let n = n as usize;

                let sprite = &self.ram[i..i+n];

                self.reg[15] = 0;
                for (sy, byte) in sprite.iter().enumerate() {
                    let dy = (y + sy) % 32;
                    for sx in 0us..8 {
                        let px = (*byte >> (7 - sx)) & 0b00000001;
                        let dx = (x + sx) % 64;
                        let idx = dy * 64 + dx;
                        self.screen[idx] ^= px;

                        // Vf is if there was a collision
                        self.reg[15] |= (self.screen[idx] == 0 && px == 1) as u8;
                    }
                }
            },
            SkipPressed(vx) => {
                let idx = self.reg[vx as usize];
                if self.keys[idx as usize] == 1 {
                    self.pc += 2;
                }
            }
            SkipNotPressed(vx) => {
                let idx = self.reg[vx as usize];
                if self.keys[idx  as usize] != 1 {
                    self.pc += 2;
                }
            }
            GetTimer(vx) => {
                self.reg[vx as usize] = self.timer;
            },
            //TODO WaitKey(vx) => {}
            SetTimer(vx) => {
                self.timer = self.reg[vx as usize];
                self.t_tick = 1.0 / 60.0;
            },
            SetSoundTimer(vx) => {
                self.sound_timer = self.reg[vx as usize];
                self.st_tick = 1.0 / 60.0;
            },
            AddToI(vx) => {
                self.i += self.reg[vx as usize] as usize;
            },
            LoadHexGlyph(vx) => {
                let x = self.reg[vx as usize];
                self.i = FONT_ADDR + x as usize * FONT_HEIGHT;
            }
            StoreBCD(vx) => {
                let mut x = self.reg[vx as usize];

                let mut place = 100;
                for i in 0us..2 {
                    self.i += i;
                    let bcd = x / place;
                    self.ram[self.i] = bcd;
                    x -= bcd;
                    place /= 10;
                }
            }
            StoreRegisters(vx) => {
                let vx = vx as usize;
                let i = self.i as usize;

                let mut dst = &mut self.ram[i..i+vx+1];
                for (x,b) in dst.iter_mut().enumerate() {
                  //println!("store reg {} to {}", x, i);
                    *b = self.reg[x];
                }
            },
            LoadRegisters(vx) => {
                let vx = vx as usize;
                let i = self.i as usize;

                let src = &self.ram[i..i+vx+1];
                for (x,b) in src.iter().enumerate() {
                  //println!("read reg {} from {}", x, i);
                    self.reg[x] = *b;
                }
            },
            other => {
                println!("Instruction not implemented {:?} skipping...", other)
            }
        }
        return false;
    }

    fn time_step(&mut self) {
        if self.timer > 0 {
            self.t_tick -= STEP_TIME;
            if self.t_tick <= 0.0 {
                self.timer -= 1;
                self.t_tick = 1.0 / 60.0;
            }
        }

        if self.sound_timer > 0 {
            self.st_tick -= STEP_TIME;
            if self.st_tick <= 0.0 {
                self.sound_timer -= 1;
                self.st_tick = 1.0 / 60.0;
            }
        }
    }

    // dt: Time in seconds since last step
    pub fn step(&mut self, dt:f32) -> bool {
        let steps = (CLOCK_SPEED_HZ as f32 * dt).round() as u32;
        let mut idle = false;

        for _ in 0u32..steps {
            let raw = {
                let codes = &self.ram[self.pc..self.pc+2];
                ((codes[0] as u16) << 8) | codes[1] as u16
            };
            let op = Op::new(raw);
            self.pc += 2;
            idle = self.exec(&op);
            self.time_step();
            if idle { break; }
        }
        return idle;
    }

    pub fn screen_rows<'a>(&'a self) -> Chunks<'a, u8> {
        self.screen.chunks(64)
    }

    pub fn print_screen(&self) {
        for row in self.screen.chunks(64) {
            println!("");
            for byte in row.iter() {
                match *byte {
                    0x0 => print!("░"),
                    _ => print!("▓")
                }
            }
        }
    }

    pub fn print_disassembly(&self) {
        for i in self.ram.chunks(2) {
            match i {
                [0, 0] => continue,
                [h, l] => {
                    let op = Op::new(((h as u16) << 8) | l as u16);
                    println!("raw: 0x{:X}", op.raw());
                    println!("instruction: {:?}", Instruction::from_op(&op));
                    println!("addr: 0x{:X}", op.addr());
                    println!("x: 0x{:X}", op.x());
                    println!("y: 0x{:X}", op.y());
                    println!("n_high: 0x{:X}", op.n_high());
                    println!("n_low: 0x{:X}", op.n_low());
                    println!("k: 0x{:X}\n", op.k());
                },
                _ => continue
            }
        }
    }
}
<|endoftext|>


<commit_before>use eventual::{self, Async};
use portaudio;
use std::borrow::Cow;
use std::sync::{mpsc, Mutex, Arc, MutexGuard};
use std::thread;
use vorbis;

use metadata::{FileFormat, Track, TrackRef};
use session::{Bitrate, Session};
use audio_decrypt::AudioDecrypt;
use util::{self, SpotifyId, Subfile};
use spirc::PlayStatus;

pub type PlayerObserver = Box<Fn(&PlayerState) + Send>;

pub struct Player {
    state: Arc<Mutex<PlayerState>>,
    observers: Arc<Mutex<Vec<PlayerObserver>>>,

    commands: mpsc::Sender<PlayerCommand>,
}

#[derive(Clone)]
pub struct PlayerState {
    status: PlayStatus,
    position_ms: u32,
    position_measured_at: i64,
    update_time: i64,
    volume: u16,

    end_of_track: bool,
}

struct PlayerInternal {
    state: Arc<Mutex<PlayerState>>,
    observers: Arc<Mutex<Vec<PlayerObserver>>>,

    session: Session,
    commands: mpsc::Receiver<PlayerCommand>,
}

#[derive(Debug)]
enum PlayerCommand {
    Load(SpotifyId, bool, u32),
    Play,
    Pause,
    Volume(u16),
    Stop,
    Seek(u32),
}

impl Player {
    pub fn new(session: Session) -> Player {
        let (cmd_tx, cmd_rx) = mpsc::channel();

        let state = Arc::new(Mutex::new(PlayerState {
            status: PlayStatus::kPlayStatusStop,
            position_ms: 0,
            position_measured_at: 0,
            update_time: util::now_ms(),
            volume: 0xFFFF,
            end_of_track: false,
        }));

        let observers = Arc::new(Mutex::new(Vec::new()));

        let internal = PlayerInternal {
            session: session,
            commands: cmd_rx,
            state: state.clone(),
            observers: observers.clone(),
        };

        thread::spawn(move || internal.run());

        Player {
            commands: cmd_tx,
            state: state,
            observers: observers,
        }
    }

    fn command(&self, cmd: PlayerCommand) {
        self.commands.send(cmd).unwrap();
    }

    pub fn load(&self, track: SpotifyId, start_playing: bool, position_ms: u32) {
        self.command(PlayerCommand::Load(track, start_playing, position_ms));
    }

    pub fn play(&self) {
        self.command(PlayerCommand::Play)
    }

    pub fn pause(&self) {
        self.command(PlayerCommand::Pause)
    }

    pub fn stop(&self) {
        self.command(PlayerCommand::Stop)
    }

    pub fn seek(&self, position_ms: u32) {
        self.command(PlayerCommand::Seek(position_ms));
    }

    pub fn state(&self) -> PlayerState {
        self.state.lock().unwrap().clone()
    }

    pub fn volume(&self, vol: u16) {
        self.command(PlayerCommand::Volume(vol));
    }

    pub fn add_observer(&self, observer: PlayerObserver) {
        self.observers.lock().unwrap().push(observer);
    }
}

fn apply_volume(volume: u16, data: &[i16]) -> Cow<[i16]> {
    // Fast path when volume is 100%
    if volume == 0xFFFF {
        Cow::Borrowed(data)
    } else {
        Cow::Owned(data.iter()
                       .map(|&x| {
                           (x as i32
                            * volume as i32
                            / 0xFFFF) as i16
                       })
                       .collect())
    }
}

impl PlayerInternal {
    fn run(self) {
        portaudio::initialize().unwrap();

        let stream = portaudio::stream::Stream::<i16, i16>::open_default(
                0, 2, 44100.0,
                portaudio::stream::FRAMES_PER_BUFFER_UNSPECIFIED,
                None
        ).unwrap();

        let mut decoder = None;

        loop {
            let playing = self.state.lock().unwrap().status == PlayStatus::kPlayStatusPlay;
            let cmd = if playing {
                self.commands.try_recv().ok()
            } else {
                Some(self.commands.recv().unwrap())
            };

            match cmd {
                Some(PlayerCommand::Load(track_id, play, position)) => {
                    self.update(|state| {
                        if state.status == PlayStatus::kPlayStatusPlay {
                            stream.stop().unwrap();
                        }
                        state.end_of_track = false;
                        state.status = if play {
                            PlayStatus::kPlayStatusPlay
                        } else {
                            PlayStatus::kPlayStatusPause
                        };
                        state.position_ms = position;
                        state.position_measured_at = util::now_ms();
                        true
                    });
                    drop(decoder);

                    let mut track = self.session.metadata::<Track>(track_id).await().unwrap();

                    if !track.available {
                        let alternatives = track.alternatives
                                                .iter()
                                                .map(|alt_id| {
                                                    self.session.metadata::<Track>(*alt_id)
                                                })
                                                .collect::<Vec<TrackRef>>();

                        track = eventual::sequence(alternatives.into_iter())
                                    .iter()
                                    .find(|alt| alt.available)
                                    .unwrap();
                    }

                    let format = match self.session.config().bitrate {
                        Bitrate::Bitrate96 => FileFormat::OGG_VORBIS_96,
                        Bitrate::Bitrate160 => FileFormat::OGG_VORBIS_160,
                        Bitrate::Bitrate320 => FileFormat::OGG_VORBIS_320,
                    };
                    let (file_id, _) = track.files.into_iter().find(|&(_, f)| f == format).unwrap();

                    let key = self.session.audio_key(track.id, file_id).await().unwrap();
                    decoder = Some(
                        vorbis::Decoder::new(
                        Subfile::new(
                        AudioDecrypt::new(key,
                        self.session.audio_file(file_id)), 0xa7)).unwrap());
                    decoder.as_mut().unwrap().time_seek(position as f64 / 1000f64).unwrap();

                    self.update(|state| {
                        state.status = if play {
                            stream.start().unwrap();
                            PlayStatus::kPlayStatusPlay
                        } else {
                            PlayStatus::kPlayStatusPause
                        };
                        state.position_ms = position;
                        state.position_measured_at = util::now_ms();

                        true
                    });
                    println!("Load Done");
                }
                Some(PlayerCommand::Seek(ms)) => {
                    decoder.as_mut().unwrap().time_seek(ms as f64 / 1000f64).unwrap();
                    self.update(|state| {
                        state.position_ms =
                            (decoder.as_mut().unwrap().time_tell().unwrap() * 1000f64) as u32;
                        state.position_measured_at = util::now_ms();

                        true
                    });
                }
                Some(PlayerCommand::Play) => {
                    self.update(|state| {
                        state.status = PlayStatus::kPlayStatusPlay;
                        true
                    });

                    stream.start().unwrap();
                }
                Some(PlayerCommand::Pause) => {
                    self.update(|state| {
                        state.status = PlayStatus::kPlayStatusPause;
                        state.update_time = util::now_ms();
                        true
                    });

                    stream.stop().unwrap();
                }
                Some(PlayerCommand::Volume(vol)) => {
                    self.update(|state| {
                        state.volume = vol;
                        true
                    });
                }
                Some(PlayerCommand::Stop) => {
                    self.update(|state| {
                        if state.status == PlayStatus::kPlayStatusPlay {
                            state.status = PlayStatus::kPlayStatusPause;
                        }
                        true
                    });

                    stream.stop().unwrap();
                    decoder = None;
                }
                None => (),
            }

            if self.state.lock().unwrap().status == PlayStatus::kPlayStatusPlay {
                match decoder.as_mut().unwrap().packets().next() {
                    Some(Ok(packet)) => {
                        let buffer = apply_volume(self.state.lock().unwrap().volume,
                                                  &packet.data);
                        match stream.write(&buffer) {
                            Ok(_) => (),
                            Err(portaudio::PaError::OutputUnderflowed) => eprintln!("Underflow"),
                            Err(e) => panic!("PA Error {}", e),
                        };
                    }
                    Some(Err(vorbis::VorbisError::Hole)) => (),
                    Some(Err(e)) => panic!("Vorbis error {:?}", e),
                    None => {
                        self.update(|state| {
                            state.status = PlayStatus::kPlayStatusStop;
                            state.end_of_track = true;
                            true
                        });

                        stream.stop().unwrap();
                        decoder = None;
                    }
                }

                self.update(|state| {
                    let now = util::now_ms();

                    if now - state.position_measured_at > 5000 {
                        state.position_ms =
                            (decoder.as_mut().unwrap().time_tell().unwrap() * 1000f64) as u32;
                        state.position_measured_at = now;

                        true
                    } else {
                        false
                    }
                });
            }
        }

        drop(stream);

        portaudio::terminate().unwrap();
    }

    fn update<F>(&self, f: F)
        where F: FnOnce(&mut MutexGuard<PlayerState>) -> bool
    {
        let mut guard = self.state.lock().unwrap();
        let update = f(&mut guard);

        let observers = self.observers.lock().unwrap();
        if update {
            guard.update_time = util::now_ms();
            let state = guard.clone();
            drop(guard);

            for observer in observers.iter() {
                observer(&state);
            }
        }
    }
}

impl PlayerState {
    pub fn status(&self) -> PlayStatus {
        self.status
    }

    pub fn position(&self) -> (u32, i64) {
        (self.position_ms, self.position_measured_at)
    }

    pub fn volume(&self) -> u16 {
        self.volume
    }

    pub fn update_time(&self) -> i64 {
        self.update_time
    }

    pub fn end_of_track(&self) -> bool {
        self.end_of_track
    }
}
<commit_msg>player: Update the position on pause/play events.<commit_after>use eventual::{self, Async};
use portaudio;
use std::borrow::Cow;
use std::sync::{mpsc, Mutex, Arc, MutexGuard};
use std::thread;
use vorbis;

use metadata::{FileFormat, Track, TrackRef};
use session::{Bitrate, Session};
use audio_decrypt::AudioDecrypt;
use util::{self, SpotifyId, Subfile};
use spirc::PlayStatus;

pub type PlayerObserver = Box<Fn(&PlayerState) + Send>;

pub struct Player {
    state: Arc<Mutex<PlayerState>>,
    observers: Arc<Mutex<Vec<PlayerObserver>>>,

    commands: mpsc::Sender<PlayerCommand>,
}

#[derive(Clone)]
pub struct PlayerState {
    status: PlayStatus,
    position_ms: u32,
    position_measured_at: i64,
    update_time: i64,
    volume: u16,

    end_of_track: bool,
}

struct PlayerInternal {
    state: Arc<Mutex<PlayerState>>,
    observers: Arc<Mutex<Vec<PlayerObserver>>>,

    session: Session,
    commands: mpsc::Receiver<PlayerCommand>,
}

#[derive(Debug)]
enum PlayerCommand {
    Load(SpotifyId, bool, u32),
    Play,
    Pause,
    Volume(u16),
    Stop,
    Seek(u32),
}

impl Player {
    pub fn new(session: Session) -> Player {
        let (cmd_tx, cmd_rx) = mpsc::channel();

        let state = Arc::new(Mutex::new(PlayerState {
            status: PlayStatus::kPlayStatusStop,
            position_ms: 0,
            position_measured_at: 0,
            update_time: util::now_ms(),
            volume: 0xFFFF,
            end_of_track: false,
        }));

        let observers = Arc::new(Mutex::new(Vec::new()));

        let internal = PlayerInternal {
            session: session,
            commands: cmd_rx,
            state: state.clone(),
            observers: observers.clone(),
        };

        thread::spawn(move || internal.run());

        Player {
            commands: cmd_tx,
            state: state,
            observers: observers,
        }
    }

    fn command(&self, cmd: PlayerCommand) {
        self.commands.send(cmd).unwrap();
    }

    pub fn load(&self, track: SpotifyId, start_playing: bool, position_ms: u32) {
        self.command(PlayerCommand::Load(track, start_playing, position_ms));
    }

    pub fn play(&self) {
        self.command(PlayerCommand::Play)
    }

    pub fn pause(&self) {
        self.command(PlayerCommand::Pause)
    }

    pub fn stop(&self) {
        self.command(PlayerCommand::Stop)
    }

    pub fn seek(&self, position_ms: u32) {
        self.command(PlayerCommand::Seek(position_ms));
    }

    pub fn state(&self) -> PlayerState {
        self.state.lock().unwrap().clone()
    }

    pub fn volume(&self, vol: u16) {
        self.command(PlayerCommand::Volume(vol));
    }

    pub fn add_observer(&self, observer: PlayerObserver) {
        self.observers.lock().unwrap().push(observer);
    }
}

fn apply_volume(volume: u16, data: &[i16]) -> Cow<[i16]> {
    // Fast path when volume is 100%
    if volume == 0xFFFF {
        Cow::Borrowed(data)
    } else {
        Cow::Owned(data.iter()
                       .map(|&x| {
                           (x as i32
                            * volume as i32
                            / 0xFFFF) as i16
                       })
                       .collect())
    }
}

impl PlayerInternal {
    fn run(self) {
        portaudio::initialize().unwrap();

        let stream = portaudio::stream::Stream::<i16, i16>::open_default(
                0, 2, 44100.0,
                portaudio::stream::FRAMES_PER_BUFFER_UNSPECIFIED,
                None
        ).unwrap();

        let mut decoder = None;

        loop {
            let playing = self.state.lock().unwrap().status == PlayStatus::kPlayStatusPlay;
            let cmd = if playing {
                self.commands.try_recv().ok()
            } else {
                Some(self.commands.recv().unwrap())
            };

            match cmd {
                Some(PlayerCommand::Load(track_id, play, position)) => {
                    self.update(|state| {
                        if state.status == PlayStatus::kPlayStatusPlay {
                            stream.stop().unwrap();
                        }
                        state.end_of_track = false;
                        state.status = PlayStatus::kPlayStatusPause;
                        state.position_ms = position;
                        state.position_measured_at = util::now_ms();
                        true
                    });
                    drop(decoder);

                    let mut track = self.session.metadata::<Track>(track_id).await().unwrap();

                    if !track.available {
                        let alternatives = track.alternatives
                                                .iter()
                                                .map(|alt_id| {
                                                    self.session.metadata::<Track>(*alt_id)
                                                })
                                                .collect::<Vec<TrackRef>>();

                        track = eventual::sequence(alternatives.into_iter())
                                    .iter()
                                    .find(|alt| alt.available)
                                    .unwrap();
                    }

                    let format = match self.session.config().bitrate {
                        Bitrate::Bitrate96 => FileFormat::OGG_VORBIS_96,
                        Bitrate::Bitrate160 => FileFormat::OGG_VORBIS_160,
                        Bitrate::Bitrate320 => FileFormat::OGG_VORBIS_320,
                    };
                    let (file_id, _) = track.files.into_iter().find(|&(_, f)| f == format).unwrap();

                    let key = self.session.audio_key(track.id, file_id).await().unwrap();
                    decoder = Some(
                        vorbis::Decoder::new(
                        Subfile::new(
                        AudioDecrypt::new(key,
                        self.session.audio_file(file_id)), 0xa7)).unwrap());
                    decoder.as_mut().unwrap().time_seek(position as f64 / 1000f64).unwrap();

                    self.update(|state| {
                        state.status = if play {
                            stream.start().unwrap();
                            PlayStatus::kPlayStatusPlay
                        } else {
                            PlayStatus::kPlayStatusPause
                        };
                        state.position_ms = position;
                        state.position_measured_at = util::now_ms();

                        true
                    });
                    println!("Load Done");
                }
                Some(PlayerCommand::Seek(ms)) => {
                    decoder.as_mut().unwrap().time_seek(ms as f64 / 1000f64).unwrap();
                    self.update(|state| {
                        state.position_ms =
                            (decoder.as_mut().unwrap().time_tell().unwrap() * 1000f64) as u32;
                        state.position_measured_at = util::now_ms();

                        true
                    });
                }
                Some(PlayerCommand::Play) => {
                    self.update(|state| {
                        state.status = PlayStatus::kPlayStatusPlay;
                        state.position_ms =
                            (decoder.as_mut().unwrap().time_tell().unwrap() * 1000f64) as u32;
                        state.position_measured_at = util::now_ms();
                        true
                    });

                    stream.start().unwrap();
                }
                Some(PlayerCommand::Pause) => {
                    self.update(|state| {
                        state.status = PlayStatus::kPlayStatusPause;
                        state.update_time = util::now_ms();
                        state.position_ms =
                            (decoder.as_mut().unwrap().time_tell().unwrap() * 1000f64) as u32;
                        state.position_measured_at = util::now_ms();
                        true
                    });

                    stream.stop().unwrap();
                }
                Some(PlayerCommand::Volume(vol)) => {
                    self.update(|state| {
                        state.volume = vol;
                        true
                    });
                }
                Some(PlayerCommand::Stop) => {
                    self.update(|state| {
                        if state.status == PlayStatus::kPlayStatusPlay {
                            state.status = PlayStatus::kPlayStatusPause;
                        }
                        state.position_ms = 0;
                        state.position_measured_at = util::now_ms();
                        true
                    });

                    stream.stop().unwrap();
                    decoder = None;
                }
                None => (),
            }

            if self.state.lock().unwrap().status == PlayStatus::kPlayStatusPlay {
                match decoder.as_mut().unwrap().packets().next() {
                    Some(Ok(packet)) => {
                        let buffer = apply_volume(self.state.lock().unwrap().volume,
                                                  &packet.data);
                        match stream.write(&buffer) {
                            Ok(_) => (),
                            Err(portaudio::PaError::OutputUnderflowed) => eprintln!("Underflow"),
                            Err(e) => panic!("PA Error {}", e),
                        };
                    }
                    Some(Err(vorbis::VorbisError::Hole)) => (),
                    Some(Err(e)) => panic!("Vorbis error {:?}", e),
                    None => {
                        self.update(|state| {
                            state.status = PlayStatus::kPlayStatusStop;
                            state.end_of_track = true;
                            true
                        });

                        stream.stop().unwrap();
                        decoder = None;
                    }
                }
            }
        }

        drop(stream);

        portaudio::terminate().unwrap();
    }

    fn update<F>(&self, f: F)
        where F: FnOnce(&mut MutexGuard<PlayerState>) -> bool
    {
        let mut guard = self.state.lock().unwrap();
        let update = f(&mut guard);

        let observers = self.observers.lock().unwrap();
        if update {
            guard.update_time = util::now_ms();
            let state = guard.clone();
            drop(guard);

            for observer in observers.iter() {
                observer(&state);
            }
        }
    }
}

impl PlayerState {
    pub fn status(&self) -> PlayStatus {
        self.status
    }

    pub fn position(&self) -> (u32, i64) {
        (self.position_ms, self.position_measured_at)
    }

    pub fn volume(&self) -> u16 {
        self.volume
    }

    pub fn update_time(&self) -> i64 {
        self.update_time
    }

    pub fn end_of_track(&self) -> bool {
        self.end_of_track
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Joe Neeman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use aho_corasick::{Automaton, AcAutomaton, FullAcAutomaton, MatchesOverlapping};
use memchr::memchr;
use memmem::{Searcher, TwoWaySearcher};

#[derive(Clone, Debug)]
pub enum Prefix {
    Empty,
    ByteSet(Vec<bool>),
    Byte(u8),
    Lit(Vec<u8>),
    Ac(FullAcAutomaton<Vec<u8>>, Vec<usize>),
    LoopWhile(Vec<bool>),
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PrefixResult {
    pub start_pos: usize,
    pub end_pos: usize,
    pub end_state: usize,
}

pub trait PrefixSearcher {
    fn skip_to(&mut self, pos: usize);
    fn search(&mut self) -> Option<PrefixResult>;
}

impl Prefix {
    pub fn from_strings<P: AsRef<[u8]>, I: Iterator<Item=(P, usize)>>(it: I) -> Prefix {
        let strings: Vec<(Vec<u8>, usize)> = it
            .filter(|x| !x.0.as_ref().is_empty())
            .map(|(s, x)| (s.as_ref().to_vec(), x))
            .collect();

        if strings.is_empty() {
            Prefix::Empty
        } else if strings.len() == 1 {
            if strings[0].0.len() == 1 {
                Prefix::Byte(strings[0].0[0])
            } else {
                Prefix::Lit(strings.into_iter().next().unwrap().0)
            }
        } else if strings.iter().map(|x| x.0.len()).min() == Some(1) {
            let mut bs = vec![false; 256];
            for (s, _) in strings.into_iter() {
                bs[s[0] as usize] = true;
            }
            Prefix::ByteSet(bs)
        } else {
            let state_map: Vec<_> = strings.iter().map(|x| x.1).collect();
            let ac = FullAcAutomaton::new(AcAutomaton::new(strings.into_iter().map(|x| x.0)));
            Prefix::Ac(ac, state_map)
        }
    }

    pub fn make_searcher<'a>(&'a self, input: &'a [u8]) -> Box<PrefixSearcher + 'a> {
        use prefix::Prefix::*;

        match self {
            &Empty => Box::new(SimpleSearcher::new((), input)),
            &ByteSet(ref bs) => Box::new(SimpleSearcher::new(&bs[..], input)),
            &Byte(b) => Box::new(SimpleSearcher::new(b, input)),
            &Lit(ref l) => Box::new(lit_searcher(l, input)),
            &LoopWhile(ref bs) => Box::new(loop_searcher(&bs[..], input)),
            &Ac(ref ac, ref map) => Box::new(AcSearcher::new(ac, map, input)),
        }
    }
}

trait SkipFn {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)>;
}

trait SimpleSkipFn {
    fn simple_skip(&self, input: &[u8]) -> Option<usize>;
}

impl<Sk: SimpleSkipFn> SkipFn for Sk {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)> {
        self.simple_skip(input).map(|x| (x, x))
    }
}

impl SimpleSkipFn for () {
    fn simple_skip(&self, _: &[u8]) -> Option<usize> { Some(0) }
}

impl SimpleSkipFn for u8 {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> { memchr(*self, input) }
}

impl<'a> SimpleSkipFn for TwoWaySearcher<'a> {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> { self.search_in(input) }
}

impl<'a> SimpleSkipFn for &'a [bool] {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> {
        input.iter().position(|c| self[*c as usize])
    }
}

struct LoopWhile<'a>(&'a [bool]);
impl<'a> SkipFn for LoopWhile<'a> {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)> {
        Some((0, input.iter().position(|c| !self.0[*c as usize]).unwrap_or(input.len())))
    }
}

struct SimpleSearcher<'a, Skip: SkipFn> {
    skip_fn: Skip,
    input: &'a [u8],
    pos: usize,
}

impl<'a, Sk: SkipFn> SimpleSearcher<'a, Sk> {
    fn new(skip_fn: Sk, input: &'a [u8]) -> SimpleSearcher<'a, Sk> {
        SimpleSearcher {
            skip_fn: skip_fn,
            input: input,
            pos: 0,
        }
    }
}

fn lit_searcher<'i, 'lit>(lit: &'lit [u8], input: &'i [u8])
-> SimpleSearcher<'i, TwoWaySearcher<'lit>> {
    SimpleSearcher {
        skip_fn: TwoWaySearcher::new(lit),
        input: input,
        pos: 0,
    }
}

fn loop_searcher<'i, 'lo>(loop_while: &'lo [bool], input: &'i [u8])
-> SimpleSearcher<'i, LoopWhile<'lo>> {
    SimpleSearcher {
        skip_fn: LoopWhile(loop_while),
        input: input,
        pos: 0,
    }
}

impl<'a, Sk: SkipFn> PrefixSearcher for SimpleSearcher<'a, Sk> {
    fn search(&mut self) -> Option<PrefixResult> {
        if self.pos > self.input.len() {
            None
        } else if let Some((start_off, end_off)) = self.skip_fn.skip(&self.input[self.pos..]) {
            let start = self.pos + start_off;
            let end = self.pos + end_off;
            self.pos += end_off + 1;

            Some(PrefixResult {
                start_pos: start,
                end_pos: end,
                end_state: 0,
            })
        } else {
            None
        }
    }

    fn skip_to(&mut self, pos: usize) { self.pos = pos; }
}

struct AcSearcher<'ac, 'i, 'st> {
    ac: &'ac FullAcAutomaton<Vec<u8>>,
    state_map: &'st [usize],
    input: &'i [u8],
    pos: usize,
    iter: MatchesOverlapping<'ac, 'i, Vec<u8>, FullAcAutomaton<Vec<u8>>>,
}

impl<'ac, 'i, 'st> AcSearcher<'ac, 'i, 'st> {
    fn new(ac: &'ac FullAcAutomaton<Vec<u8>>, state_map: &'st [usize], input: &'i [u8])
    -> AcSearcher<'ac, 'i, 'st> {
        AcSearcher {
            ac: ac,
            state_map: state_map,
            input: input,
            pos: 0,
            iter: ac.find_overlapping(input),
        }
    }
}

impl<'ac, 'i, 'st> PrefixSearcher for AcSearcher<'ac, 'i, 'st> {
    fn skip_to(&mut self, pos: usize) {
        self.pos = pos;
        let input: &'i [u8] = if pos > self.input.len() {
            &[]
        } else {
            &self.input[self.pos..]
        };
        self.iter = self.ac.find_overlapping(input);
    }

    fn search(&mut self) -> Option<PrefixResult> {
        self.iter.next().map(|mat| PrefixResult {
            start_pos: mat.start,
            end_pos: mat.end,
            end_state: self.state_map[mat.pati],
        })
    }
}

#[cfg(test)]
mod tests {
    use ::prefix::*;

    impl<'a> Iterator for Box<PrefixSearcher + 'a> {
        type Item = PrefixResult;
        fn next(&mut self) -> Option<PrefixResult> {
            self.search()
        }
    }

    fn search(pref: Prefix, input: &str) -> Vec<PrefixResult> {
        pref.make_searcher(input.as_bytes()).collect::<Vec<_>>()
    }

    fn result(pos: usize) -> PrefixResult {
        PrefixResult {
            start_pos: pos,
            end_pos: pos,
            end_state: 0,
        }
    }

    fn results(posns: Vec<usize>) -> Vec<PrefixResult> {
        posns.into_iter().map(result).collect()
    }

    #[test]
    fn test_empty_search() {
        assert_eq!(search(Prefix::Empty, "blah"), results(vec![0, 1, 2, 3, 4]));
        assert_eq!(search(Prefix::Empty, ""), results(vec![0]));
    }

    #[test]
    fn test_byte_search() {
        assert_eq!(search(Prefix::Byte(b'a'), "abracadabra"), results(vec![0, 3, 5, 7, 10]));
        assert_eq!(search(Prefix::Byte(b'a'), "abracadabr"), results(vec![0, 3, 5, 7]));
        assert_eq!(search(Prefix::Byte(b'a'), ""), vec![]);
    }

    #[test]
    fn test_str_search() {
        fn lit_pref(s: &str) -> Prefix {
            Prefix::Lit(s.as_bytes().to_vec())
        }
        assert_eq!(search(lit_pref("aa"), "baa baa black sheep aa"), results(vec![1, 5, 20]));
        assert_eq!(search(lit_pref("aa"), "aaa baaa black sheep"), results(vec![0, 1, 5, 6]));
        assert_eq!(search(lit_pref("aa"), ""), vec![]);
    }

    #[test]
    fn test_byteset_search() {
        fn bs_pref(s: &str) -> Prefix {
            let mut bytes = vec![false; 256];
            for &b in s.as_bytes().iter() {
                bytes[b as usize] = true;
            }
            Prefix::ByteSet(bytes)
        }
        assert_eq!(search(bs_pref("aeiou"), "quick brown"), results(vec![1, 2, 8]));
        assert_eq!(search(bs_pref("aeiou"), "aabaa"), results(vec![0, 1, 3, 4]));
        assert_eq!(search(bs_pref("aeiou"), ""), vec![]);
    }

    fn pair_results(posns: Vec<(usize, usize)>) -> Vec<PrefixResult> {
        posns.into_iter()
            .map(|(s, e)| PrefixResult { start_pos: s, end_pos: e, end_state: 0 })
            .collect()
    }

    #[test]
    fn test_loop_search() {
        fn loop_pref(s: &str) -> Prefix {
            let mut bytes = vec![false; 256];
            for &b in s.as_bytes().iter() {
                bytes[b as usize] = true;
            }
            Prefix::LoopWhile(bytes)
        }
        assert_eq!(search(loop_pref("aeiou"), "quick"),
            pair_results(vec![(0, 0), (1, 3), (4, 4), (5, 5)]));
        assert_eq!(search(loop_pref("aeiou"), "aabaa"),
            pair_results(vec![(0, 2), (3, 5)]));
        assert_eq!(search(loop_pref("aeiou"), ""), pair_results(vec![(0, 0)]));
    }

    #[test]
    fn test_ac_search() {
        fn ac_pref(strs: Vec<&str>) -> Prefix {
            let len = strs.len();
            let pref = Prefix::from_strings(strs.into_iter().zip(0..len));
            assert!(matches!(pref, Prefix::Ac(_, _)));
            pref
        }

        assert_eq!(search(ac_pref(vec!["baa", "aa"]), "baa aaa black sheep"),
            vec![
                PrefixResult { start_pos: 0, end_pos: 3, end_state: 0 },
                PrefixResult { start_pos: 1, end_pos: 3, end_state: 1 },
                PrefixResult { start_pos: 4, end_pos: 6, end_state: 1 },
                PrefixResult { start_pos: 5, end_pos: 7, end_state: 1 },
            ]);
        assert_eq!(search(ac_pref(vec!["baa", "aa"]), ""), vec![]);
    }

    #[test]
    fn test_prefix_choice() {
        use ::prefix::Prefix::*;

        fn pref(strs: Vec<&str>) -> Prefix {
            let len = strs.len();
            Prefix::from_strings(strs.into_iter().zip(0..len))
        }

        assert!(matches!(pref(vec![]), Empty));
        assert!(matches!(pref(vec![""]), Empty));
        assert!(matches!(pref(vec!["a"]), Byte(_)));
        assert!(matches!(pref(vec!["", "a", ""]), Byte(_)));
        assert!(matches!(pref(vec!["abc"]), Lit(_)));
        assert!(matches!(pref(vec!["abc", ""]), Lit(_)));
        assert!(matches!(pref(vec!["a", "b", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["a", "b", "", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["a", "baa", "", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["ab", "baa", "", "cb"]), Ac(_, _)));
    }
}

<commit_msg>Begin adding documentation.<commit_after>// Copyright 2015 Joe Neeman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use aho_corasick::{Automaton, AcAutomaton, FullAcAutomaton, MatchesOverlapping};
use memchr::memchr;
use memmem::{Searcher, TwoWaySearcher};

/// A `Prefix` is the first part of a DFA. Anything matching the DFA should start with
/// something matching the `Prefix`.
///
/// The purpose of a `Prefix` is that scanning through the input looking for the `Prefix` should be
/// much faster than running the DFA naively.
#[derive(Clone, Debug)]
pub enum Prefix {
    // Matches every position.
    Empty,
    // Matches a single byte in a particular set.
    ByteSet(Vec<bool>),
    // Matches one specific byte.
    Byte(u8),
    // Matches a specific sequence of bytes.
    Lit(Vec<u8>),
    // Matches one of several sequences of bytes. The sequences are contained in the
    // `FullAcAutomaton`. The `Vec<usize>` tells us which state the DFA should start in after
    // matching each sequence. That is, `vec[i] == s` if after finding sequence `i` we should
    // start in state `s`.
    Ac(FullAcAutomaton<Vec<u8>>, Vec<usize>),
    // Matches a maximal (but possibly non-empty) sequence of bytes each of which belong to a
    // particular set of bytes.
    //
    // The interesting thing about this prefix is that it will only look for non-overlapping
    // matches (whereas most prefixes look for overlapping matches). This is intended for use when
    // the first state of a DFA contains self-transitions. In that case, if there's a sequence of
    // bytes that keeps us in the first state then there's no point in trying to start in the
    // middle of that sequence of bytes: even if that would give a match, we would get an earlier
    // match from starting at the beginning of the sequence.
    LoopWhile(Vec<bool>),
}

/// The result of scanning through the input for a `Prefix`.
///
/// The semi-open interval `[start_pos, end_pos)` is the part of the interval that was consumed by
/// the `Prefix`. The state `end_state` is the DFA state at which we should start to continue
/// matching; that is, the DFA should begin at position `end_pos` in state `end_state`.
///
/// Note that some `Prefix`es return empty intervals (`start_pos == end_pos`). This doesn't mean
/// necessarily that the `Prefix` didn't match any input, only that it's simpler (and fast) just
/// to have the DFA re-match from the beginning.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PrefixResult {
    pub start_pos: usize,
    pub end_pos: usize,
    pub end_state: usize,
}

/// Encapsulates the `Prefix` and the input string, and allows iteration over all matches.
pub trait PrefixSearcher {
    /// Moves the "cursor" to the given position in the input.
    fn skip_to(&mut self, pos: usize);
    /// From the current position in the input, finds the next substring matching the `Prefix`
    /// and advances the "cursor" past that point.
    fn search(&mut self) -> Option<PrefixResult>;
}

impl Prefix {
    /// Converts a set of `(string, state)` pairs into a `Prefix` that matches any of the strings.
    ///
    /// The `state` part of each `(string, state)` pair is the DFA state that we would be in after
    /// matching the `string`.
    pub fn from_strings<P: AsRef<[u8]>, I: Iterator<Item=(P, usize)>>(it: I) -> Prefix {
        let strings: Vec<(Vec<u8>, usize)> = it
            .filter(|x| !x.0.as_ref().is_empty())
            .map(|(s, x)| (s.as_ref().to_vec(), x))
            .collect();

        if strings.is_empty() {
            Prefix::Empty
        } else if strings.len() == 1 {
            if strings[0].0.len() == 1 {
                Prefix::Byte(strings[0].0[0])
            } else {
                Prefix::Lit(strings.into_iter().next().unwrap().0)
            }
        } else if strings.iter().map(|x| x.0.len()).min() == Some(1) {
            let mut bs = vec![false; 256];
            for (s, _) in strings.into_iter() {
                bs[s[0] as usize] = true;
            }
            Prefix::ByteSet(bs)
        } else {
            let state_map: Vec<_> = strings.iter().map(|x| x.1).collect();
            let ac = FullAcAutomaton::new(AcAutomaton::new(strings.into_iter().map(|x| x.0)));
            Prefix::Ac(ac, state_map)
        }
    }

    /// Takes an input string and prepares for quickly finding matches in it.
    pub fn make_searcher<'a>(&'a self, input: &'a [u8]) -> Box<PrefixSearcher + 'a> {
        use prefix::Prefix::*;

        match self {
            &Empty => Box::new(SimpleSearcher::new((), input)),
            &ByteSet(ref bs) => Box::new(SimpleSearcher::new(&bs[..], input)),
            &Byte(b) => Box::new(SimpleSearcher::new(b, input)),
            &Lit(ref l) => Box::new(lit_searcher(l, input)),
            &LoopWhile(ref bs) => Box::new(loop_searcher(&bs[..], input)),
            &Ac(ref ac, ref map) => Box::new(AcSearcher::new(ac, map, input)),
        }
    }
}

trait SkipFn {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)>;
}

trait SimpleSkipFn {
    fn simple_skip(&self, input: &[u8]) -> Option<usize>;
}

impl<Sk: SimpleSkipFn> SkipFn for Sk {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)> {
        self.simple_skip(input).map(|x| (x, x))
    }
}

impl SimpleSkipFn for () {
    fn simple_skip(&self, _: &[u8]) -> Option<usize> { Some(0) }
}

impl SimpleSkipFn for u8 {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> { memchr(*self, input) }
}

impl<'a> SimpleSkipFn for TwoWaySearcher<'a> {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> { self.search_in(input) }
}

impl<'a> SimpleSkipFn for &'a [bool] {
    fn simple_skip(&self, input: &[u8]) -> Option<usize> {
        input.iter().position(|c| self[*c as usize])
    }
}

struct LoopWhile<'a>(&'a [bool]);
impl<'a> SkipFn for LoopWhile<'a> {
    fn skip(&self, input: &[u8]) -> Option<(usize, usize)> {
        Some((0, input.iter().position(|c| !self.0[*c as usize]).unwrap_or(input.len())))
    }
}

struct SimpleSearcher<'a, Skip: SkipFn> {
    skip_fn: Skip,
    input: &'a [u8],
    pos: usize,
}

impl<'a, Sk: SkipFn> SimpleSearcher<'a, Sk> {
    fn new(skip_fn: Sk, input: &'a [u8]) -> SimpleSearcher<'a, Sk> {
        SimpleSearcher {
            skip_fn: skip_fn,
            input: input,
            pos: 0,
        }
    }
}

fn lit_searcher<'i, 'lit>(lit: &'lit [u8], input: &'i [u8])
-> SimpleSearcher<'i, TwoWaySearcher<'lit>> {
    SimpleSearcher {
        skip_fn: TwoWaySearcher::new(lit),
        input: input,
        pos: 0,
    }
}

fn loop_searcher<'i, 'lo>(loop_while: &'lo [bool], input: &'i [u8])
-> SimpleSearcher<'i, LoopWhile<'lo>> {
    SimpleSearcher {
        skip_fn: LoopWhile(loop_while),
        input: input,
        pos: 0,
    }
}

impl<'a, Sk: SkipFn> PrefixSearcher for SimpleSearcher<'a, Sk> {
    fn search(&mut self) -> Option<PrefixResult> {
        if self.pos > self.input.len() {
            None
        } else if let Some((start_off, end_off)) = self.skip_fn.skip(&self.input[self.pos..]) {
            let start = self.pos + start_off;
            let end = self.pos + end_off;
            self.pos += end_off + 1;

            Some(PrefixResult {
                start_pos: start,
                end_pos: end,
                end_state: 0,
            })
        } else {
            None
        }
    }

    fn skip_to(&mut self, pos: usize) { self.pos = pos; }
}

struct AcSearcher<'ac, 'i, 'st> {
    ac: &'ac FullAcAutomaton<Vec<u8>>,
    state_map: &'st [usize],
    input: &'i [u8],
    pos: usize,
    iter: MatchesOverlapping<'ac, 'i, Vec<u8>, FullAcAutomaton<Vec<u8>>>,
}

impl<'ac, 'i, 'st> AcSearcher<'ac, 'i, 'st> {
    fn new(ac: &'ac FullAcAutomaton<Vec<u8>>, state_map: &'st [usize], input: &'i [u8])
    -> AcSearcher<'ac, 'i, 'st> {
        AcSearcher {
            ac: ac,
            state_map: state_map,
            input: input,
            pos: 0,
            iter: ac.find_overlapping(input),
        }
    }
}

impl<'ac, 'i, 'st> PrefixSearcher for AcSearcher<'ac, 'i, 'st> {
    fn skip_to(&mut self, pos: usize) {
        self.pos = pos;
        let input: &'i [u8] = if pos > self.input.len() {
            &[]
        } else {
            &self.input[self.pos..]
        };
        self.iter = self.ac.find_overlapping(input);
    }

    fn search(&mut self) -> Option<PrefixResult> {
        self.iter.next().map(|mat| PrefixResult {
            start_pos: mat.start,
            end_pos: mat.end,
            end_state: self.state_map[mat.pati],
        })
    }
}

#[cfg(test)]
mod tests {
    use ::prefix::*;

    impl<'a> Iterator for Box<PrefixSearcher + 'a> {
        type Item = PrefixResult;
        fn next(&mut self) -> Option<PrefixResult> {
            self.search()
        }
    }

    fn search(pref: Prefix, input: &str) -> Vec<PrefixResult> {
        pref.make_searcher(input.as_bytes()).collect::<Vec<_>>()
    }

    fn result(pos: usize) -> PrefixResult {
        PrefixResult {
            start_pos: pos,
            end_pos: pos,
            end_state: 0,
        }
    }

    fn results(posns: Vec<usize>) -> Vec<PrefixResult> {
        posns.into_iter().map(result).collect()
    }

    #[test]
    fn test_empty_search() {
        assert_eq!(search(Prefix::Empty, "blah"), results(vec![0, 1, 2, 3, 4]));
        assert_eq!(search(Prefix::Empty, ""), results(vec![0]));
    }

    #[test]
    fn test_byte_search() {
        assert_eq!(search(Prefix::Byte(b'a'), "abracadabra"), results(vec![0, 3, 5, 7, 10]));
        assert_eq!(search(Prefix::Byte(b'a'), "abracadabr"), results(vec![0, 3, 5, 7]));
        assert_eq!(search(Prefix::Byte(b'a'), ""), vec![]);
    }

    #[test]
    fn test_str_search() {
        fn lit_pref(s: &str) -> Prefix {
            Prefix::Lit(s.as_bytes().to_vec())
        }
        assert_eq!(search(lit_pref("aa"), "baa baa black sheep aa"), results(vec![1, 5, 20]));
        assert_eq!(search(lit_pref("aa"), "aaa baaa black sheep"), results(vec![0, 1, 5, 6]));
        assert_eq!(search(lit_pref("aa"), ""), vec![]);
    }

    #[test]
    fn test_byteset_search() {
        fn bs_pref(s: &str) -> Prefix {
            let mut bytes = vec![false; 256];
            for &b in s.as_bytes().iter() {
                bytes[b as usize] = true;
            }
            Prefix::ByteSet(bytes)
        }
        assert_eq!(search(bs_pref("aeiou"), "quick brown"), results(vec![1, 2, 8]));
        assert_eq!(search(bs_pref("aeiou"), "aabaa"), results(vec![0, 1, 3, 4]));
        assert_eq!(search(bs_pref("aeiou"), ""), vec![]);
    }

    fn pair_results(posns: Vec<(usize, usize)>) -> Vec<PrefixResult> {
        posns.into_iter()
            .map(|(s, e)| PrefixResult { start_pos: s, end_pos: e, end_state: 0 })
            .collect()
    }

    #[test]
    fn test_loop_search() {
        fn loop_pref(s: &str) -> Prefix {
            let mut bytes = vec![false; 256];
            for &b in s.as_bytes().iter() {
                bytes[b as usize] = true;
            }
            Prefix::LoopWhile(bytes)
        }
        assert_eq!(search(loop_pref("aeiou"), "quick"),
            pair_results(vec![(0, 0), (1, 3), (4, 4), (5, 5)]));
        assert_eq!(search(loop_pref("aeiou"), "aabaa"),
            pair_results(vec![(0, 2), (3, 5)]));
        assert_eq!(search(loop_pref("aeiou"), ""), pair_results(vec![(0, 0)]));
    }

    #[test]
    fn test_ac_search() {
        fn ac_pref(strs: Vec<&str>) -> Prefix {
            let len = strs.len();
            let pref = Prefix::from_strings(strs.into_iter().zip(0..len));
            assert!(matches!(pref, Prefix::Ac(_, _)));
            pref
        }

        assert_eq!(search(ac_pref(vec!["baa", "aa"]), "baa aaa black sheep"),
            vec![
                PrefixResult { start_pos: 0, end_pos: 3, end_state: 0 },
                PrefixResult { start_pos: 1, end_pos: 3, end_state: 1 },
                PrefixResult { start_pos: 4, end_pos: 6, end_state: 1 },
                PrefixResult { start_pos: 5, end_pos: 7, end_state: 1 },
            ]);
        assert_eq!(search(ac_pref(vec!["baa", "aa"]), ""), vec![]);
    }

    #[test]
    fn test_prefix_choice() {
        use ::prefix::Prefix::*;

        fn pref(strs: Vec<&str>) -> Prefix {
            let len = strs.len();
            Prefix::from_strings(strs.into_iter().zip(0..len))
        }

        assert!(matches!(pref(vec![]), Empty));
        assert!(matches!(pref(vec![""]), Empty));
        assert!(matches!(pref(vec!["a"]), Byte(_)));
        assert!(matches!(pref(vec!["", "a", ""]), Byte(_)));
        assert!(matches!(pref(vec!["abc"]), Lit(_)));
        assert!(matches!(pref(vec!["abc", ""]), Lit(_)));
        assert!(matches!(pref(vec!["a", "b", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["a", "b", "", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["a", "baa", "", "c"]), ByteSet(_)));
        assert!(matches!(pref(vec!["ab", "baa", "", "cb"]), Ac(_, _)));
    }
}

<|endoftext|>


<commit_before>use std::env;
use std::fs::{File, read_dir};
use std::io::Write;
use std::path::PathBuf;
use std::result::Result;
use serde::{Deserialize, Serialize};
use serde_json::ser::Serializer;
use serde_protobuf::descriptor::{Descriptors, MessageDescriptor};
use serde_protobuf::de::Deserializer;
use serde_protobuf::error::{Error, ErrorKind};
use serde_value::Value;
use protobuf::{CodedInputStream, parse_from_reader};
use protobuf::error::ProtobufError;

pub enum PqrsError {
    EofError(String),
    SerdeError(String),
    ProtobufError(String),
}

pub fn process_single(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    match msg_type {
        "guess" => return guess_message(data, out),
        _ => return named_message(data, msg_type, out),
    }
}

fn named_message(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    let mut loc_msg_type = String::new();
    let ch = msg_type.chars().nth(0).unwrap();
    if ch != '.' {
        loc_msg_type.push('.');
    }
    loc_msg_type.push_str(msg_type);

    let (descriptors, _) = get_descriptors(false);

    let stream = CodedInputStream::from_bytes(data);
    let mut deserializer = Deserializer::for_named_message(&descriptors, &loc_msg_type, stream).unwrap();
    let mut serializer = Serializer::new(out);
    match deser(&mut deserializer) {
        Ok(value) => value.serialize(&mut serializer).unwrap(),
        Err(e) => return Err(e),
    }
    Ok(())
}

fn guess_message(data: &[u8], out: &mut Write) -> Result<(), PqrsError> {
    let (descriptors, message_descriptors) = get_descriptors(true);

    let mut serializer = Serializer::new(out);
    for md in message_descriptors {
        let stream = CodedInputStream::from_bytes(data);
        let mut deserializer = Deserializer::new(&descriptors, &md, stream);
        match deser(&mut deserializer) {
            Ok(Value::Map(value)) => {
                let mut unknowns_found = 0;
                for (_, v) in &value {
                    match v {
                        &Value::Unit => unknowns_found += 1,
                        _ => continue,
                    }
                }
                if unknowns_found == 0 {
                    value.serialize(&mut serializer).unwrap();
                }
            },
            Ok(_) => continue,
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

fn deser(deserializer: &mut Deserializer) -> Result<Value, PqrsError> {
    match Value::deserialize(deserializer) {
        Ok(x) => return Ok(x),
        Err(Error(ErrorKind::Protobuf(ProtobufError::WireError(msg)), _)) => {
            if msg == "unexpected EOF" {
                return Err(PqrsError::EofError(msg));
            }
            return Err(PqrsError::ProtobufError(msg));
        },
        Err(e) => return Err(PqrsError::SerdeError(String::from(e.description()))),
    };
}

fn get_descriptors(with_message_descriptors: bool) -> (Descriptors, Vec<MessageDescriptor>) {
    let mut descriptors = Descriptors::new();
    let mut message_descriptors = Vec::new();

    for fdset_path in discover_fdsets() {
        let mut fdset_file = File::open(fdset_path.as_path()).unwrap();
        let fdset_proto = parse_from_reader(&mut fdset_file).unwrap();
        descriptors.add_file_set_proto(&fdset_proto);
        if with_message_descriptors {
            for file_proto in fdset_proto.get_file().iter() {
                for message_proto in file_proto.get_message_type().iter() {
                    message_descriptors.push(MessageDescriptor::from_proto(&fdset_path.to_string_lossy().into_owned().as_str(), message_proto));
                }
            }
        }
    }
    descriptors.resolve_refs();
    (descriptors, message_descriptors)
}

fn discover_fdsets() -> Vec<PathBuf> {
    let mut home = env::home_dir().expect("Could not find $HOME");
    home.push(".pq");

    read_dir(home.as_path()).unwrap()
        .map(|x| {
            x.unwrap().path()
        })
        .collect::<Vec<_>>()
}
<commit_msg>Choose largest BTreeMap as the winning message type<commit_after>use std::env;
use std::fs::{File, read_dir};
use std::io::Write;
use std::path::PathBuf;
use std::result::Result;
use serde::{Deserialize, Serialize};
use serde_json::ser::Serializer;
use serde_protobuf::descriptor::{Descriptors, MessageDescriptor};
use serde_protobuf::de::Deserializer;
use serde_protobuf::error::{Error, ErrorKind};
use serde_value::Value;
use protobuf::{CodedInputStream, parse_from_reader};
use protobuf::error::ProtobufError;

pub enum PqrsError {
    EofError(String),
    SerdeError(String),
    ProtobufError(String),
}

pub fn process_single(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    match msg_type {
        "guess" => return guess_message(data, out),
        _ => return named_message(data, msg_type, out),
    }
}

fn named_message(data: &[u8], msg_type: &str, out: &mut Write) -> Result<(), PqrsError> {
    let mut loc_msg_type = String::new();
    let ch = msg_type.chars().nth(0).unwrap();
    if ch != '.' {
        loc_msg_type.push('.');
    }
    loc_msg_type.push_str(msg_type);

    let (descriptors, _) = get_descriptors(false);

    let stream = CodedInputStream::from_bytes(data);
    let mut deserializer = Deserializer::for_named_message(&descriptors, &loc_msg_type, stream).unwrap();
    let mut serializer = Serializer::new(out);
    match deser(&mut deserializer) {
        Ok(value) => value.serialize(&mut serializer).unwrap(),
        Err(e) => return Err(e),
    }
    Ok(())
}

fn guess_message(data: &[u8], out: &mut Write) -> Result<(), PqrsError> {
    let (descriptors, message_descriptors) = get_descriptors(true);

    let mut serializer = Serializer::new(out);
    let mut contenders = Vec::new();
    for md in message_descriptors {
        let stream = CodedInputStream::from_bytes(data);
        let mut deserializer = Deserializer::new(&descriptors, &md, stream);
        match deser(&mut deserializer) {
            Ok(Value::Map(value)) => {
                let mut unknowns_found = 0;
                for (_, v) in &value {
                    match v {
                        &Value::Unit => unknowns_found += 1,
                        _ => continue,
                    }
                }
                if unknowns_found == 0 {
                    contenders.push(value);
                    //value.serialize(&mut serializer).unwrap();
                }
            },
            Ok(_) => continue,
            Err(e) => return Err(e),
        }
    }
    let contender_max = contenders.iter().max_by_key(|x| x.len());
    contender_max.serialize(&mut serializer).unwrap();
    Ok(())
}

fn deser(deserializer: &mut Deserializer) -> Result<Value, PqrsError> {
    match Value::deserialize(deserializer) {
        Ok(x) => return Ok(x),
        Err(Error(ErrorKind::Protobuf(ProtobufError::WireError(msg)), _)) => {
            if msg == "unexpected EOF" {
                return Err(PqrsError::EofError(msg));
            }
            return Err(PqrsError::ProtobufError(msg));
        },
        Err(e) => return Err(PqrsError::SerdeError(String::from(e.description()))),
    };
}

fn get_descriptors(with_message_descriptors: bool) -> (Descriptors, Vec<MessageDescriptor>) {
    let mut descriptors = Descriptors::new();
    let mut message_descriptors = Vec::new();

    for fdset_path in discover_fdsets() {
        let mut fdset_file = File::open(fdset_path.as_path()).unwrap();
        let fdset_proto = parse_from_reader(&mut fdset_file).unwrap();
        descriptors.add_file_set_proto(&fdset_proto);
        if with_message_descriptors {
            for file_proto in fdset_proto.get_file().iter() {
                for message_proto in file_proto.get_message_type().iter() {
                    message_descriptors.push(MessageDescriptor::from_proto(&fdset_path.to_string_lossy().into_owned().as_str(), message_proto));
                }
            }
        }
    }
    descriptors.resolve_refs();
    (descriptors, message_descriptors)
}

fn discover_fdsets() -> Vec<PathBuf> {
    let mut home = env::home_dir().expect("Could not find $HOME");
    home.push(".pq");

    read_dir(home.as_path()).unwrap()
        .map(|x| {
            x.unwrap().path()
        })
        .collect::<Vec<_>>()
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::cmp;
use std::fmt;
use std::ops::Range;

use std::collections::Bound::{Excluded, Included, Unbounded};

use std::collections::btree_map;
use std::collections::BTreeMap;

#[derive(Clone, Default, PartialEq, PartialOrd)]
pub struct RangeSet {
    inner: BTreeMap<u64, u64>,
}

impl RangeSet {
    pub fn insert(&mut self, item: Range<u64>) {
        let mut start = item.start;
        let mut end = item.end;

        while let Some(r) = self.prev_to(start) {
            if range_contains(&r, start) && range_contains(&r, end) {
                return;
            }

            if range_contains(&item, r.start) &&
               range_contains(&item, r.end) {
                self.inner.remove(&r.start);
                continue;
            }

            if range_contains(&r, start) && !range_contains(&r, end) {
                self.inner.remove(&r.start);
                start = cmp::min(start, r.start);
            }

            if !range_contains(&r, start) && range_contains(&r, end) {
                self.inner.remove(&r.start);
                end = cmp::max(end, r.end);
            }

            if !range_contains(&r, start) && !range_contains(&r, end) {
                break;
            }
        }

        while let Some(r) = self.next_to(start) {
            if range_contains(&r, start) && range_contains(&r, end) {
                return;
            }

            if range_contains(&item, r.start) &&
               range_contains(&item, r.end) {
                self.inner.remove(&r.start);
                continue;
            }

            if range_contains(&r, start) && !range_contains(&r, end) {
                self.inner.remove(&r.start);
                start = cmp::min(start, r.start);
            }

            if !range_contains(&r, start) && range_contains(&r, end) {
                self.inner.remove(&r.start);
                end = cmp::max(end, r.end);
            }

            if !range_contains(&r, start) && !range_contains(&r, end) {
                break;
            }
        }

        self.inner.insert(start, end);
    }

    pub fn push_item(&mut self, item: u64) {
        let r = Range {
            start: item,
            end: item + 1,
        };

        self.insert(r);
    }

    pub fn iter(&self) -> Iter {
        Iter {
            inner: self.inner.iter(),
        }
    }

    pub fn flatten(&self) -> Flatten {
        Flatten {
            inner: self.inner.iter(),
            next: 0,
            end: 0,
        }
    }

    pub fn clear(&mut self) {
        self.inner = BTreeMap::new();
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

    fn prev_to(&self, item: u64) -> Option<Range<u64>> {
        match self.inner.range((Unbounded, Included(item))).rev().next() {
            Some((start, end)) => Some(*start..*end),
            None => None,
        }
    }

    fn next_to(&self, item: u64) -> Option<Range<u64>> {
        match self.inner.range((Included(item), Unbounded)).next() {
            Some((start, end)) => Some(*start..*end),
            None => None,
        }
    }
}

impl fmt::Debug for RangeSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.iter().map(|mut r| { r.end -= 1; r })
                                     .collect::<Vec<Range<u64>>>())
    }
}

pub struct Iter<'a> {
    inner: btree_map::Iter<'a, u64, u64>,
}

impl<'a> Iterator for Iter<'a> {
    type Item = Range<u64>;

    fn next(&mut self) -> Option<Range<u64>> {
        let (&start, &end) = self.inner.next()?;
        Some(start..end)
    }
}

impl<'a> DoubleEndedIterator for Iter<'a> {
    fn next_back(&mut self) -> Option<Range<u64>> {
        let (&start, &end) = self.inner.next_back()?;
        Some(start..end)
    }
}

impl<'a> ExactSizeIterator for Iter<'a> {
    fn len(&self) -> usize {
        self.inner.len()
    }
}

pub struct Flatten<'a> {
    inner: btree_map::Iter<'a, u64, u64>,
    next: u64,
    end: u64,
}

impl<'a> Iterator for Flatten<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        if self.next == self.end {
            let (&start, &end) = self.inner.next()?;

            self.next = start;
            self.end = end;
        }

        let next = self.next;
        self.next += 1;

        Some(next)
    }
}

fn range_contains(r: &Range<u64>, item: u64) -> bool {
    (match Included(r.start) {
        Included(start) => start <= item,
        Excluded(start) => start < item,
        Unbounded => true,
    })
    &&
    (match Included(r.end) {
        Included(end) => item <= end,
        Excluded(end) => item < end,
        Unbounded => true,
    })
}

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

    #[test]
    fn insert_non_overlapping() {
        let mut r = RangeSet::default();
        assert_eq!(r.len(), 0);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[]);

        r.insert(4..7);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);

        r.insert(9..12);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
    }

    #[test]
    fn insert_contained() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(4..7);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(4..6);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(5..6);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(10..11);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(9..11);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
    }

    #[test]
    fn insert_overlapping() {
        let mut r = RangeSet::default();

        r.insert(3..6);
        r.insert(9..12);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 9, 10, 11]);

        r.insert(5..7);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 9, 10, 11]);

        r.insert(10..15);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 9, 10, 11, 12, 13, 14]);

        r.insert(2..5);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14]);

        r.insert(8..10);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14]);

        r.insert(6..10);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]);
    }

    #[test]
    fn insert_overlapping_multi() {
        let mut r = RangeSet::default();

        r.insert(3..6);
        r.insert(16..20);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 16, 17, 18, 19]);

        r.insert(10..11);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 10, 16, 17, 18, 19]);

        r.insert(13..14);
        assert_eq!(r.len(), 4);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 10, 13, 16, 17, 18, 19]);

        r.insert(4..17);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]);
    }

    #[test]
    fn prev_to() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);

        assert_eq!(r.prev_to(2), None);
        assert_eq!(r.prev_to(4), Some(4..7));
        assert_eq!(r.prev_to(15), Some(9..12));
        assert_eq!(r.prev_to(5), Some(4..7));
        assert_eq!(r.prev_to(8), Some(4..7));
    }

    #[test]
    fn next_to() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);

        assert_eq!(r.next_to(2), Some(4..7));
        assert_eq!(r.next_to(12), None);
        assert_eq!(r.next_to(15), None);
        assert_eq!(r.next_to(5), Some(9..12));
        assert_eq!(r.next_to(8), Some(9..12));
    }

    #[test]
    fn push_item() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.push_item(15);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[4, 5, 6, 9, 10, 11, 15]);

        r.push_item(15);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[4, 5, 6, 9, 10, 11, 15]);

        r.push_item(1);
        assert_eq!(r.len(), 4);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 4, 5, 6, 9, 10, 11, 15]);

        r.push_item(12);
        r.push_item(13);
        r.push_item(14);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15]);

        r.push_item(2);
        r.push_item(3);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15]);

        r.push_item(8);
        r.push_item(7);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);
    }
}
<commit_msg>ranges: use .next_back() instead of .rev().next()<commit_after>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::cmp;
use std::fmt;
use std::ops::Range;

use std::collections::Bound::{Excluded, Included, Unbounded};

use std::collections::btree_map;
use std::collections::BTreeMap;

#[derive(Clone, Default, PartialEq, PartialOrd)]
pub struct RangeSet {
    inner: BTreeMap<u64, u64>,
}

impl RangeSet {
    pub fn insert(&mut self, item: Range<u64>) {
        let mut start = item.start;
        let mut end = item.end;

        while let Some(r) = self.prev_to(start) {
            if range_contains(&r, start) && range_contains(&r, end) {
                return;
            }

            if range_contains(&item, r.start) &&
               range_contains(&item, r.end) {
                self.inner.remove(&r.start);
                continue;
            }

            if range_contains(&r, start) && !range_contains(&r, end) {
                self.inner.remove(&r.start);
                start = cmp::min(start, r.start);
            }

            if !range_contains(&r, start) && range_contains(&r, end) {
                self.inner.remove(&r.start);
                end = cmp::max(end, r.end);
            }

            if !range_contains(&r, start) && !range_contains(&r, end) {
                break;
            }
        }

        while let Some(r) = self.next_to(start) {
            if range_contains(&r, start) && range_contains(&r, end) {
                return;
            }

            if range_contains(&item, r.start) &&
               range_contains(&item, r.end) {
                self.inner.remove(&r.start);
                continue;
            }

            if range_contains(&r, start) && !range_contains(&r, end) {
                self.inner.remove(&r.start);
                start = cmp::min(start, r.start);
            }

            if !range_contains(&r, start) && range_contains(&r, end) {
                self.inner.remove(&r.start);
                end = cmp::max(end, r.end);
            }

            if !range_contains(&r, start) && !range_contains(&r, end) {
                break;
            }
        }

        self.inner.insert(start, end);
    }

    pub fn push_item(&mut self, item: u64) {
        let r = Range {
            start: item,
            end: item + 1,
        };

        self.insert(r);
    }

    pub fn iter(&self) -> Iter {
        Iter {
            inner: self.inner.iter(),
        }
    }

    pub fn flatten(&self) -> Flatten {
        Flatten {
            inner: self.inner.iter(),
            next: 0,
            end: 0,
        }
    }

    pub fn clear(&mut self) {
        self.inner = BTreeMap::new();
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

    fn prev_to(&self, item: u64) -> Option<Range<u64>> {
        match self.inner.range((Unbounded, Included(item))).next_back() {
            Some((start, end)) => Some(*start..*end),
            None => None,
        }
    }

    fn next_to(&self, item: u64) -> Option<Range<u64>> {
        match self.inner.range((Included(item), Unbounded)).next() {
            Some((start, end)) => Some(*start..*end),
            None => None,
        }
    }
}

impl fmt::Debug for RangeSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.iter().map(|mut r| { r.end -= 1; r })
                                     .collect::<Vec<Range<u64>>>())
    }
}

pub struct Iter<'a> {
    inner: btree_map::Iter<'a, u64, u64>,
}

impl<'a> Iterator for Iter<'a> {
    type Item = Range<u64>;

    fn next(&mut self) -> Option<Range<u64>> {
        let (&start, &end) = self.inner.next()?;
        Some(start..end)
    }
}

impl<'a> DoubleEndedIterator for Iter<'a> {
    fn next_back(&mut self) -> Option<Range<u64>> {
        let (&start, &end) = self.inner.next_back()?;
        Some(start..end)
    }
}

impl<'a> ExactSizeIterator for Iter<'a> {
    fn len(&self) -> usize {
        self.inner.len()
    }
}

pub struct Flatten<'a> {
    inner: btree_map::Iter<'a, u64, u64>,
    next: u64,
    end: u64,
}

impl<'a> Iterator for Flatten<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        if self.next == self.end {
            let (&start, &end) = self.inner.next()?;

            self.next = start;
            self.end = end;
        }

        let next = self.next;
        self.next += 1;

        Some(next)
    }
}

fn range_contains(r: &Range<u64>, item: u64) -> bool {
    (match Included(r.start) {
        Included(start) => start <= item,
        Excluded(start) => start < item,
        Unbounded => true,
    })
    &&
    (match Included(r.end) {
        Included(end) => item <= end,
        Excluded(end) => item < end,
        Unbounded => true,
    })
}

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

    #[test]
    fn insert_non_overlapping() {
        let mut r = RangeSet::default();
        assert_eq!(r.len(), 0);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[]);

        r.insert(4..7);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);

        r.insert(9..12);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
    }

    #[test]
    fn insert_contained() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(4..7);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(4..6);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(5..6);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(10..11);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.insert(9..11);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
    }

    #[test]
    fn insert_overlapping() {
        let mut r = RangeSet::default();

        r.insert(3..6);
        r.insert(9..12);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 9, 10, 11]);

        r.insert(5..7);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 9, 10, 11]);

        r.insert(10..15);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 9, 10, 11, 12, 13, 14]);

        r.insert(2..5);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14]);

        r.insert(8..10);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14]);

        r.insert(6..10);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]);
    }

    #[test]
    fn insert_overlapping_multi() {
        let mut r = RangeSet::default();

        r.insert(3..6);
        r.insert(16..20);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 16, 17, 18, 19]);

        r.insert(10..11);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 10, 16, 17, 18, 19]);

        r.insert(13..14);
        assert_eq!(r.len(), 4);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 10, 13, 16, 17, 18, 19]);

        r.insert(4..17);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]);
    }

    #[test]
    fn prev_to() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);

        assert_eq!(r.prev_to(2), None);
        assert_eq!(r.prev_to(4), Some(4..7));
        assert_eq!(r.prev_to(15), Some(9..12));
        assert_eq!(r.prev_to(5), Some(4..7));
        assert_eq!(r.prev_to(8), Some(4..7));
    }

    #[test]
    fn next_to() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);

        assert_eq!(r.next_to(2), Some(4..7));
        assert_eq!(r.next_to(12), None);
        assert_eq!(r.next_to(15), None);
        assert_eq!(r.next_to(5), Some(9..12));
        assert_eq!(r.next_to(8), Some(9..12));
    }

    #[test]
    fn push_item() {
        let mut r = RangeSet::default();

        r.insert(4..7);
        r.insert(9..12);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);

        r.push_item(15);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[4, 5, 6, 9, 10, 11, 15]);

        r.push_item(15);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[4, 5, 6, 9, 10, 11, 15]);

        r.push_item(1);
        assert_eq!(r.len(), 4);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 4, 5, 6, 9, 10, 11, 15]);

        r.push_item(12);
        r.push_item(13);
        r.push_item(14);
        assert_eq!(r.len(), 3);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15]);

        r.push_item(2);
        r.push_item(3);
        assert_eq!(r.len(), 2);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15]);

        r.push_item(8);
        r.push_item(7);
        assert_eq!(r.len(), 1);
        assert_eq!(&r.flatten().collect::<Vec<u64>>(),
                   &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);
    }
}
<|endoftext|>