<commit_before>use super::PaletteIndex;
use super::TilePattern;
use super::SCREEN_BUFFER_SIZE;
use super::SCREEN_WIDTH;
use super::SCREEN_HEIGHT;
use super::ppu_reg::PPUReg;
use super::ppu_memory::PPUMemory;
use memory::MemSegment;
use std::cmp;

const TILES_PER_LINE: usize = 34;

#[derive(Debug, Copy, Clone, PartialEq)]
struct TileAttribute {
    bits: u8,
}

impl TileAttribute {
    fn new(bits: u8) -> TileAttribute {
        TileAttribute { bits: bits }
    }

    fn get_palette(&self, x: u16, y: u16) -> u8 {
        let y = y % SCREEN_HEIGHT as u16;
        let mut at = self.bits;
        if y & 0x10 != 0 {
            at >>= 4
        }
        if x & 0x10 != 0 {
            at >>= 2
        }
        at & 0x03
    }

    #[cfg(feature="mousepick")]
    fn get_palette_mask(&self, x: u16, y: u16) -> u8 {
        let y = y % SCREEN_HEIGHT as u16;
        let mut at = 0xFF;
        if y & 0x10 != 0 {
            at &= 0b1111_0000;
        } else {
            at &= 0b0000_1111;
        }

        if x & 0x10 != 0 {
            at &= 0b1100_1100;
        } else {
            at &= 0b0011_0011;
        }
        at
    }
}

impl Default for TileAttribute {
    fn default() -> TileAttribute {
        TileAttribute { bits: 0 }
    }
}

pub struct BackgroundRenderer {
    idx: Box<[[u8; TILES_PER_LINE]; SCREEN_HEIGHT]>,
    tile: Box<[[TilePattern; TILES_PER_LINE]; SCREEN_HEIGHT]>,
    attr: Box<[[u8; TILES_PER_LINE]; SCREEN_HEIGHT]>,
}

fn draw_segment(pattern_line: &[TilePattern; TILES_PER_LINE],
                attr_line: &[u8; TILES_PER_LINE],
                pixel_line: &mut [PaletteIndex],
                fine_x_scroll: usize,
                start: usize,
                stop: usize) {
    for (pixel, item) in pixel_line.iter_mut().enumerate().take(stop).skip(start) {
        let displayed_pixel = pixel + fine_x_scroll;
        render_single_pixel(pattern_line, attr_line, displayed_pixel, item);
    }
}

fn render_single_pixel(pattern_line: &[TilePattern],
                       attr_line: &[u8],
                       displayed_pixel: usize,
                       item: &mut PaletteIndex) {
    let tile_idx = displayed_pixel / 8;
    let pattern = pattern_line[tile_idx];
    let fine_x = displayed_pixel as u32 & 0x07;
    let color_id = pattern.get_color_in_pattern(fine_x);

    let palette_id = attr_line[tile_idx];

    *item = PaletteIndex::from_packed(color_id | palette_id);
}

impl BackgroundRenderer {
    pub fn render(&mut self,
                  buffer: &mut [PaletteIndex; SCREEN_BUFFER_SIZE],
                  start: usize,
                  stop: usize,
                  reg: &PPUReg) {
        let mut current_scanline = start / SCREEN_WIDTH;
        let mut last_scanline_boundary = current_scanline * SCREEN_WIDTH;
        let next_scanline = current_scanline + 1;
        let mut next_scanline_boundary = next_scanline * SCREEN_WIDTH;

        let mut current = start;
        let fine_x_scroll = reg.scroll_x_fine() as usize;
        while current < stop {
            let segment_start = current - last_scanline_boundary;
            let segment_end = cmp::min(next_scanline_boundary, stop) - last_scanline_boundary;

            let pattern_line = &self.tile[current_scanline];
            let attr_line = &self.attr[current_scanline];
            let pixel_line = &mut buffer[last_scanline_boundary..next_scanline_boundary];

            draw_segment(pattern_line,
                         attr_line,
                         pixel_line,
                         fine_x_scroll,
                         segment_start,
                         segment_end);
            current_scanline += 1;
            last_scanline_boundary = next_scanline_boundary;
            current = next_scanline_boundary;
            next_scanline_boundary += SCREEN_WIDTH;
        }
    }

    pub fn run_cycle(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        // Match to update vram addresses
        self.update_vram_address(cyc, sl, reg);
        // VRAM Accesses
        self.read_data(cyc, sl, reg, mem);
    }

    fn update_vram_address(&self, cyc: u16, sl: i16, reg: &mut PPUReg) {
        if sl < 240 {
            match cyc {
                280 if sl == -1 => self.copy_vertical(reg),
                256 => self.increment_y(reg),
                257 => self.copy_horizontal(reg),
                8 | 16 | 24 | 32 | 40 | 48 | 56 | 64 | 72 | 80 | 88 | 96 | 104 | 112 | 120 |
                128 | 136 | 144 | 152 | 160 | 168 | 176 | 184 | 192 | 200 | 208 | 216 | 224 |
                232 | 240 | 248 | 328 | 336 => self.increment_x(reg),
                _ => (),
            }
        }
    }

    fn read_data(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        if sl == -1 {
            match cyc {
                // Fetches for next scanline
                321 | 329 => self.fetch_nametable((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                323 | 331 => self.fetch_attribute((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                325 | 333 => self.fetch_tile_pattern((cyc - 320) / 8, (sl + 1) % 240, reg, mem),

                // The two garbage nametable fetches at the end of every scanline
                337 | 339 => self.garbage_nametable_fetch(reg, mem),
                _ => (),
            }
        } else if sl < 240 {
            match cyc {
                // Normal fetches
                1 | 9 | 17 | 25 | 33 | 41 | 49 | 57 | 65 | 73 | 81 | 89 | 97 | 105 | 113 |
                121 | 129 | 137 | 145 | 153 | 161 | 169 | 177 | 185 | 193 | 201 | 209 | 217 |
                225 | 233 | 241 | 249 => self.fetch_nametable((cyc + 16) / 8, sl, reg, mem),
                3 | 11 | 19 | 27 | 35 | 43 | 51 | 59 | 67 | 75 | 83 | 91 | 99 | 107 | 115 |
                123 | 131 | 139 | 147 | 155 | 163 | 171 | 179 | 187 | 195 | 203 | 211 | 219 |
                227 | 235 | 243 | 251 => self.fetch_attribute((cyc + 16) / 8, sl, reg, mem),
                5 | 13 | 21 | 29 | 37 | 45 | 53 | 61 | 69 | 77 | 85 | 93 | 101 | 109 | 117 |
                125 | 133 | 141 | 149 | 157 | 165 | 173 | 181 | 189 | 197 | 205 | 213 | 221 |
                229 | 237 | 245 | 253 => self.fetch_tile_pattern((cyc + 16) / 8, sl, reg, mem),

                // Fetches for next scanline
                321 | 329 => self.fetch_nametable((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                323 | 331 => self.fetch_attribute((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                325 | 333 => self.fetch_tile_pattern((cyc - 320) / 8, (sl + 1) % 240, reg, mem),

                // The two garbage nametable fetches at the end of every scanline
                337 | 339 => self.garbage_nametable_fetch(reg, mem),
                _ => (),
            }
        }
    }

    fn copy_vertical(&self, reg: &mut PPUReg) {
        let vertical_mask = 0b_111_10_11111_00000;
        reg.v = (reg.v & !vertical_mask) | (reg.t & vertical_mask);
    }

    fn copy_horizontal(&self, reg: &mut PPUReg) {
        let horizontal_mask = 0b_000_01_00000_11111;
        reg.v = (reg.v & !horizontal_mask) | (reg.t & horizontal_mask);
    }

    fn increment_x(&self, reg: &mut PPUReg) {
        if (reg.v & 0x001F) == 31 {
            reg.v &= !0x001F; //clear coarse x
            reg.v ^= 0x0400; //Switch nametable
        } else {
            reg.v += 1; //increment coarse x
        }
    }

    fn increment_y(&self, reg: &mut PPUReg) {
        if (reg.v & 0x7000) != 0x7000 {
            reg.v += 0x1000; //Increment fine Y
        } else {
            reg.v &= !0x7000; //Clear fine Y
            let mut coarse_y = (reg.v & 0x03E0) >> 5;
            if coarse_y == 29 {
                coarse_y = 0;
                reg.v ^= 0x0800; //Switch vertical nametable
            } else if coarse_y == 31 {
                coarse_y = 0; //Clear coarse_y, but do not switch nametable
            } else {
                coarse_y += 1;
            }
            reg.v = (reg.v & !0x03E0) | (coarse_y << 5); //copy coarse_y back into V.
        }
    }

    fn fetch_nametable(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        let nametable_addr = 0x2000 | (reg.v & 0x0FFF);
        self.idx[y as usize][tile_x as usize] = mem.read(nametable_addr);
    }

    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn fetch_attribute(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        let addr = 0x23C0 | (reg.v & 0x0C00) | ((reg.v >> 4) & 0x38) | ((reg.v >> 2) & 0x07);
        let tile_attribute = TileAttribute::new(mem.read(addr));
        let palette_id = tile_attribute.get_palette(
            tile_x * 8 + reg.get_scroll_x(),
            y as u16 + reg.get_scroll_y());
        self.attr[y as usize][tile_x as usize] = palette_id << 3;
    }

    fn fetch_tile_pattern(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        self.tile[y as usize][tile_x as usize] =
            mem.read_tile_pattern(self.idx[y as usize][tile_x as usize],
                                  reg.scroll_y_fine(),
                                  reg.ppuctrl.background_table());
    }

    fn garbage_nametable_fetch(&mut self, reg: &PPUReg, mem: &mut PPUMemory) {
        let nametable_addr = 0x2000 | (reg.v & 0x0FFF);
        mem.read(nametable_addr);
    }

    #[cfg(feature="mousepick")]
    pub fn mouse_pick(&self, reg: &PPUReg, px_x: i32, px_y: i32) {
        let scanline = px_y as usize;
        let tile = (px_x / 8) as usize;
        let tile_id = self.idx[scanline][tile];
        let attr = self.attr[scanline][tile];
        let scrolled_x = px_x as u16 + reg.get_scroll_x() as u16;
        let scrolled_y = px_y as u16 + reg.get_scroll_y() as u16;
        let palette = attr.get_palette(scrolled_x, scrolled_y);
        let palette_mask = attr.get_palette_mask(scrolled_x, scrolled_y);
        println!("{:03}/{:03}: Tile: {:03}, Attribute: {:08b} & {:08b}, Palette: {}",
                 scrolled_x / 8,
                 scrolled_y / 8,
                 tile_id,
                 attr.bits,
                 palette_mask,
                 palette);
    }
}

impl Default for BackgroundRenderer {
    fn default() -> BackgroundRenderer {
        // Work around the 32-element array limitation
        let (idx, tiles, attrs) = unsafe {
            use std::ptr;
            use std::mem;

            let mut idx: [[u8; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();
            let mut tiles: [[TilePattern; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();
            let mut attrs: [[u8; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();

            for element in idx.iter_mut() {
                let idx_line: [u8; TILES_PER_LINE] = [0; TILES_PER_LINE];
                ptr::write(element, idx_line);
            }
            for element in tiles.iter_mut() {
                let tile_line: [TilePattern; TILES_PER_LINE] = [Default::default(); TILES_PER_LINE];
                ptr::write(element, tile_line);
            }
            for element in attrs.iter_mut() {
                let attr_line: [u8; TILES_PER_LINE] = [0; TILES_PER_LINE];
                ptr::write(element, attr_line);
            }

            (idx, tiles, attrs)
        };

        BackgroundRenderer {
            idx: Box::new(idx),
            tile: Box::new(tiles),
            attr: Box::new(attrs),
        }
    }
}
<commit_msg>Fix graphical bug in Zelda.<commit_after>use super::PaletteIndex;
use super::TilePattern;
use super::SCREEN_BUFFER_SIZE;
use super::SCREEN_WIDTH;
use super::SCREEN_HEIGHT;
use super::ppu_reg::PPUReg;
use super::ppu_memory::PPUMemory;
use memory::MemSegment;
use std::cmp;

const TILES_PER_LINE: usize = 34;

#[derive(Debug, Copy, Clone, PartialEq)]
struct TileAttribute {
    bits: u8,
}

impl TileAttribute {
    fn new(bits: u8) -> TileAttribute {
        TileAttribute { bits: bits }
    }

    fn get_palette(&self, x: u16, y: u16) -> u8 {
        let y = y % SCREEN_HEIGHT as u16;
        let mut at = self.bits;
        if y & 0x10 != 0 {
            at >>= 4
        }
        if x & 0x10 != 0 {
            at >>= 2
        }
        at & 0x03
    }

    #[cfg(feature="mousepick")]
    fn get_palette_mask(&self, x: u16, y: u16) -> u8 {
        let y = y % SCREEN_HEIGHT as u16;
        let mut at = 0xFF;
        if y & 0x10 != 0 {
            at &= 0b1111_0000;
        } else {
            at &= 0b0000_1111;
        }

        if x & 0x10 != 0 {
            at &= 0b1100_1100;
        } else {
            at &= 0b0011_0011;
        }
        at
    }
}

impl Default for TileAttribute {
    fn default() -> TileAttribute {
        TileAttribute { bits: 0 }
    }
}

pub struct BackgroundRenderer {
    idx: Box<[[u8; TILES_PER_LINE]; SCREEN_HEIGHT]>,
    tile: Box<[[TilePattern; TILES_PER_LINE]; SCREEN_HEIGHT]>,
    attr: Box<[[u8; TILES_PER_LINE]; SCREEN_HEIGHT]>,
}

fn draw_segment(pattern_line: &[TilePattern; TILES_PER_LINE],
                attr_line: &[u8; TILES_PER_LINE],
                pixel_line: &mut [PaletteIndex],
                fine_x_scroll: usize,
                start: usize,
                stop: usize) {
    for (pixel, item) in pixel_line.iter_mut().enumerate().take(stop).skip(start) {
        let displayed_pixel = pixel + fine_x_scroll;
        render_single_pixel(pattern_line, attr_line, displayed_pixel, item);
    }
}

fn render_single_pixel(pattern_line: &[TilePattern],
                       attr_line: &[u8],
                       displayed_pixel: usize,
                       item: &mut PaletteIndex) {
    let tile_idx = displayed_pixel / 8;
    let pattern = pattern_line[tile_idx];
    let fine_x = displayed_pixel as u32 & 0x07;
    let color_id = pattern.get_color_in_pattern(fine_x);

    let palette_id = attr_line[tile_idx];

    *item = PaletteIndex::from_packed(color_id | palette_id);
}

impl BackgroundRenderer {
    pub fn render(&mut self,
                  buffer: &mut [PaletteIndex; SCREEN_BUFFER_SIZE],
                  start: usize,
                  stop: usize,
                  reg: &PPUReg) {
        let mut current_scanline = start / SCREEN_WIDTH;
        let mut last_scanline_boundary = current_scanline * SCREEN_WIDTH;
        let next_scanline = current_scanline + 1;
        let mut next_scanline_boundary = next_scanline * SCREEN_WIDTH;

        let mut current = start;
        let fine_x_scroll = reg.scroll_x_fine() as usize;
        while current < stop {
            let segment_start = current - last_scanline_boundary;
            let segment_end = cmp::min(next_scanline_boundary, stop) - last_scanline_boundary;

            let pattern_line = &self.tile[current_scanline];
            let attr_line = &self.attr[current_scanline];
            let pixel_line = &mut buffer[last_scanline_boundary..next_scanline_boundary];

            draw_segment(pattern_line,
                         attr_line,
                         pixel_line,
                         fine_x_scroll,
                         segment_start,
                         segment_end);
            current_scanline += 1;
            last_scanline_boundary = next_scanline_boundary;
            current = next_scanline_boundary;
            next_scanline_boundary += SCREEN_WIDTH;
        }
    }

    pub fn run_cycle(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        // Match to update vram addresses
        self.update_vram_address(cyc, sl, reg);
        // VRAM Accesses
        self.read_data(cyc, sl, reg, mem);
    }

    fn update_vram_address(&self, cyc: u16, sl: i16, reg: &mut PPUReg) {
        if sl < 240 {
            match cyc {
                280 if sl == -1 => self.copy_vertical(reg),
                256 => self.increment_y(reg),
                257 => self.copy_horizontal(reg),
                8 | 16 | 24 | 32 | 40 | 48 | 56 | 64 | 72 | 80 | 88 | 96 | 104 | 112 | 120 |
                128 | 136 | 144 | 152 | 160 | 168 | 176 | 184 | 192 | 200 | 208 | 216 | 224 |
                232 | 240 | 248 | 328 | 336 => self.increment_x(reg),
                _ => (),
            }
        }
    }

    fn read_data(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        if sl == -1 {
            match cyc {
                // Fetches for next scanline
                321 | 329 => self.fetch_nametable((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                323 | 331 => self.fetch_attribute((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                325 | 333 => self.fetch_tile_pattern((cyc - 320) / 8, (sl + 1) % 240, reg, mem),

                // The two garbage nametable fetches at the end of every scanline
                337 | 339 => self.garbage_nametable_fetch(reg, mem),
                _ => (),
            }
        } else if sl < 240 {
            match cyc {
                // Normal fetches
                1 | 9 | 17 | 25 | 33 | 41 | 49 | 57 | 65 | 73 | 81 | 89 | 97 | 105 | 113 |
                121 | 129 | 137 | 145 | 153 | 161 | 169 | 177 | 185 | 193 | 201 | 209 | 217 |
                225 | 233 | 241 | 249 => self.fetch_nametable((cyc + 16) / 8, sl, reg, mem),
                3 | 11 | 19 | 27 | 35 | 43 | 51 | 59 | 67 | 75 | 83 | 91 | 99 | 107 | 115 |
                123 | 131 | 139 | 147 | 155 | 163 | 171 | 179 | 187 | 195 | 203 | 211 | 219 |
                227 | 235 | 243 | 251 => self.fetch_attribute((cyc + 16) / 8, sl, reg, mem),
                5 | 13 | 21 | 29 | 37 | 45 | 53 | 61 | 69 | 77 | 85 | 93 | 101 | 109 | 117 |
                125 | 133 | 141 | 149 | 157 | 165 | 173 | 181 | 189 | 197 | 205 | 213 | 221 |
                229 | 237 | 245 | 253 => self.fetch_tile_pattern((cyc + 16) / 8, sl, reg, mem),

                // Fetches for next scanline
                321 | 329 => self.fetch_nametable((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                323 | 331 => self.fetch_attribute((cyc - 320) / 8, (sl + 1) % 240, reg, mem),
                325 | 333 => self.fetch_tile_pattern((cyc - 320) / 8, (sl + 1) % 240, reg, mem),

                // The two garbage nametable fetches at the end of every scanline
                337 | 339 => self.garbage_nametable_fetch(reg, mem),
                _ => (),
            }
        }
    }

    fn copy_vertical(&self, reg: &mut PPUReg) {
        let vertical_mask = 0b_111_10_11111_00000;
        reg.v = (reg.v & !vertical_mask) | (reg.t & vertical_mask);
    }

    fn copy_horizontal(&self, reg: &mut PPUReg) {
        let horizontal_mask = 0b_000_01_00000_11111;
        reg.v = (reg.v & !horizontal_mask) | (reg.t & horizontal_mask);
    }

    fn increment_x(&self, reg: &mut PPUReg) {
        if (reg.v & 0x001F) == 31 {
            reg.v &= !0x001F; //clear coarse x
            reg.v ^= 0x0400; //Switch nametable
        } else {
            reg.v += 1; //increment coarse x
        }
    }

    fn increment_y(&self, reg: &mut PPUReg) {
        if (reg.v & 0x7000) != 0x7000 {
            reg.v += 0x1000; //Increment fine Y
        } else {
            reg.v &= !0x7000; //Clear fine Y
            let mut coarse_y = (reg.v & 0x03E0) >> 5;
            if coarse_y == 29 {
                coarse_y = 0;
                reg.v ^= 0x0800; //Switch vertical nametable
            } else if coarse_y == 31 {
                coarse_y = 0; //Clear coarse_y, but do not switch nametable
            } else {
                coarse_y += 1;
            }
            reg.v = (reg.v & !0x03E0) | (coarse_y << 5); //copy coarse_y back into V.
        }
    }

    fn fetch_nametable(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        let nametable_addr = 0x2000 | (reg.v & 0x0FFF);
        self.idx[y as usize][tile_x as usize] = mem.read(nametable_addr);
    }

    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn fetch_attribute(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        let addr = 0x23C0 | (reg.v & 0x0C00) | ((reg.v >> 4) & 0x38) | ((reg.v >> 2) & 0x07);
        let tile_attribute = TileAttribute::new(mem.read(addr));
        let palette_id = tile_attribute.get_palette(
            tile_x * 8 + reg.get_scroll_x(),
            y as u16 + reg.get_scroll_y());
        self.attr[y as usize][tile_x as usize] = palette_id << 2;
    }

    fn fetch_tile_pattern(&mut self, tile_x: u16, y: i16, reg: &PPUReg, mem: &mut PPUMemory) {
        self.tile[y as usize][tile_x as usize] =
            mem.read_tile_pattern(self.idx[y as usize][tile_x as usize],
                                  reg.scroll_y_fine(),
                                  reg.ppuctrl.background_table());
    }

    fn garbage_nametable_fetch(&mut self, reg: &PPUReg, mem: &mut PPUMemory) {
        let nametable_addr = 0x2000 | (reg.v & 0x0FFF);
        mem.read(nametable_addr);
    }

    #[cfg(feature="mousepick")]
    pub fn mouse_pick(&self, reg: &PPUReg, px_x: i32, px_y: i32) {
        let scanline = px_y as usize;
        let tile = (px_x / 8) as usize;
        let tile_id = self.idx[scanline][tile];
        let attr = self.attr[scanline][tile];
        let scrolled_x = px_x as u16 + reg.get_scroll_x() as u16;
        let scrolled_y = px_y as u16 + reg.get_scroll_y() as u16;
        let palette = attr.get_palette(scrolled_x, scrolled_y);
        let palette_mask = attr.get_palette_mask(scrolled_x, scrolled_y);
        println!("{:03}/{:03}: Tile: {:03}, Attribute: {:08b} & {:08b}, Palette: {}",
                 scrolled_x / 8,
                 scrolled_y / 8,
                 tile_id,
                 attr.bits,
                 palette_mask,
                 palette);
    }
}

impl Default for BackgroundRenderer {
    fn default() -> BackgroundRenderer {
        // Work around the 32-element array limitation
        let (idx, tiles, attrs) = unsafe {
            use std::ptr;
            use std::mem;

            let mut idx: [[u8; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();
            let mut tiles: [[TilePattern; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();
            let mut attrs: [[u8; TILES_PER_LINE]; SCREEN_HEIGHT] = mem::uninitialized();

            for element in idx.iter_mut() {
                let idx_line: [u8; TILES_PER_LINE] = [0; TILES_PER_LINE];
                ptr::write(element, idx_line);
            }
            for element in tiles.iter_mut() {
                let tile_line: [TilePattern; TILES_PER_LINE] = [Default::default(); TILES_PER_LINE];
                ptr::write(element, tile_line);
            }
            for element in attrs.iter_mut() {
                let attr_line: [u8; TILES_PER_LINE] = [0; TILES_PER_LINE];
                ptr::write(element, attr_line);
            }

            (idx, tiles, attrs)
        };

        BackgroundRenderer {
            idx: Box::new(idx),
            tile: Box::new(tiles),
            attr: Box::new(attrs),
        }
    }
}
<|endoftext|>


<commit_before>use crate::providers::gcp;
use crate::providers::MetadataProvider;
use mockito;

#[test]
fn basic_hostname() {
    let ep = "/instance/hostname";
    let hostname = "test-hostname";

    let mut provider = gcp::GcpProvider::try_new().unwrap();
    provider.client = provider.client.max_attempts(1);

    provider.hostname().unwrap_err();

    let _m = mockito::mock("GET", ep).with_status(503).create();
    provider.hostname().unwrap_err();

    let _m = mockito::mock("GET", ep)
        .with_status(200)
        .with_body(hostname)
        .create();
    let v = provider.hostname().unwrap();
    assert_eq!(v, Some(hostname.to_string()));

    let _m = mockito::mock("GET", ep).with_status(404).create();
    let v = provider.hostname().unwrap();
    assert_eq!(v, None);
}
<commit_msg>tests/gcp: avoid flake<commit_after>use crate::providers::gcp;
use crate::providers::MetadataProvider;
use mockito;

#[test]
fn basic_hostname() {
    let ep = "/instance/hostname";
    let hostname = "test-hostname";

    let mut provider = gcp::GcpProvider::try_new().unwrap();
    provider.client = provider.client.max_attempts(1);

    let _m = mockito::mock("GET", ep).with_status(503).create();
    provider.hostname().unwrap_err();

    let _m = mockito::mock("GET", ep)
        .with_status(200)
        .with_body(hostname)
        .create();
    let v = provider.hostname().unwrap();
    assert_eq!(v, Some(hostname.to_string()));

    let _m = mockito::mock("GET", ep).with_status(404).create();
    let v = provider.hostname().unwrap();
    assert_eq!(v, None);

    mockito::reset();
    provider.hostname().unwrap_err();
}
<|endoftext|>


<commit_before>use std::env;
use std::io::{stdout, stderr};

use argparse::{ArgumentParser, Store, StoreTrue};
use unshare::{Command, Namespace};

use launcher::Context;
use launcher::wrap::Wrapper;
use options::build_mode::BuildMode;
use process_util::{capture_fd3, copy_env_vars, squash_stdio};
use container::util::version_from_symlink;

pub fn build_container(context: &Context, name: &String, mode: BuildMode)
    -> Result<String, String>
{
    use options::build_mode::BuildMode::*;
    let ver = match mode {
        Normal => try!(build_internal(context, name, &[])),
        NoImage => try!(build_internal(context, name,
            &[String::from("--no-image-download")])),
        NoBuild => format!("{}.{}", &name, try!(get_version(context, &name))),
        NoVersion => {
            try!(version_from_symlink(format!(".vagga/{}", name)))
        }
    };
    Ok(ver)
}

/// Similar to build_container but never actually builds
pub fn get_version(context: &Context, name: &str) -> Result<String, String> {
    let mut cmd: Command = Wrapper::new(None, &context.settings);
    cmd.arg("_version_hash");
    cmd.arg("--short");
    cmd.arg("--fd3");
    cmd.arg(name);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &context.settings);
    // TODO(tailhook) move these to copy_env_vars, or at least reuse in build?
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    if let Ok(x) = env::var("HOME") {
        cmd.env("_VAGGA_HOME", x);
    }
    if let Some(x) = env::var_os("VAGGA_SETTINGS") {
        cmd.env("VAGGA_SETTINGS", x);
    }
    cmd.unshare(
        [Namespace::Mount, Namespace::Ipc, Namespace::Pid].iter().cloned());
    try!(cmd.map_users_for(
        try!(context.config.get_container(name)),
        &context.settings));

    capture_fd3(cmd)
    .and_then(|x| String::from_utf8(x)
                  .map_err(|e| format!("Can't decode version: {}", e)))
}

fn build_internal(context: &Context, name: &str, args: &[String])
    -> Result<String, String>
{
    let mut cmd = Wrapper::new(None, &context.settings);
    try!(squash_stdio(&mut cmd));
    cmd.arg0("vagga_wrapper");
    cmd.arg("_build");
    cmd.arg(name);
    cmd.args(&args);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &context.settings);
    // TODO(tailhook) move these to copy_env_vars, or at least reuse in ver?
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    if let Ok(x) = env::var("HOME") {
        cmd.env("_VAGGA_HOME", x);
    }
    if let Some(x) = env::var_os("VAGGA_SETTINGS") {
        cmd.env("VAGGA_SETTINGS", x);
    }
    cmd.unshare(
        [Namespace::Mount, Namespace::Ipc, Namespace::Pid].iter().cloned());
    try!(cmd.map_users_for(
        try!(context.config.get_container(name)),
        &context.settings));

    capture_fd3(cmd)
    .and_then(|x| String::from_utf8(x)
                  .map_err(|e| format!("Can't decode version: {}", e)))
}

pub fn build_command(context: &Context, args: Vec<String>)
    -> Result<i32, String>
{
    let mut name: String = "".to_string();
    let mut force: bool = false;
    {
        let mut cmdline = args.clone();
        cmdline.insert(0, "vagga _build".to_string());
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Internal vagga tool to setup basic system sandbox
            ");
        ap.refer(&mut name)
            .add_argument("container_name", Store,
                "Container name to build");
        ap.refer(&mut force)
            .add_option(&["--force"], StoreTrue,
                "Force build even if container is considered up to date");
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let mut args = Vec::new();
    if force {
        args.push("--force".to_string());
    }

    build_internal(context, &name, &args)
    .map(|v| debug!("Container {:?} build with version {:?}", name, v))
    .map(|()| 0)
}
<commit_msg>_build command should respect --no-image-download option (#319)<commit_after>use std::env;
use std::io::{stdout, stderr};

use argparse::{ArgumentParser, Store, StoreTrue};
use unshare::{Command, Namespace};

use launcher::Context;
use launcher::wrap::Wrapper;
use options::build_mode::BuildMode;
use process_util::{capture_fd3, copy_env_vars, squash_stdio};
use container::util::version_from_symlink;

pub fn build_container(context: &Context, name: &String, mode: BuildMode)
    -> Result<String, String>
{
    use options::build_mode::BuildMode::*;
    let ver = match mode {
        Normal => try!(build_internal(context, name, &[])),
        NoImage => try!(build_internal(context, name,
            &["--no-image-download".to_string()])),
        NoBuild => format!("{}.{}", &name, try!(get_version(context, &name))),
        NoVersion => {
            try!(version_from_symlink(format!(".vagga/{}", name)))
        }
    };
    Ok(ver)
}

/// Similar to build_container but never actually builds
pub fn get_version(context: &Context, name: &str) -> Result<String, String> {
    let mut cmd: Command = Wrapper::new(None, &context.settings);
    cmd.arg("_version_hash");
    cmd.arg("--short");
    cmd.arg("--fd3");
    cmd.arg(name);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &context.settings);
    // TODO(tailhook) move these to copy_env_vars, or at least reuse in build?
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    if let Ok(x) = env::var("HOME") {
        cmd.env("_VAGGA_HOME", x);
    }
    if let Some(x) = env::var_os("VAGGA_SETTINGS") {
        cmd.env("VAGGA_SETTINGS", x);
    }
    cmd.unshare(
        [Namespace::Mount, Namespace::Ipc, Namespace::Pid].iter().cloned());
    try!(cmd.map_users_for(
        try!(context.config.get_container(name)),
        &context.settings));

    capture_fd3(cmd)
    .and_then(|x| String::from_utf8(x)
                  .map_err(|e| format!("Can't decode version: {}", e)))
}

fn build_internal(context: &Context, name: &str, args: &[String])
    -> Result<String, String>
{
    let mut cmd = Wrapper::new(None, &context.settings);
    try!(squash_stdio(&mut cmd));
    cmd.arg0("vagga_wrapper");
    cmd.arg("_build");
    cmd.arg(name);
    cmd.args(&args);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &context.settings);
    // TODO(tailhook) move these to copy_env_vars, or at least reuse in ver?
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    if let Ok(x) = env::var("HOME") {
        cmd.env("_VAGGA_HOME", x);
    }
    if let Some(x) = env::var_os("VAGGA_SETTINGS") {
        cmd.env("VAGGA_SETTINGS", x);
    }
    cmd.unshare(
        [Namespace::Mount, Namespace::Ipc, Namespace::Pid].iter().cloned());
    try!(cmd.map_users_for(
        try!(context.config.get_container(name)),
        &context.settings));

    capture_fd3(cmd)
    .and_then(|x| String::from_utf8(x)
                  .map_err(|e| format!("Can't decode version: {}", e)))
}

pub fn build_command(context: &Context, args: Vec<String>)
    -> Result<i32, String>
{
    let mut name: String = "".to_string();
    let mut force: bool = false;
    {
        let mut cmdline = args.clone();
        cmdline.insert(0, "vagga _build".to_string());
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Internal vagga tool to setup basic system sandbox
            ");
        ap.refer(&mut name)
            .add_argument("container_name", Store,
                "Container name to build");
        ap.refer(&mut force)
            .add_option(&["--force"], StoreTrue,
                "Force build even if container is considered up to date");
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let mut args = Vec::new();
    if force {
        args.push("--force".to_string());
    }
    if let BuildMode::NoImage = context.build_mode {
        args.push("--no-image-download".to_string());
    }

    build_internal(context, &name, &args)
    .map(|v| debug!("Container {:?} build with version {:?}", name, v))
    .map(|()| 0)
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
#[cfg(unix)]
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;

#[cfg(unix)]
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };

    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;

    let bold: fn(w: &str) -> String = if should_colorize {
        |w| Style::new().bold().paint(w).to_string()
    } else {
        |w| w.to_string()
    };

    println!("{} {}", bold("Vendor:"), info.vendor);
    println!("{} {}", bold("Product:"), info.product);
    println!("{} {}", bold("Version:"), info.version);
    println!("{} {}", bold("URL:"), info.url);
    println!("{}", bold("Interfaces:"));

    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    columns: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        columns.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, columns.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            array_brackets: Style::new(),
            object_brackets: Style::new(),
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

#[cfg(unix)]
fn varlink_bridge(address: Option<&str>) -> Result<()> {
    use std::os::unix::io::{AsRawFd, FromRawFd};

    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn main() -> Result<()> {
    // TODO: windows
    // #[cfg(windows)]
    // let _enabled = colored_json::enable_ansi_support();

    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("cols")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("cols")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );

    #[cfg(unix)]
    {
        app = app.arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        );
    }

    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        #[cfg(unix)]
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");
            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<commit_msg>varlink-cli/src/main.rs: windows: colored_json::enable_ansi_support();<commit_after>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
#[cfg(unix)]
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;

#[cfg(unix)]
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };

    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;

    let bold: fn(w: &str) -> String = if should_colorize {
        |w| Style::new().bold().paint(w).to_string()
    } else {
        |w| w.to_string()
    };

    println!("{} {}", bold("Vendor:"), info.vendor);
    println!("{} {}", bold("Product:"), info.product);
    println!("{} {}", bold("Version:"), info.version);
    println!("{} {}", bold("URL:"), info.url);
    println!("{}", bold("Interfaces:"));

    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    columns: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        columns.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, columns.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            array_brackets: Style::new(),
            object_brackets: Style::new(),
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

#[cfg(unix)]
fn varlink_bridge(address: Option<&str>) -> Result<()> {
    use std::os::unix::io::{AsRawFd, FromRawFd};

    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn main() -> Result<()> {
    #[cfg(windows)]
    let _enabled = colored_json::enable_ansi_support();

    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("cols")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("cols")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );

    #[cfg(unix)]
    {
        app = app.arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        );
    }

    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        #[cfg(unix)]
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");
            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use decoders::basics::*;

#[derive(Debug, Copy, Clone)]
pub struct BitPumpLSB<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u32, // This benchmarks better than u64 and u128
  nbits: u32,
}

impl<'a> BitPumpLSB<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpLSB {
    BitPumpLSB {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpMSB<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64, // This benchmarks better than u32 and u128
  nbits: u32,
}

impl<'a> BitPumpMSB<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpMSB {
    BitPumpMSB {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpMSB32<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPumpMSB32<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpMSB32 {
    BitPumpMSB32 {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }

  pub fn get_pos(&self) -> usize {
    self.pos - ((self.nbits >> 3) as usize)
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpJPEG<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
  finished: bool,
}

impl<'a> BitPumpJPEG<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpJPEG {
    BitPumpJPEG {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
      finished: false,
    }
  }
}

pub trait BitPump {
  fn peek_bits(&mut self, num: u32) -> u32;
  fn consume_bits(&mut self, num: u32);

  fn get_bits(&mut self, num: u32) -> u32 {
    if num == 0 {
      return 0
    }

    let val = self.peek_bits(num);
    self.consume_bits(num);

    val
  }

  fn peek_ibits(&mut self, num: u32) -> i32 {
    self.peek_bits(num) as i32
  }

  fn get_ibits(&mut self, num: u32) -> i32 {
    self.get_bits(num) as i32
  }

  // Sign extend ibits
  fn get_ibits_sextended(&mut self, num: u32) -> i32 {
    let val = self.get_ibits(num);
    val.wrapping_shl(32 - num).wrapping_shr(32 - num)
  }
}

impl<'a> BitPump for BitPumpLSB<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u32 = LEu16(self.buffer, self.pos) as u32;
      self.bits = ((inbits << 16) | (self.bits << (16-self.nbits))) >> (16-self.nbits);
      self.pos += 2;
      self.nbits += 16;
    }
    (self.bits & (0x0ffffu32 >> (16-num))) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits >>= num;
  }
}

impl<'a> BitPump for BitPumpMSB<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u64 = BEu32(self.buffer, self.pos) as u64;
      self.bits = (self.bits << 32) | inbits;
      self.pos += 4;
      self.nbits += 32;
    }
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl<'a> BitPump for BitPumpMSB32<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u64 = LEu32(self.buffer, self.pos) as u64;
      self.bits = (self.bits << 32) | inbits;
      self.pos += 4;
      self.nbits += 32;
    }
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl<'a> BitPump for BitPumpJPEG<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits && !self.finished {
      if self.pos < self.buffer.len()-4 &&
         self.buffer[self.pos+0] != 0xff &&
         self.buffer[self.pos+1] != 0xff &&
         self.buffer[self.pos+2] != 0xff &&
         self.buffer[self.pos+3] != 0xff {
        let inbits: u64 = BEu32(self.buffer, self.pos) as u64;
        self.bits = (self.bits << 32) | inbits;
        self.pos += 4;
        self.nbits += 32;
      } else {
        // Read 32 bits the hard way
        let mut read_bytes = 0;
        while read_bytes < 4 && !self.finished {
          let byte = {
            if self.pos >= self.buffer.len() {
              self.finished = true;
              0
            } else {
              let nextbyte = self.buffer[self.pos];
              if nextbyte != 0xff {
                nextbyte
              } else if self.buffer[self.pos+1] == 0x00 {
                self.pos += 1; // Skip the extra byte used to mark 255
                nextbyte
              } else {
                self.finished = true;
                0
              }
            }
          };
          self.bits = (self.bits << 8) | (byte as u64);
          self.pos += 1;
          self.nbits += 8;
          read_bytes += 1;
        }
      }
    }
    if num > self.nbits && self.finished {
      // Stuff with zeroes to not fail to read
      self.bits <<= 32;
      self.nbits += 32;
    }

    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

#[derive(Debug, Copy, Clone)]
pub struct ByteStream<'a> {
  buffer: &'a [u8],
  pos: usize,
  endian: Endian,
}

impl<'a> ByteStream<'a> {
  pub fn new(src: &'a [u8], endian: Endian) -> ByteStream {
    ByteStream {
      buffer: src,
      pos: 0,
      endian: endian,
    }
  }

  pub fn get_pos(&self) -> usize { self.pos }

  pub fn peek_u8(&self) -> u8 { self.buffer[self.pos] }
  pub fn get_u8(&mut self) -> u8 {
    let val = self.peek_u8();
    self.pos += 1;
    val
  }

  pub fn peek_u16(&self) -> u16 { self.endian.ru16(self.buffer, self.pos) }
  pub fn get_u16(&mut self) -> u16 {
    let val = self.peek_u16();
    self.pos += 2;
    val
  }

//  pub fn peek_u32(&self) -> u32 { self.endian.ru32(self.buffer, self.pos) }
//  pub fn get_u32(&mut self) -> u32 {
//    let val = self.peek_u32();
//    self.pos += 4;
//    val
//  }

  pub fn consume_bytes(&mut self, num: usize) {
    self.pos += num
  }

  pub fn skip_to_marker(&mut self) -> Result<usize, String> {
    let mut skip_count = 0;
    while !(self.buffer[self.pos] == 0xFF &&
            self.buffer[self.pos+1] != 0 &&
            self.buffer[self.pos+1] != 0xFF) {
      self.pos += 1;
      skip_count += 1;
      if self.pos >= self.buffer.len() {
        return Err("No marker found inside rest of buffer".to_string())
      }
    }
    self.pos += 1; // Make the next byte the marker
    Ok(skip_count+1)
  }
}
<commit_msg>Go back to u64 in BitPumpLSB<commit_after>use decoders::basics::*;

#[derive(Debug, Copy, Clone)]
pub struct BitPumpLSB<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPumpLSB<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpLSB {
    BitPumpLSB {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpMSB<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPumpMSB<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpMSB {
    BitPumpMSB {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpMSB32<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPumpMSB32<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpMSB32 {
    BitPumpMSB32 {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }

  pub fn get_pos(&self) -> usize {
    self.pos - ((self.nbits >> 3) as usize)
  }
}

#[derive(Debug, Copy, Clone)]
pub struct BitPumpJPEG<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
  finished: bool,
}

impl<'a> BitPumpJPEG<'a> {
  pub fn new(src: &'a [u8]) -> BitPumpJPEG {
    BitPumpJPEG {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
      finished: false,
    }
  }
}

pub trait BitPump {
  fn peek_bits(&mut self, num: u32) -> u32;
  fn consume_bits(&mut self, num: u32);

  fn get_bits(&mut self, num: u32) -> u32 {
    if num == 0 {
      return 0
    }

    let val = self.peek_bits(num);
    self.consume_bits(num);

    val
  }

  fn peek_ibits(&mut self, num: u32) -> i32 {
    self.peek_bits(num) as i32
  }

  fn get_ibits(&mut self, num: u32) -> i32 {
    self.get_bits(num) as i32
  }

  // Sign extend ibits
  fn get_ibits_sextended(&mut self, num: u32) -> i32 {
    let val = self.get_ibits(num);
    val.wrapping_shl(32 - num).wrapping_shr(32 - num)
  }
}

impl<'a> BitPump for BitPumpLSB<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u64 = LEu32(self.buffer, self.pos) as u64;
      self.bits = ((inbits << 32) | (self.bits << (32-self.nbits))) >> (32-self.nbits);
      self.pos += 4;
      self.nbits += 32;
    }
    (self.bits & (0x0ffffffffu64 >> (32-num))) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits >>= num;
  }
}

impl<'a> BitPump for BitPumpMSB<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u64 = BEu32(self.buffer, self.pos) as u64;
      self.bits = (self.bits << 32) | inbits;
      self.pos += 4;
      self.nbits += 32;
    }
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl<'a> BitPump for BitPumpMSB32<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      let inbits: u64 = LEu32(self.buffer, self.pos) as u64;
      self.bits = (self.bits << 32) | inbits;
      self.pos += 4;
      self.nbits += 32;
    }
    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

impl<'a> BitPump for BitPumpJPEG<'a> {
  fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits && !self.finished {
      if self.pos < self.buffer.len()-4 &&
         self.buffer[self.pos+0] != 0xff &&
         self.buffer[self.pos+1] != 0xff &&
         self.buffer[self.pos+2] != 0xff &&
         self.buffer[self.pos+3] != 0xff {
        let inbits: u64 = BEu32(self.buffer, self.pos) as u64;
        self.bits = (self.bits << 32) | inbits;
        self.pos += 4;
        self.nbits += 32;
      } else {
        // Read 32 bits the hard way
        let mut read_bytes = 0;
        while read_bytes < 4 && !self.finished {
          let byte = {
            if self.pos >= self.buffer.len() {
              self.finished = true;
              0
            } else {
              let nextbyte = self.buffer[self.pos];
              if nextbyte != 0xff {
                nextbyte
              } else if self.buffer[self.pos+1] == 0x00 {
                self.pos += 1; // Skip the extra byte used to mark 255
                nextbyte
              } else {
                self.finished = true;
                0
              }
            }
          };
          self.bits = (self.bits << 8) | (byte as u64);
          self.pos += 1;
          self.nbits += 8;
          read_bytes += 1;
        }
      }
    }
    if num > self.nbits && self.finished {
      // Stuff with zeroes to not fail to read
      self.bits <<= 32;
      self.nbits += 32;
    }

    (self.bits >> (self.nbits-num)) as u32
  }

  fn consume_bits(&mut self, num: u32) {
    self.nbits -= num;
    self.bits &= (1 << self.nbits) - 1;
  }
}

#[derive(Debug, Copy, Clone)]
pub struct ByteStream<'a> {
  buffer: &'a [u8],
  pos: usize,
  endian: Endian,
}

impl<'a> ByteStream<'a> {
  pub fn new(src: &'a [u8], endian: Endian) -> ByteStream {
    ByteStream {
      buffer: src,
      pos: 0,
      endian: endian,
    }
  }

  pub fn get_pos(&self) -> usize { self.pos }

  pub fn peek_u8(&self) -> u8 { self.buffer[self.pos] }
  pub fn get_u8(&mut self) -> u8 {
    let val = self.peek_u8();
    self.pos += 1;
    val
  }

  pub fn peek_u16(&self) -> u16 { self.endian.ru16(self.buffer, self.pos) }
  pub fn get_u16(&mut self) -> u16 {
    let val = self.peek_u16();
    self.pos += 2;
    val
  }

//  pub fn peek_u32(&self) -> u32 { self.endian.ru32(self.buffer, self.pos) }
//  pub fn get_u32(&mut self) -> u32 {
//    let val = self.peek_u32();
//    self.pos += 4;
//    val
//  }

  pub fn consume_bytes(&mut self, num: usize) {
    self.pos += num
  }

  pub fn skip_to_marker(&mut self) -> Result<usize, String> {
    let mut skip_count = 0;
    while !(self.buffer[self.pos] == 0xFF &&
            self.buffer[self.pos+1] != 0 &&
            self.buffer[self.pos+1] != 0xFF) {
      self.pos += 1;
      skip_count += 1;
      if self.pos >= self.buffer.len() {
        return Err("No marker found inside rest of buffer".to_string())
      }
    }
    self.pos += 1; // Make the next byte the marker
    Ok(skip_count+1)
  }
}
<|endoftext|>