<commit_before>use failure::Error;
use lucet_runtime_internals::module::DlModule;
use lucetc::Lucetc;
use std::fs::File;
use std::io::prelude::*;
use std::sync::Arc;
use tempdir::TempDir;

pub fn stack_testcase(num_locals: usize) -> Result<Arc<DlModule>, Error> {
    let workdir = TempDir::new("stack_test").expect("create working directory");

    let wasm_path = workdir.path().join("out.wasm");

    let mut wasm_file = File::create(&wasm_path)?;
    wasm_file.write_all(generate_test_wat(num_locals).as_bytes())?;

    let native_build = Lucetc::new(wasm_path)?;

    let so_file = workdir.path().join("out.so");

    native_build.shared_object_file(so_file.clone())?;

    let dlmodule = DlModule::load(so_file)?;

    Ok(dlmodule)
}

fn generate_test_wat(num_locals: usize) -> String {
    assert!(num_locals > 2);

    let mut module =
        "(module (func $localpalooza (export \"localpalooza\") (param i32) (result i32)\n"
            .to_owned();

    // Declare locals:
    module.push_str("(local ");
    for _ in 0..num_locals {
        module.push_str("i32 ");
    }
    module.push_str(")\n");

    // Use each local for the first time:
    for i in 1..num_locals {
        module.push_str(&format!(
            "(set_local {} (i32.add (get_local {}) (get_local {})))\n",
            i,
            i - 1,
            i
        ));
    }

    // Use each local for a second time, so they get pushed to the stack between uses:
    for i in 2..(num_locals - 1) {
        module.push_str(&format!(
            "(set_local {} (i32.add (get_local {}) (get_local {})))\n",
            i,
            i - 1,
            i
        ));
    }

    // Keep locals alive across a recursive call. Make as many recursive calls as the first
    // argument to the func:
    module.push_str("(if (get_local 0)\n");
    module.push_str(&format!(
        "  (then (set_local {} (i32.add (get_local {})\n",
        num_locals - 1,
        num_locals - 2
    ));
    module.push_str("      (call $localpalooza (i32.sub (get_local 0) (i32.const 1))))))\n");
    module.push_str(&format!(
        "  (else (set_local {} (i32.add (get_local {}) (get_local {})))))\n",
        num_locals - 1,
        num_locals - 2,
        num_locals - 3
    ));

    module.push_str(&format!("(get_local {})\n", num_locals - 1));
    module.push_str("))\n");
    module
}

#[macro_export]
macro_rules! stack_tests {
    ( $TestRegion:path ) => {
        use lucet_runtime::{
            DlModule, Error, InstanceHandle, Limits, Region, TrapCodeType, UntypedRetVal, Val,
        };
        use std::sync::Arc;
        use $TestRegion as TestRegion;
        use $crate::stack::stack_testcase;

        fn run(
            module: Arc<DlModule>,
            recursion_depth: libc::c_int,
        ) -> Result<UntypedRetVal, Error> {
            let region = TestRegion::create(1, &Limits::default()).expect("region can be created");
            let mut inst = region
                .new_instance(module)
                .expect("instance can be created");

            inst.run(b"localpalooza", &[recursion_depth.into()])
        }

        fn expect_ok(module: Arc<DlModule>, recursion_depth: libc::c_int) {
            assert!(run(module, recursion_depth).is_ok());
        }

        fn expect_stack_overflow(
            module: Arc<DlModule>,
            recursion_depth: libc::c_int,
            probestack: bool,
        ) {
            match run(module, recursion_depth) {
                Err(Error::RuntimeFault(details)) => {
                    // We should get a nonfatal trap due to the stack overflow.
                    assert_eq!(details.fatal, false);
                    assert_eq!(details.trapcode.ty, TrapCodeType::StackOverflow);
                    if probestack {
                        // Make sure we overflowed in the stack probe as expected
                        //
                        // TODO: this no longer differentiates between different stack overflow
                        // sites after moving the stack probe into lucetc; figure out a way to
                        // provide that information or just wait till we can do the stack probe as a
                        // global symbol again
                        let addr_details =
                            details.rip_addr_details.expect("can look up addr details");
                        assert!(addr_details.in_module_code);
                    }
                }
                res => panic!("unexpected result: {:?}", res),
            }
        }

        #[test]
        fn expect_ok_locals3_1() {
            expect_ok(stack_testcase(3).expect("generate stack_testcase 3"), 1);
        }

        // The test with 64 locals should cause a stack overflow on the 481st recursion.  The trap
        // table knows about all of the instructions in the function that manipulate the stack, so
        // the catch mechanism for this is the usual one.

        #[test]
        fn expect_ok_locals64_1() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 1);
        }

        #[test]
        fn expect_ok_locals64_2() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 2);
        }

        #[test]
        fn expect_ok_locals64_480() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 480);
        }

        #[test]
        fn expect_stack_overflow_locals64_481() {
            expect_stack_overflow(
                stack_testcase(64).expect("generate stack_testcase 64"),
                481,
                true,
            );
        }

        // 1050 locals is about 1 page (4k) on the stack - just enough for Cranelift to use probestack to grow
        // the stack. The 31st recursion should cause a stack overflow.

        #[test]
        fn expect_ok_locals_1page_1() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                1,
            );
        }

        #[test]
        fn expect_ok_locals_1page_2() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                2,
            );
        }

        #[test]
        fn expect_ok_locals_1page_30() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                30,
            );
        }

        #[test]
        fn expect_stack_overflow_locals_1page_31() {
            expect_stack_overflow(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                31,
                true,
            );
        }

        // This test has 5000 locals - over 4 pages worth. Cranelift will use probestack here as well. The
        // 6th recursion should cause a stack overflow.

        #[test]
        fn expect_ok_locals_multipage_1() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                1,
            );
        }

        #[test]
        fn expect_ok_locals_multipage_2() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                2,
            );
        }

        #[test]
        fn expect_ok_locals_multipage_5() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                5,
            );
        }

        #[test]
        fn expect_stack_overflow_locals_multipage_6() {
            expect_stack_overflow(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                6,
                true,
            );
        }
    };
}
<commit_msg>stack tests: no need for c_int type anymore<commit_after>use failure::Error;
use lucet_runtime_internals::module::DlModule;
use lucetc::Lucetc;
use std::fs::File;
use std::io::prelude::*;
use std::sync::Arc;
use tempdir::TempDir;

pub fn stack_testcase(num_locals: usize) -> Result<Arc<DlModule>, Error> {
    let workdir = TempDir::new("stack_test").expect("create working directory");

    let wasm_path = workdir.path().join("out.wasm");

    let mut wasm_file = File::create(&wasm_path)?;
    wasm_file.write_all(generate_test_wat(num_locals).as_bytes())?;

    let native_build = Lucetc::new(wasm_path)?;

    let so_file = workdir.path().join("out.so");

    native_build.shared_object_file(so_file.clone())?;

    let dlmodule = DlModule::load(so_file)?;

    Ok(dlmodule)
}

fn generate_test_wat(num_locals: usize) -> String {
    assert!(num_locals > 2);

    let mut module =
        "(module (func $localpalooza (export \"localpalooza\") (param i32) (result i32)\n"
            .to_owned();

    // Declare locals:
    module.push_str("(local ");
    for _ in 0..num_locals {
        module.push_str("i32 ");
    }
    module.push_str(")\n");

    // Use each local for the first time:
    for i in 1..num_locals {
        module.push_str(&format!(
            "(set_local {} (i32.add (get_local {}) (get_local {})))\n",
            i,
            i - 1,
            i
        ));
    }

    // Use each local for a second time, so they get pushed to the stack between uses:
    for i in 2..(num_locals - 1) {
        module.push_str(&format!(
            "(set_local {} (i32.add (get_local {}) (get_local {})))\n",
            i,
            i - 1,
            i
        ));
    }

    // Keep locals alive across a recursive call. Make as many recursive calls as the first
    // argument to the func:
    module.push_str("(if (get_local 0)\n");
    module.push_str(&format!(
        "  (then (set_local {} (i32.add (get_local {})\n",
        num_locals - 1,
        num_locals - 2
    ));
    module.push_str("      (call $localpalooza (i32.sub (get_local 0) (i32.const 1))))))\n");
    module.push_str(&format!(
        "  (else (set_local {} (i32.add (get_local {}) (get_local {})))))\n",
        num_locals - 1,
        num_locals - 2,
        num_locals - 3
    ));

    module.push_str(&format!("(get_local {})\n", num_locals - 1));
    module.push_str("))\n");
    module
}

#[macro_export]
macro_rules! stack_tests {
    ( $TestRegion:path ) => {
        use lucet_runtime::{
            DlModule, Error, InstanceHandle, Limits, Region, TrapCodeType, UntypedRetVal, Val,
        };
        use std::sync::Arc;
        use $TestRegion as TestRegion;
        use $crate::stack::stack_testcase;

        fn run(
            module: Arc<DlModule>,
            recursion_depth: i32,
        ) -> Result<UntypedRetVal, Error> {
            let region = TestRegion::create(1, &Limits::default()).expect("region can be created");
            let mut inst = region
                .new_instance(module)
                .expect("instance can be created");

            inst.run(b"localpalooza", &[recursion_depth.into()])
        }

        fn expect_ok(module: Arc<DlModule>, recursion_depth: i32) {
            assert!(run(module, recursion_depth).is_ok());
        }

        fn expect_stack_overflow(
            module: Arc<DlModule>,
            recursion_depth: i32,
            probestack: bool,
        ) {
            match run(module, recursion_depth) {
                Err(Error::RuntimeFault(details)) => {
                    // We should get a nonfatal trap due to the stack overflow.
                    assert_eq!(details.fatal, false);
                    assert_eq!(details.trapcode.ty, TrapCodeType::StackOverflow);
                    if probestack {
                        // Make sure we overflowed in the stack probe as expected
                        //
                        // TODO: this no longer differentiates between different stack overflow
                        // sites after moving the stack probe into lucetc; figure out a way to
                        // provide that information or just wait till we can do the stack probe as a
                        // global symbol again
                        let addr_details =
                            details.rip_addr_details.expect("can look up addr details");
                        assert!(addr_details.in_module_code);
                    }
                }
                res => panic!("unexpected result: {:?}", res),
            }
        }

        #[test]
        fn expect_ok_locals3_1() {
            expect_ok(stack_testcase(3).expect("generate stack_testcase 3"), 1);
        }

        // The test with 64 locals should cause a stack overflow on the 481st recursion.  The trap
        // table knows about all of the instructions in the function that manipulate the stack, so
        // the catch mechanism for this is the usual one.

        #[test]
        fn expect_ok_locals64_1() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 1);
        }

        #[test]
        fn expect_ok_locals64_2() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 2);
        }

        #[test]
        fn expect_ok_locals64_480() {
            expect_ok(stack_testcase(64).expect("generate stack_testcase 64"), 480);
        }

        #[test]
        fn expect_stack_overflow_locals64_481() {
            expect_stack_overflow(
                stack_testcase(64).expect("generate stack_testcase 64"),
                481,
                true,
            );
        }

        // 1050 locals is about 1 page (4k) on the stack - just enough for Cranelift to use probestack to grow
        // the stack. The 31st recursion should cause a stack overflow.

        #[test]
        fn expect_ok_locals_1page_1() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                1,
            );
        }

        #[test]
        fn expect_ok_locals_1page_2() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                2,
            );
        }

        #[test]
        fn expect_ok_locals_1page_30() {
            expect_ok(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                30,
            );
        }

        #[test]
        fn expect_stack_overflow_locals_1page_31() {
            expect_stack_overflow(
                stack_testcase(1050).expect("generate stack_testcase 1050"),
                31,
                true,
            );
        }

        // This test has 5000 locals - over 4 pages worth. Cranelift will use probestack here as well. The
        // 6th recursion should cause a stack overflow.

        #[test]
        fn expect_ok_locals_multipage_1() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                1,
            );
        }

        #[test]
        fn expect_ok_locals_multipage_2() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                2,
            );
        }

        #[test]
        fn expect_ok_locals_multipage_5() {
            expect_ok(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                5,
            );
        }

        #[test]
        fn expect_stack_overflow_locals_multipage_6() {
            expect_stack_overflow(
                stack_testcase(5000).expect("generate stack_testcase 5000"),
                6,
                true,
            );
        }
    };
}
<|endoftext|>


<commit_before>use chrono::UTC;
use chrono::datetime::DateTime;

use company::Company;

use postgres::Connection;

use rs_es::query::Filter;
use rs_es::units::JsonVal;

pub fn visibility_filters(conn: &Connection, company_id: &Option<i32>) -> Vec<Filter> {
  let now = DateTime::timestamp(&UTC::now());

  let visibility_rules = Filter::build_bool()
                                .with_must(
                                  vec![
                                    Filter::build_term("accepted", true)
                                           .build(),
                                    Filter::build_range("batch_starts_at")
                                           .with_lt(JsonVal::from(now))
                                           .with_format("epoch_second")
                                           .build(),
                                    Filter::build_range("batch_ends_at")
                                           .with_gte(JsonVal::from(now))
                                           .with_format("epoch_second")
                                           .build()
                                  ])
                                .build();

  let company = match *company_id {
    Some(company_id) => Company::find(conn, &company_id),
    None             => None,
  };

  match company {
    Some(company) => {
      let presented_talents = Filter::build_bool()
                                     .with_must(
                                       vec![
                                         <Filter as VectorOfTerms<i32>>::build_terms(
                                           "ids", &company.presented_talents)
                                       ].into_iter()
                                        .flat_map(|x| x)
                                        .collect::<Vec<Filter>>())
                                     .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents])
               .build()
      ]
    },
    None => vec![visibility_rules]
  }
}

pub trait VectorOfTerms<T> {
  fn build_terms(key: &str, values: &Vec<T>) -> Vec<Filter>;
}

macro_rules! build_vector_of_terms_impl {
  ($t:ty) => {
    impl VectorOfTerms<$t> for Filter {
      /// Extract all given items into multiple filters (if present)
      /// i.e. build_terms("field", vec![1, 2]) => vec![Filter(1), Filter(2)]
      /// This enable us to operate on these values with boolean values
      fn build_terms(key: &str, values: &Vec<$t>) -> Vec<Filter> {
        if values.is_empty() {
          return vec![];
        }

        vec![
          Filter::build_terms(key, values.iter()
                                         .map(|v| JsonVal::from(*v))
                                         .collect::<Vec<JsonVal>>()).build()
        ]
      }
    }
  }
}

build_vector_of_terms_impl!(i32);
build_vector_of_terms_impl!(&'static str);
<commit_msg>filters: add implementation for VectorOfTerms<String><commit_after>use chrono::UTC;
use chrono::datetime::DateTime;

use company::Company;

use postgres::Connection;

use rs_es::query::Filter;
use rs_es::units::JsonVal;

pub fn visibility_filters(conn: &Connection, company_id: &Option<i32>) -> Vec<Filter> {
  let now = DateTime::timestamp(&UTC::now());

  let visibility_rules = Filter::build_bool()
                                .with_must(
                                  vec![
                                    Filter::build_term("accepted", true)
                                           .build(),
                                    Filter::build_range("batch_starts_at")
                                           .with_lt(JsonVal::from(now))
                                           .with_format("epoch_second")
                                           .build(),
                                    Filter::build_range("batch_ends_at")
                                           .with_gte(JsonVal::from(now))
                                           .with_format("epoch_second")
                                           .build()
                                  ])
                                .build();

  let company = match *company_id {
    Some(company_id) => Company::find(conn, &company_id),
    None             => None,
  };

  match company {
    Some(company) => {
      let presented_talents = Filter::build_bool()
                                     .with_must(
                                       vec![
                                         <Filter as VectorOfTerms<i32>>::build_terms(
                                           "ids", &company.presented_talents)
                                       ].into_iter()
                                        .flat_map(|x| x)
                                        .collect::<Vec<Filter>>())
                                     .build();
      vec![
        Filter::build_bool()
               .with_should(vec![visibility_rules, presented_talents])
               .build()
      ]
    },
    None => vec![visibility_rules]
  }
}

pub trait VectorOfTerms<T> {
  fn build_terms(key: &str, values: &Vec<T>) -> Vec<Filter>;
}

macro_rules! build_vector_of_terms_impl {
  ($t:ty) => {
    impl VectorOfTerms<$t> for Filter {
      /// Extract all given items into multiple filters (if present)
      /// i.e. build_terms("field", vec![1, 2]) => vec![Filter(1), Filter(2)]
      /// This enable us to operate on these values with boolean values
      fn build_terms(key: &str, values: &Vec<$t>) -> Vec<Filter> {
        if values.is_empty() {
          return vec![];
        }

        vec![
          Filter::build_terms(key, values.iter()
                                         .map(|v| JsonVal::from(v.clone()))
                                         .collect::<Vec<JsonVal>>()).build()
        ]
      }
    }
  }
}

build_vector_of_terms_impl!(i32);
build_vector_of_terms_impl!(String);
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};

use gameboy;
use gameboy::{Interconnect, Interrupt, Irq, Memory};
use gameboy::frame::{Color, Frame};

const VRAM_SIZE: usize = 0x4000;
const SPRITE_DATA_SIZE: usize = 0xA0;

pub struct GpuStat {
    pub coincidence_interrupt_enabled: bool,
    pub OAM_interrupt_enabled: bool,
    pub VBlank_interrupt_enabled: bool,
    pub HBlank_interrupt_enabled: bool,
}

impl GpuStat {
    pub fn new() -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: false,
            OAM_interrupt_enabled: false,
            VBlank_interrupt_enabled: false,
            HBlank_interrupt_enabled: false,
        }
    }

    pub fn from_u8(b: u8) -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: b & 0x40 == 0x40,
            OAM_interrupt_enabled: b & 0x20 == 0x20,
            VBlank_interrupt_enabled: b & 0x10 == 0x10,
            HBlank_interrupt_enabled: b & 0x08 == 0x08,
        }
    }

    pub fn to_u8(&self, gpu: &Gpu) -> u8 {
        (if self.coincidence_interrupt_enabled {
            0x40
        } else {
            0
        }) | (if self.OAM_interrupt_enabled { 0x20 } else { 0 }) |
        (if self.VBlank_interrupt_enabled {
            0x10
        } else {
            0
        }) |
        (if self.HBlank_interrupt_enabled {
            0x08
        } else {
            0
        }) | (if gpu.ly == gpu.lyc { 0x04 } else { 0 }) | gpu.mode.to_u8()
    }
}

#[derive(Eq, PartialEq, Clone)]
pub enum GpuMode {
    HBlank,
    VBlank,
    SearchingRam,
    TransferringData,
}

impl GpuMode {
    fn to_u8(&self) -> u8 {
        match *self {
            GpuMode::HBlank => 0x00,
            GpuMode::VBlank => 0x01,
            GpuMode::SearchingRam => 0x02,
            GpuMode::TransferringData => 0x03,
        }
    }
}

pub enum SpriteShape {
    Square,
    Rectangle,
}

pub struct Gpu {
    pub enabled: bool,
    pub ram: Memory,
    pub sprite_data: Memory,
    control_register: u8,
    stat: GpuStat,
    scroll_y: u8,
    scroll_x: u8,
    window_y: u8,
    window_x: u8,
    pub ly: u8,
    lyc: u8,
    bg_palette: u8,
    palette0: u8,
    palette1: u8,

    cycles: isize,

    pub frame: Frame,
    pub backbuffer: Frame,
    mode: GpuMode,

    counter: u8,
    tile_base: usize,
    background_base: usize,
    sprite_shape: SpriteShape,
}

impl Gpu {
    pub fn new() -> Gpu {
        Gpu {
            enabled: true,
            ram: Memory::new(VRAM_SIZE),
            sprite_data: Memory::new(SPRITE_DATA_SIZE),
            control_register: 0x00,
            stat: GpuStat::new(),
            scroll_y: 0x00,
            scroll_x: 0x00,
            window_y: 0x00,
            window_x: 0x00,
            ly: 0x00,
            lyc: 0x00,
            bg_palette: 0x00,
            palette0: 0x00,
            palette1: 0x00,
            cycles: 0x00,
            frame: Frame::new(),
            backbuffer: Frame::new(),
            mode: GpuMode::HBlank,
            counter: 0,
            tile_base: 0x00,
            background_base: 0xC00,
            sprite_shape: SpriteShape::Square,
        }
    }

    pub fn reset(&mut self) {
        *self = Gpu::new();
    }

    fn set_status(&mut self, irq: &mut Irq) {
        if !self.enabled {
            self.cycles = 0x1C8;
            self.ly = 0x00;
            self.stat = GpuStat::from_u8(0xFD);
        }

        let current_mode = self.mode.clone();

        if self.ly >= 0x90 {
            self.mode = GpuMode::VBlank;
        } else {
            if self.cycles >= 0x178 {
                self.mode = GpuMode::SearchingRam;
            } else if self.cycles >= 0xCC {
                self.mode = GpuMode::TransferringData;
            } else {
                self.mode = GpuMode::HBlank;
            }
        }

        if self.mode != current_mode {
            irq.request(Interrupt::Lcd);
        }

        self.check_coincidence(irq);
    }

    pub fn step(&mut self, irq: &mut Irq, cycles: usize) -> Result<(), String> {
        self.set_status(irq);

        let cycles = cycles as isize;
        if self.enabled {
            self.cycles -= cycles;
        } else {
            return Ok(());
        }

        if self.cycles <= 0 {
            self.ly += 0x01;

            self.cycles = 0x1C8;

            if self.ly == 0x90 {
                irq.request(Interrupt::Vblank);
            } else if self.ly > 0x99 {
                self.ly = 0x00;
            } else if self.ly < 0x90 {
                self.render_scanline();
            }
        }

        Ok(())
    }

    fn render_scanline(&mut self) {
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        // self.clear_scanline(line);
        self.render_background(line);
        self.render_sprites(line);
        self.frame = self.backbuffer.clone();
    }

    fn clear_scanline(&mut self, line: usize) {
        for i in 0..160 {
            self.backbuffer.pixels[line * 160 + i] = Color::new(0xFF, 0xFF, 0xFF, 0xFF);
        }
    }

    fn render_background(&mut self, line: usize) {
        let bg_base = self.background_base;
        let tile_base = self.tile_base;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..gameboy::SCREEN_WIDTH {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number = self.ram[bg_base + (bg_map_row * 0x20 + bg_map_col)] as usize;

            let line_offset = (line % 0x08) << 0x01;

            let tile_data_start = tile_base +
                                  (if tile_base == 0x00 {
                raw_tile_number
            } else {
                (raw_tile_number as i8 as i16 + 0x80) as usize
            }) * 0x10 + line_offset;

            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.backbuffer.pixels[self.ly as usize * 160 + i as usize] = c;
        }
    }

    fn render_sprites(&mut self, line: usize) {
        for i in 0..40 {
            let sprite_table_entry_base = i * 0x04;
            let s_y = self.sprite_data[sprite_table_entry_base] as i16 - 0x10;
            let s_x = self.sprite_data[sprite_table_entry_base + 0x01] as i16 - 0x08;
            let sprite_height = if let SpriteShape::Rectangle = self.sprite_shape {
                0x10
            } else {
                0x08
            };
            if line < s_y as usize || line >= s_y as usize + sprite_height {
                continue;
            }
            let tile_number = self.sprite_data[sprite_table_entry_base + 0x02] as i16;
            let attributes = self.sprite_data[sprite_table_entry_base + 0x03];
            let above_background = attributes & 0x80 == 0x00;
            let flip_y = attributes & 0x40 == 0x40;
            let flip_x = attributes & 0x20 == 0x20;

            let sprite_y = (line as i16 - s_y) << 0x01;
            let tile_data_start = (tile_number * 0x10) + sprite_y;
            for x in 0..8 {
                if s_x + x < 0 || s_x + x >= 160 {
                    continue;
                }
                let shift = if flip_x { x } else { 0x07 - x };
                let tile_data1 = (self.ram[tile_data_start as usize] >> shift) & 0x01;
                let tile_data2 = (self.ram[tile_data_start as usize + 0x01] >> shift) & 0x01;
                let total_row_data = (tile_data2 << 1) | tile_data1;
                let color_value = total_row_data;
                if color_value == 0x00 {
                    continue;
                }
                let c = self.get_sprite_color_for_byte(color_value as u8,
                                                       ((attributes & 0x10) >> 0x04) as u8);
                self.backbuffer.pixels[line as usize * 160 + (s_x as usize + x as usize)] = c;
            }
        }
    }

    fn get_background_color_for_byte(&self, b: u8) -> Color {
        let palette_index = match b {
            0x00 => self.bg_palette & 0x0003,
            0x01 => (self.bg_palette & 0x000C) >> 0x02,
            0x02 => (self.bg_palette & 0x0030) >> 0x04,
            0x03 => (self.bg_palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    fn get_sprite_color_for_byte(&self, b: u8, palette_entry: u8) -> Color {
        let palette = if palette_entry == 0x00 {
            self.palette0
        } else {
            self.palette1
        };

        let palette_index = match b {
            0x00 => palette & 0x0003,
            0x01 => (palette & 0x000C) >> 0x02,
            0x02 => (palette & 0x0030) >> 0x04,
            0x03 => (palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        match addr {
            0x40 => self.control_register,
            0x41 => self.stat.to_u8(self),
            0x42 => self.scroll_y,
            0x43 => self.scroll_x,
            0x4A => self.window_y,
            0x4B => self.window_x,
            0x44 => self.ly,
            0x45 => self.lyc,
            0x47 => self.bg_palette,
            0x48 => self.palette0,
            0x49 => self.palette1,
            _ => panic!("tried to read GPU memory that is not mapped"),
        }
    }

    pub fn write_u8(&mut self, addr: u16, val: u8) {
        match addr {
            0x40 => {
                self.control_register = val;
                self.enabled = self.control_register & 0x80 == 0x80;
                self.tile_base = if self.control_register & 0x10 == 0x10 {
                    0x00
                } else {
                    0x800
                };
                self.background_base = if self.control_register & 0x08 == 0x08 {
                    0x1C00
                } else {
                    0x1800
                };
                self.sprite_shape = if self.control_register & 0x04 == 0x04 {
                    SpriteShape::Rectangle
                } else {
                    SpriteShape::Square
                };
            }
            0x41 => self.stat = GpuStat::from_u8(val),
            0x42 => self.scroll_y = val,
            0x43 => self.scroll_x = val,
            0x44 => (),
            0x45 => {
                self.lyc = val;
            }
            0x47 => self.bg_palette = val,
            0x48 => self.palette0 = val,
            0x49 => self.palette1 = val,
            0x4A => self.window_y = val,
            0x4B => self.window_x = val,
            _ => panic!("tried to write GPU memory that is not mapped: {:04}", addr),
        }
    }

    pub fn write_u16(&mut self, addr: u16, val: u16) {
        let addr = addr as usize;
        LittleEndian::write_u16(&mut self.ram[addr..], val)
    }

    fn check_coincidence(&mut self, irq: &mut Irq) {
        // If LY == LYC then set the coincidence flag
        if self.lyc == self.ly && self.stat.coincidence_interrupt_enabled {
            irq.request(Interrupt::Lcd);
        }
    }
}
<commit_msg>Invert the sprite scanline rendering logic and leave a comment RE Kirby disappearing<commit_after>
use byteorder::{ByteOrder, LittleEndian};

use std;

use gameboy;
use gameboy::{Interconnect, Interrupt, Irq, Memory};
use gameboy::frame::{Color, Frame};

const VRAM_SIZE: usize = 0x4000;
const SPRITE_DATA_SIZE: usize = 0xA0;

pub struct GpuStat {
    pub coincidence_interrupt_enabled: bool,
    pub OAM_interrupt_enabled: bool,
    pub VBlank_interrupt_enabled: bool,
    pub HBlank_interrupt_enabled: bool,
}

impl GpuStat {
    pub fn new() -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: false,
            OAM_interrupt_enabled: false,
            VBlank_interrupt_enabled: false,
            HBlank_interrupt_enabled: false,
        }
    }

    pub fn from_u8(b: u8) -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: b & 0x40 == 0x40,
            OAM_interrupt_enabled: b & 0x20 == 0x20,
            VBlank_interrupt_enabled: b & 0x10 == 0x10,
            HBlank_interrupt_enabled: b & 0x08 == 0x08,
        }
    }

    pub fn to_u8(&self, gpu: &Gpu) -> u8 {
        (if self.coincidence_interrupt_enabled {
            0x40
        } else {
            0
        }) | (if self.OAM_interrupt_enabled { 0x20 } else { 0 }) |
        (if self.VBlank_interrupt_enabled {
            0x10
        } else {
            0
        }) |
        (if self.HBlank_interrupt_enabled {
            0x08
        } else {
            0
        }) | (if gpu.ly == gpu.lyc { 0x04 } else { 0 }) | gpu.mode.to_u8()
    }
}

#[derive(Eq, PartialEq, Clone)]
pub enum GpuMode {
    HBlank,
    VBlank,
    SearchingRam,
    TransferringData,
}

impl GpuMode {
    fn to_u8(&self) -> u8 {
        match *self {
            GpuMode::HBlank => 0x00,
            GpuMode::VBlank => 0x01,
            GpuMode::SearchingRam => 0x02,
            GpuMode::TransferringData => 0x03,
        }
    }
}

pub enum SpriteShape {
    Square,
    Rectangle,
}

pub struct Gpu {
    pub enabled: bool,
    pub ram: Memory,
    pub sprite_data: Memory,
    control_register: u8,
    stat: GpuStat,
    scroll_y: u8,
    scroll_x: u8,
    window_y: u8,
    window_x: u8,
    pub ly: u8,
    lyc: u8,
    bg_palette: u8,
    palette0: u8,
    palette1: u8,

    cycles: isize,

    pub frame: Frame,
    pub backbuffer: Frame,
    mode: GpuMode,

    counter: u8,
    tile_base: usize,
    background_base: usize,
    sprite_shape: SpriteShape,
}

impl Gpu {
    pub fn new() -> Gpu {
        Gpu {
            enabled: true,
            ram: Memory::new(VRAM_SIZE),
            sprite_data: Memory::new(SPRITE_DATA_SIZE),
            control_register: 0x00,
            stat: GpuStat::new(),
            scroll_y: 0x00,
            scroll_x: 0x00,
            window_y: 0x00,
            window_x: 0x00,
            ly: 0x00,
            lyc: 0x00,
            bg_palette: 0x00,
            palette0: 0x00,
            palette1: 0x00,
            cycles: 0x00,
            frame: Frame::new(),
            backbuffer: Frame::new(),
            mode: GpuMode::HBlank,
            counter: 0,
            tile_base: 0x00,
            background_base: 0xC00,
            sprite_shape: SpriteShape::Square,
        }
    }

    pub fn reset(&mut self) {
        *self = Gpu::new();
    }

    fn set_status(&mut self, irq: &mut Irq) {
        if !self.enabled {
            self.cycles = 0x1C8;
            self.ly = 0x00;
            self.stat = GpuStat::from_u8(0xFD);
        }

        let current_mode = self.mode.clone();

        if self.ly >= 0x90 {
            self.mode = GpuMode::VBlank;
        } else {
            if self.cycles >= 0x178 {
                self.mode = GpuMode::SearchingRam;
            } else if self.cycles >= 0xCC {
                self.mode = GpuMode::TransferringData;
            } else {
                self.mode = GpuMode::HBlank;
            }
        }

        if self.mode != current_mode {
            irq.request(Interrupt::Lcd);
        }

        self.check_coincidence(irq);
    }

    pub fn step(&mut self, irq: &mut Irq, cycles: usize) -> Result<(), String> {
        self.set_status(irq);

        let cycles = cycles as isize;
        if self.enabled {
            self.cycles -= cycles;
        } else {
            return Ok(());
        }

        if self.cycles <= 0 {
            self.ly += 0x01;

            self.cycles = 0x1C8;

            if self.ly == 0x90 {
                irq.request(Interrupt::Vblank);
            } else if self.ly > 0x99 {
                self.ly = 0x00;
            } else if self.ly < 0x90 {
                self.render_scanline();
            }
        }

        Ok(())
    }

    fn render_scanline(&mut self) {
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        // self.clear_scanline(line);
        self.render_background(line);
        self.render_sprites(line);
        self.frame = self.backbuffer.clone();
    }

    fn clear_scanline(&mut self, line: usize) {
        for i in 0..160 {
            self.backbuffer.pixels[line * 160 + i] = Color::new(0xFF, 0xFF, 0xFF, 0xFF);
        }
    }

    fn render_background(&mut self, line: usize) {
        let bg_base = self.background_base;
        let tile_base = self.tile_base;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..gameboy::SCREEN_WIDTH {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number = self.ram[bg_base + (bg_map_row * 0x20 + bg_map_col)] as usize;

            let line_offset = (line % 0x08) << 0x01;

            let tile_data_start = tile_base +
                                  (if tile_base == 0x00 {
                raw_tile_number
            } else {
                (raw_tile_number as i8 as i16 + 0x80) as usize
            }) * 0x10 + line_offset;

            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.backbuffer.pixels[self.ly as usize * 160 + i as usize] = c;
        }
    }

    fn render_sprites(&mut self, line: usize) {
        for i in 0..40 {
            let sprite_table_entry_base = i * 0x04;
            let s_y = self.sprite_data[sprite_table_entry_base] as i16 - 0x10;
            let s_x = self.sprite_data[sprite_table_entry_base + 0x01] as i16 - 0x08;
            let sprite_height = if let SpriteShape::Rectangle = self.sprite_shape {
                0x10
            } else {
                0x08
            };

            // Ignore this scanline if the sprite doesn't sit within it
            // Should we be casting to usize here? If the number is negative the sign bit
            // will be throwing this calculation off entirely - perhaps this is why Kirby
            // disappears when sprite Y position is negative?
            if s_y as usize <= line && (s_y as usize + sprite_height) > line {
                let tile_number = self.sprite_data[sprite_table_entry_base + 0x02] as i16;
                let attributes = self.sprite_data[sprite_table_entry_base + 0x03];
                let above_background = attributes & 0x80 == 0x00;
                let flip_y = attributes & 0x40 == 0x40;
                let flip_x = attributes & 0x20 == 0x20;

                let sprite_y = (line as i16 - s_y) << 0x01;
                let tile_data_start = (tile_number * 0x10) + sprite_y;
                for x in 0..8 {
                    if s_x + x < 0 || s_x + x >= 160 {
                        continue;
                    }
                    let shift = if flip_x { x } else { 0x07 - x };
                    let tile_data1 = (self.ram[tile_data_start as usize] >> shift) & 0x01;
                    let tile_data2 = (self.ram[tile_data_start as usize + 0x01] >> shift) & 0x01;
                    let total_row_data = (tile_data2 << 1) | tile_data1;
                    let color_value = total_row_data;
                    if color_value == 0x00 {
                        continue;
                    }
                    let c = self.get_sprite_color_for_byte(color_value as u8,
                                                           ((attributes & 0x10) >> 0x04) as u8);
                    self.backbuffer.pixels[line as usize * 160 + (s_x as usize + x as usize)] = c;
                }
            }
        }
    }

    fn get_background_color_for_byte(&self, b: u8) -> Color {
        let palette_index = match b {
            0x00 => self.bg_palette & 0x0003,
            0x01 => (self.bg_palette & 0x000C) >> 0x02,
            0x02 => (self.bg_palette & 0x0030) >> 0x04,
            0x03 => (self.bg_palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    fn get_sprite_color_for_byte(&self, b: u8, palette_entry: u8) -> Color {
        let palette = if palette_entry == 0x00 {
            self.palette0
        } else {
            self.palette1
        };

        let palette_index = match b {
            0x00 => palette & 0x0003,
            0x01 => (palette & 0x000C) >> 0x02,
            0x02 => (palette & 0x0030) >> 0x04,
            0x03 => (palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        match addr {
            0x40 => self.control_register,
            0x41 => self.stat.to_u8(self),
            0x42 => self.scroll_y,
            0x43 => self.scroll_x,
            0x4A => self.window_y,
            0x4B => self.window_x,
            0x44 => self.ly,
            0x45 => self.lyc,
            0x47 => self.bg_palette,
            0x48 => self.palette0,
            0x49 => self.palette1,
            _ => panic!("tried to read GPU memory that is not mapped"),
        }
    }

    pub fn write_u8(&mut self, addr: u16, val: u8) {
        match addr {
            0x40 => {
                self.control_register = val;
                self.enabled = self.control_register & 0x80 == 0x80;
                self.tile_base = if self.control_register & 0x10 == 0x10 {
                    0x00
                } else {
                    0x800
                };
                self.background_base = if self.control_register & 0x08 == 0x08 {
                    0x1C00
                } else {
                    0x1800
                };
                self.sprite_shape = if self.control_register & 0x04 == 0x04 {
                    SpriteShape::Rectangle
                } else {
                    SpriteShape::Square
                };
            }
            0x41 => self.stat = GpuStat::from_u8(val),
            0x42 => self.scroll_y = val,
            0x43 => self.scroll_x = val,
            0x44 => (),
            0x45 => {
                self.lyc = val;
            }
            0x47 => self.bg_palette = val,
            0x48 => self.palette0 = val,
            0x49 => self.palette1 = val,
            0x4A => self.window_y = val,
            0x4B => self.window_x = val,
            _ => panic!("tried to write GPU memory that is not mapped: {:04}", addr),
        }
    }

    pub fn write_u16(&mut self, addr: u16, val: u16) {
        let addr = addr as usize;
        LittleEndian::write_u16(&mut self.ram[addr..], val)
    }

    fn check_coincidence(&mut self, irq: &mut Irq) {
        // If LY == LYC then set the coincidence flag
        if self.lyc == self.ly && self.stat.coincidence_interrupt_enabled {
            irq.request(Interrupt::Lcd);
        }
    }
}
<|endoftext|>


<commit_before>extern crate rs6502;
extern crate sdl2;

mod console;
mod position;
mod text;

use std::io::{self, BufRead, Write};

use rs6502::{CodeSegment, Cpu, Disassembler};
use sdl2::render::Renderer;

pub use self::console::Console;
pub use self::position::Position;
pub use self::text::Text;

const HELPTEXT: &'static str = "

HAKKA
-----

Commands:

break | b: addr
       - Toggles a breakpoint at a specific address. If the
         program counter hits this address, execution
         stops.

step | s
       - Executes the current instruction before breaking
         execution again.

continue | c
       - Resumes execution of code within the virtual machine.

list
       - Lists the code surrounding the current program
         counter.

memset | set: addr args [args, ...]
       - memset sets the value of memory directly.

memdmp | dmp: page || start end
       - memdmp dumps a single memory page, or a specified
         memory range from <start> to <end> (inclusive).

monitor | mon: start end
       - monitor dumps the memory between start and end (inclusive)
         every second. Press ENTER to stop the monitor.

source
       - Lists the code currently running in the virtual
         machine. A '>' symbol indicates the current
         program counter.

registers | reg
       - Lists the CPU registers and their current values.

help
       - Lists this help text.
";

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine<'a> {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    pub console: Console<'a>,
    segments: Vec<CodeSegment>,
    clock_rate: Option<u32>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
    step: bool,
}

impl<'a> VirtualMachine<'a> {
    pub fn new<CR>(cpu: Cpu, clock_rate: CR, mut console: Console<'a>) -> VirtualMachine<'a>
        where CR: Into<Option<u32>>
    {
        console.println("Welcome to hakka. Type 'help' for instructions.");
        console.println("");

        VirtualMachine {
            cpu: cpu,
            console: console,
            segments: Vec::new(),
            clock_rate: clock_rate.into(),
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
            step: false,
        }
    }

    pub fn render(&mut self, mut renderer: &mut Renderer) {
        self.console.render(renderer);
    }

    pub fn load_code_segments(&mut self, segments: Vec<CodeSegment>) {
        if segments.len() == 0 {
            return;
        }
        self.segments = segments;
        for segment in &self.segments {
            self.cpu.load(&segment.code, segment.address);
        }

        self.cpu.registers.PC = self.segments[0].address;
    }

    /// Cycles the Virtual Machine CPU according to the clock rate
    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while (n < clock_rate && !self.broken) || self.step {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    self.console.println("");
                    self.console
                        .println(format!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC));
                    self.console.toggle();
                }
                // If we stepped, dump the local disassembly
                if self.step {
                    self.dump_local_disassembly();
                }
                self.step = false;
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            if self.step {
                self.dump_local_disassembly();
            }
            self.step = false;
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                self.console.println("");
                self.console.println(format!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC));
                self.console.toggle();
            }
        }
    }

    pub fn execute_command<S>(&mut self, cmd: S)
        where S: Into<String>
    {
        let mut input: String = cmd.into().trim().into();

        if input == "r" || input == "repeat" {
            input = self.last_command.clone();
        }

        let parts = input.split(" ").collect::<Vec<_>>();

        if input.len() == 0 {
            if self.monitor.enabled {
                self.monitor.enabled = false;
            }
        } else if parts[0] == "clear" || parts[0] == "cls" {
            self.console.clear();
        } else if parts[0] == "help" {
            self.console.print_lines(format!("{}", HELPTEXT));
        } else if parts[0] == "source" {
            self.dump_disassembly();
        } else if parts[0] == "list" {
            self.dump_local_disassembly();
        } else if parts[0] == "monitor" || parts[0] == "mon" {
            if self.monitor.enabled {
                self.monitor.enabled = false;
            } else {
                self.enable_memory_monitor(&input);
            }
        } else if parts[0] == "memset" || parts[0] == "set" {
            if parts.len() < 3 {
                self.console
                    .println("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
            } else if parts.len() == 3 {
                if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                        self.cpu.memory[dst as usize] = src;
                    } else {
                        self.console.println("ERR: Unable to parse source byte value");
                    }
                } else {
                    self.console.println("ERR: Unable to parse destination byte value");
                }
            } else {
                if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    for p in &parts[2..] {
                        if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                            self.cpu.memory[dst] = byte;
                            dst += 0x01;
                        } else {
                            self.console.println("ERR: Unable to parse source byte value");
                        }
                    }
                } else {
                    self.console.println("ERR: Unable to parse destination byte value");
                }
            }
        } else if parts[0] == "memdmp" || parts[0] == "dmp" {
            // 1 argument assumes 1 memory "page"
            if parts.len() == 2 {
                if let Ok(page_number) = parts[1].parse() {

                    self.dump_memory_page(page_number);
                } else {
                    self.console.println("ERR: Unable to parse memory page");
                }
            } else if parts.len() == 3 {
                // A memory range instead
                if let Ok(start) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    if let Ok(end) = u16::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                        self.dump_memory_range(start, end);
                    } else {
                        self.console.println("ERR: Unable to parse end address value");
                    }
                } else {
                    self.console.println("ERR: Unable to parse start address value");
                }
            }
        } else if parts[0] == "registers" || parts[0] == "reg" {
            self.dump_registers();
        } else if parts[0] == "break" || parts[0] == "b" {
            // 1 argument assumes 1 memory "page"
            if parts.len() == 2 {
                if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                    if addr <= u16::max_value() as usize {
                        if self.breakpoints[addr] > 0 {
                            self.console.println(format!("Removed breakpoint at {:04X}", addr));
                            self.breakpoints[addr] = 0;
                        } else {
                            self.console.println(format!("Added breakpoint at {:04X}", addr));
                            self.breakpoints[addr] = 1;
                        }
                    } else {
                        self.console.println("ERR: Value outside addressable range.");
                    }
                } else {
                    self.broken = true;
                    println!("Execution stopped");
                }
            } else {
                self.console.println("ERR: Requires 1 argument");
            }
        } else if parts[0] == "continue" || parts[0] == "c" {
            self.broken = false;
            self.console.println("Execution resumed");
        } else if parts[0] == "step" || parts[0] == "s" {
            self.broken = true;
            self.step = true;
        } else {
            self.console.println("Unknown command");
        }

        // Don't assign a blank command as a last command
        if input.len() > 0 {
            self.last_command = input.clone();
        }
    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            self.console.println("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;
        }
    }

    fn dump_disassembly(&mut self) {
        self.console.println("");

        for segment in &self.segments {
            self.console.println(format!(".ORG ${:04X}", segment.address));
            let disassembler = Disassembler::with_offset(segment.address);
            let pairs = disassembler.disassemble_with_addresses(&segment.code);
            let lines = self.highlight_lines(self.cpu.registers.PC as usize,
                                             pairs,
                                             segment.address,
                                             false);
            self.console.print_lines(lines.join(""));
        }
        self.console.println("");
    }

    fn dump_local_disassembly(&mut self) {
        self.console.println("");

        let result = {
            let pc = self.cpu.registers.PC as usize;
            let local_segment = self.get_local_segment(pc);
            let disassembler = Disassembler::with_offset(local_segment.address);
            let pairs = disassembler.disassemble_with_addresses(&local_segment.code);
            self.highlight_lines(pc, pairs, local_segment.address, true)
        };
        self.console.print_lines(result.join(""));
        self.console.println("");
    }

    pub fn dump_memory_page(&mut self, page: usize) {
        let mut addr = page * 0x100;
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            self.console.print(format!("{:04X}: ", addr));
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.wrap_line();
            addr += 0x08;
        }
    }

    pub fn dump_memory(&mut self) {
        self.console.println("");
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.println("");
        }
    }

    fn dump_memory_range(&mut self, start: u16, end: u16) {
        self.console.println("");
        let start = start as usize;
        let end = end as usize;
        for chunk in self.cpu.memory[start..end + 0x01].chunks(8) {
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.println("");
        }
    }

    fn dump_registers(&mut self) {
        self.console.println("");
        self.console.println(format!("A: {} ({:04X})", self.cpu.registers.A, self.cpu.registers.A));
        self.console.println(format!("X: {} ({:04X})", self.cpu.registers.X, self.cpu.registers.X));
        self.console.println(format!("Y: {} ({:04X})", self.cpu.registers.Y, self.cpu.registers.Y));
        self.console.println(format!("PC: {} ({:04X})",
                                     self.cpu.registers.PC,
                                     self.cpu.registers.PC));
        self.console.println(format!("S: {} ({:04X})",
                                     self.cpu.stack.pointer,
                                     self.cpu.stack.pointer));
    }

    fn get_local_segment(&self, pc: usize) -> &CodeSegment {
        for segment in &self.segments {
            let addr = segment.address as usize;
            if pc >= addr && pc <= addr + segment.code.len() {
                return segment;
            }
        }

        &self.segments[0]
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       segment_start: u16,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc as isize - segment_start as isize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as isize) < start || (pair.1 as isize) > base + 0x0A {
                    continue;
                }
            }
            let current_line = pc as u16 == segment_start + pair.1;
            let breakpoint = self.breakpoints[segment_start as usize + pair.1 as usize] > 0x00;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else if !breakpoint && current_line {
                result.push(format!(">   {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}
<commit_msg>Pack memory monitor output closer together<commit_after>extern crate rs6502;
extern crate sdl2;

mod console;
mod position;
mod text;

use std::io::{self, BufRead, Write};

use rs6502::{CodeSegment, Cpu, Disassembler};
use sdl2::render::Renderer;

pub use self::console::Console;
pub use self::position::Position;
pub use self::text::Text;

const HELPTEXT: &'static str = "

HAKKA
-----

Commands:

break | b: addr
       - Toggles a breakpoint at a specific address. If the
         program counter hits this address, execution
         stops.

step | s
       - Executes the current instruction before breaking
         execution again.

continue | c
       - Resumes execution of code within the virtual machine.

list
       - Lists the code surrounding the current program
         counter.

memset | set: addr args [args, ...]
       - memset sets the value of memory directly.

memdmp | dmp: page || start end
       - memdmp dumps a single memory page, or a specified
         memory range from <start> to <end> (inclusive).

monitor | mon: start end
       - monitor dumps the memory between start and end (inclusive)
         every second. Press ENTER to stop the monitor.

source
       - Lists the code currently running in the virtual
         machine. A '>' symbol indicates the current
         program counter.

registers | reg
       - Lists the CPU registers and their current values.

help
       - Lists this help text.
";

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine<'a> {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    pub console: Console<'a>,
    segments: Vec<CodeSegment>,
    clock_rate: Option<u32>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
    step: bool,
}

impl<'a> VirtualMachine<'a> {
    pub fn new<CR>(cpu: Cpu, clock_rate: CR, mut console: Console<'a>) -> VirtualMachine<'a>
        where CR: Into<Option<u32>>
    {
        console.println("Welcome to hakka. Type 'help' for instructions.");
        console.println("");

        VirtualMachine {
            cpu: cpu,
            console: console,
            segments: Vec::new(),
            clock_rate: clock_rate.into(),
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
            step: false,
        }
    }

    pub fn render(&mut self, mut renderer: &mut Renderer) {
        self.console.render(renderer);
    }

    pub fn load_code_segments(&mut self, segments: Vec<CodeSegment>) {
        if segments.len() == 0 {
            return;
        }
        self.segments = segments;
        for segment in &self.segments {
            self.cpu.load(&segment.code, segment.address);
        }

        self.cpu.registers.PC = self.segments[0].address;
    }

    /// Cycles the Virtual Machine CPU according to the clock rate
    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while (n < clock_rate && !self.broken) || self.step {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    self.console.println("");
                    self.console
                        .println(format!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC));
                    self.console.toggle();
                }
                // If we stepped, dump the local disassembly
                if self.step {
                    self.dump_local_disassembly();
                }
                self.step = false;
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            if self.step {
                self.dump_local_disassembly();
            }
            self.step = false;
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                self.console.println("");
                self.console.println(format!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC));
                self.console.toggle();
            }
        }
    }

    pub fn execute_command<S>(&mut self, cmd: S)
        where S: Into<String>
    {
        let mut input: String = cmd.into().trim().into();

        if input == "r" || input == "repeat" {
            input = self.last_command.clone();
        }

        let parts = input.split(" ").collect::<Vec<_>>();

        if input.len() == 0 {
            if self.monitor.enabled {
                self.monitor.enabled = false;
            }
        } else if parts[0] == "clear" || parts[0] == "cls" {
            self.console.clear();
        } else if parts[0] == "help" {
            self.console.print_lines(format!("{}", HELPTEXT));
        } else if parts[0] == "source" {
            self.dump_disassembly();
        } else if parts[0] == "list" {
            self.dump_local_disassembly();
        } else if parts[0] == "monitor" || parts[0] == "mon" {
            if self.monitor.enabled {
                self.monitor.enabled = false;
            } else {
                self.enable_memory_monitor(&input);
            }
        } else if parts[0] == "memset" || parts[0] == "set" {
            if parts.len() < 3 {
                self.console
                    .println("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
            } else if parts.len() == 3 {
                if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                        self.cpu.memory[dst as usize] = src;
                    } else {
                        self.console.println("ERR: Unable to parse source byte value");
                    }
                } else {
                    self.console.println("ERR: Unable to parse destination byte value");
                }
            } else {
                if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    for p in &parts[2..] {
                        if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                            self.cpu.memory[dst] = byte;
                            dst += 0x01;
                        } else {
                            self.console.println("ERR: Unable to parse source byte value");
                        }
                    }
                } else {
                    self.console.println("ERR: Unable to parse destination byte value");
                }
            }
        } else if parts[0] == "memdmp" || parts[0] == "dmp" {
            // 1 argument assumes 1 memory "page"
            if parts.len() == 2 {
                if let Ok(page_number) = parts[1].parse() {

                    self.dump_memory_page(page_number);
                } else {
                    self.console.println("ERR: Unable to parse memory page");
                }
            } else if parts.len() == 3 {
                // A memory range instead
                if let Ok(start) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                    if let Ok(end) = u16::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                        self.dump_memory_range(start, end);
                    } else {
                        self.console.println("ERR: Unable to parse end address value");
                    }
                } else {
                    self.console.println("ERR: Unable to parse start address value");
                }
            }
        } else if parts[0] == "registers" || parts[0] == "reg" {
            self.dump_registers();
        } else if parts[0] == "break" || parts[0] == "b" {
            // 1 argument assumes 1 memory "page"
            if parts.len() == 2 {
                if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                    if addr <= u16::max_value() as usize {
                        if self.breakpoints[addr] > 0 {
                            self.console.println(format!("Removed breakpoint at {:04X}", addr));
                            self.breakpoints[addr] = 0;
                        } else {
                            self.console.println(format!("Added breakpoint at {:04X}", addr));
                            self.breakpoints[addr] = 1;
                        }
                    } else {
                        self.console.println("ERR: Value outside addressable range.");
                    }
                } else {
                    self.broken = true;
                    println!("Execution stopped");
                }
            } else {
                self.console.println("ERR: Requires 1 argument");
            }
        } else if parts[0] == "continue" || parts[0] == "c" {
            self.broken = false;
            self.console.println("Execution resumed");
        } else if parts[0] == "step" || parts[0] == "s" {
            self.broken = true;
            self.step = true;
        } else {
            self.console.println("Unknown command");
        }

        // Don't assign a blank command as a last command
        if input.len() > 0 {
            self.last_command = input.clone();
        }
    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            self.console.println("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;
        }
    }

    fn dump_disassembly(&mut self) {
        self.console.println("");

        for segment in &self.segments {
            self.console.println(format!(".ORG ${:04X}", segment.address));
            let disassembler = Disassembler::with_offset(segment.address);
            let pairs = disassembler.disassemble_with_addresses(&segment.code);
            let lines = self.highlight_lines(self.cpu.registers.PC as usize,
                                             pairs,
                                             segment.address,
                                             false);
            self.console.print_lines(lines.join(""));
        }
        self.console.println("");
    }

    fn dump_local_disassembly(&mut self) {
        self.console.println("");

        let result = {
            let pc = self.cpu.registers.PC as usize;
            let local_segment = self.get_local_segment(pc);
            let disassembler = Disassembler::with_offset(local_segment.address);
            let pairs = disassembler.disassemble_with_addresses(&local_segment.code);
            self.highlight_lines(pc, pairs, local_segment.address, true)
        };
        self.console.print_lines(result.join(""));
        self.console.println("");
    }

    pub fn dump_memory_page(&mut self, page: usize) {
        let mut addr = page * 0x100;
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            self.console.print(format!("{:04X}: ", addr));
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.wrap_line();
            addr += 0x08;
        }
    }

    pub fn dump_memory(&mut self) {
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.wrap_line();
        }
    }

    fn dump_memory_range(&mut self, start: u16, end: u16) {
        self.console.println("");
        let start = start as usize;
        let end = end as usize;
        for chunk in self.cpu.memory[start..end + 0x01].chunks(8) {
            for b in chunk {
                self.console.print(format!("{:02X} ", *b));
            }
            self.console.println("");
        }
    }

    fn dump_registers(&mut self) {
        self.console.println("");
        self.console.println(format!("A: {} ({:04X})", self.cpu.registers.A, self.cpu.registers.A));
        self.console.println(format!("X: {} ({:04X})", self.cpu.registers.X, self.cpu.registers.X));
        self.console.println(format!("Y: {} ({:04X})", self.cpu.registers.Y, self.cpu.registers.Y));
        self.console.println(format!("PC: {} ({:04X})",
                                     self.cpu.registers.PC,
                                     self.cpu.registers.PC));
        self.console.println(format!("S: {} ({:04X})",
                                     self.cpu.stack.pointer,
                                     self.cpu.stack.pointer));
    }

    fn get_local_segment(&self, pc: usize) -> &CodeSegment {
        for segment in &self.segments {
            let addr = segment.address as usize;
            if pc >= addr && pc <= addr + segment.code.len() {
                return segment;
            }
        }

        &self.segments[0]
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       segment_start: u16,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc as isize - segment_start as isize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as isize) < start || (pair.1 as isize) > base + 0x0A {
                    continue;
                }
            }
            let current_line = pc as u16 == segment_start + pair.1;
            let breakpoint = self.breakpoints[segment_start as usize + pair.1 as usize] > 0x00;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else if !breakpoint && current_line {
                result.push(format!(">   {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate criterion;
extern crate os_pipe;
extern crate pipe;

use criterion::{black_box, Bencher, Criterion, ParameterizedBenchmark, Throughput};
use std::convert::TryInto;
use std::io::prelude::*;
use std::thread;

const TOTAL_TO_SEND: usize = 1 * 1024 * 1024;

fn send_recv_size<F, R, W>(mut f: F, reads: usize) -> impl FnMut(&mut Bencher, &&usize)
where
    F: FnMut() -> (R, W),
    F: Send + 'static,
    R: Read + Send + 'static,
    W: Write + Send + 'static,
{
    return move |b: &mut Bencher, &&size: &&usize| {
        let f = &mut f;
        let buf: Vec<u8> = (0..size).map(|i| i as u8).collect();
        b.iter(move || {
            let (mut reader, mut writer) = f();
            let t = thread::spawn(move || {
                let mut buf = vec![0; size / reads];
                while let Ok(_) = reader.read_exact(&mut buf) {}
            });

            for _ in 0..(TOTAL_TO_SEND / size) {
                writer.write_all(black_box(&buf[..])).unwrap();
            }
            drop(writer);
            t.join().expect("writing failed");
        })
    };
}

fn pipe_send(c: &mut Criterion) {
    const KB: usize = 1024;
    const SIZES: &[usize] = &[4 * KB, 64 * KB];
    //&[4 * KB, 8 * KB, 16 * KB, 32 * KB, 64 * KB],

    let bench = ParameterizedBenchmark::new( "pipe-rs", send_recv_size(|| pipe::pipe(), 1), SIZES)
        .throughput(|_| Throughput::Bytes(TOTAL_TO_SEND.try_into().unwrap()));
    c.bench("pipe_send", bench);

    let bench = ParameterizedBenchmark::new( "pipe-rs (16 reads per write)", send_recv_size(|| pipe::pipe(), 16), SIZES)
        .throughput(|_| Throughput::Bytes(TOTAL_TO_SEND.try_into().unwrap()));
    c.bench("pipe_send", bench);

    let bench = ParameterizedBenchmark::new( "os_pipe", send_recv_size(|| os_pipe::pipe().unwrap(), 1), SIZES)
        .throughput(|_| Throughput::Bytes(TOTAL_TO_SEND.try_into().unwrap()));
    c.bench("pipe_send", bench);

    let bench = ParameterizedBenchmark::new( "os_pipe (16 reads per write)", send_recv_size(|| os_pipe::pipe().unwrap(), 16), SIZES)
        .throughput(|_| Throughput::Bytes(TOTAL_TO_SEND.try_into().unwrap()));
    c.bench("pipe_send", bench);
}

criterion_group!(benches, pipe_send);
criterion_main!(benches);
<commit_msg>parameterize test by read count<commit_after>#[macro_use]
extern crate criterion;
extern crate os_pipe;
extern crate pipe;

use criterion::{black_box, Bencher, Criterion, ParameterizedBenchmark, Throughput};
use std::convert::TryInto;
use std::io::prelude::*;
use std::thread;

const TOTAL_TO_SEND: usize = 1 * 1024 * 1024;

fn send_recv_size<F, R, W>(mut f: F) -> impl FnMut(&mut Bencher, &&(usize, usize))
where
    F: FnMut() -> (R, W),
    F: Send + 'static,
    R: Read + Send + 'static,
    W: Write + Send + 'static,
{
    return move |b: &mut Bencher, &&(size, reads)| {
        let f = &mut f;
        let buf: Vec<u8> = (0..size).map(|i| i as u8).collect();
        b.iter(move || {
            let (mut reader, mut writer) = f();
            let t = thread::spawn(move || {
                let mut buf = vec![0; size / reads];
                while let Ok(_) = reader.read_exact(&mut buf) {}
            });

            for _ in 0..(TOTAL_TO_SEND / size) {
                writer.write_all(black_box(&buf[..])).unwrap();
            }
            drop(writer);
            t.join().expect("writing failed");
        })
    };
}

fn pipe_send(c: &mut Criterion) {
    const KB: usize = 1024;
    const SIZES: &[(usize, usize)] = &[
        (4 * KB, 1),
        (4 * KB, 16),
        (8 * KB, 1),
        (16 * KB, 1),
        (32 * KB, 1),
        (64 * KB, 1),
        (64 * KB, 16),
    ];
    let bench = ParameterizedBenchmark::new(
        "pipe-rs",
        send_recv_size(|| pipe::pipe()),
        SIZES,
    );
    let bench = bench
        .with_function("os_pipe", send_recv_size(|| os_pipe::pipe().unwrap()))
        .throughput(|_| Throughput::Bytes(TOTAL_TO_SEND.try_into().unwrap()));
    c.bench("pipe_send", bench);
}

criterion_group!(benches, pipe_send);
criterion_main!(benches);
<|endoftext|>


<commit_before>use std::io::{Read, Write};
use serialize::json::{ToJson, Json};
use std::str::{from_utf8};
use super::response::Response;
use super::event::{Event, is_event_json};
use super::request::Request;
use super::error::Error;
use std::borrow::ToOwned;

pub enum Message {
    Response(Response),
    Event(Event),
}

pub struct Handler<T> {
    socket: T,
    buffer: Vec<u8>,
}

impl<T> Handler<T> where T: Read + Write {
    pub fn new(socket: T) -> Handler<T> {
        Handler { socket: socket, buffer: Vec::new() }
    }

    pub fn read(&mut self) -> Result<Message, Error> {
        loop {
            if let Some((offset, len)) = self.find_message() {
                return self.make_message(offset, len);
            }

            try!(self.retrieve_from_socket());
        }
    }

    /// Retrieve new data from the socket
    fn retrieve_from_socket(&mut self) -> Result<(), Error> {
        let mut buf = [0; 1024];
        let bytes = try!(self.socket.read(&mut buf));
        for b in buf[..bytes].iter() {
            self.buffer.push(*b);
        }
        Ok(())
    }

    /// Find where a message is located
    fn find_message(&self) -> Option<(usize, usize)> {
        let mut offset = 0;
        let mut message_len = 0;

        while offset < self.buffer.len() {
            // check for a number
            if self.buffer[offset] < 0x40 && self.buffer[offset] >= 0x30 {
                message_len *= 10;
                message_len += (self.buffer[offset] - 0x30) as usize;
                offset += 1;

            // skip newline and carriage return
            } else if self.buffer[offset] == 0xa || self.buffer[offset] == 0xd {
                offset += 1;
            } else {
                break;
            }
        }

        if message_len > 0 && self.buffer.len() >= offset + message_len {
            debug!("Found message in buffer starting at {} with length {}", offset, message_len);
            Some((offset, message_len))
        } else {
            debug!("Found no complete message in buffer");
            None
        }
    }

    fn make_message(&mut self, offset: usize, length: usize) -> Result<Message, Error> {
        let end = offset + length;
        assert!(self.buffer.len() >= end);

        // check the result of our conversion
        let json_try = match from_utf8(&self.buffer[offset..end]) {
            Ok(json_str) => Ok(Json::from_str(json_str)),
            Err(e) => Err(e),
        };

        // first make sure we have a correct internal state
        self.buffer = self.buffer[offset+length..].to_owned(); // iter().collect();

        let json = try!(try!(json_try));

        if is_event_json(&json) {
            let evt = try!(Event::from_json(&json));
            debug!("Valid event received: {}", json);
            Ok(Message::Event(evt))
        } else {
            let resp = try!(Response::from_json(&json));
            debug!("Valid response received: {}", json);
            Ok(Message::Response(resp))
        }
    }

    pub fn write(&mut self, request: Request) -> Result<(), Error> {
        let encoded = request.to_json().to_string();
        debug!("Sending message: {}", encoded);

        let bytes = encoded.as_bytes();
        try!(self.socket.write_all(format!("{}", bytes.len()).as_bytes()));
        try!(self.socket.write_all(bytes));
        Ok(())
    }
}
<commit_msg>ASCII numbers run to 0x3a, not 0x40<commit_after>use std::io::{Read, Write};
use serialize::json::{ToJson, Json};
use std::str::{from_utf8};
use super::response::Response;
use super::event::{Event, is_event_json};
use super::request::Request;
use super::error::Error;
use std::borrow::ToOwned;

pub enum Message {
    Response(Response),
    Event(Event),
}

pub struct Handler<T> {
    socket: T,
    buffer: Vec<u8>,
}

impl<T> Handler<T> where T: Read + Write {
    pub fn new(socket: T) -> Handler<T> {
        Handler { socket: socket, buffer: Vec::new() }
    }

    pub fn read(&mut self) -> Result<Message, Error> {
        loop {
            if let Some((offset, len)) = self.find_message() {
                return self.make_message(offset, len);
            }

            try!(self.retrieve_from_socket());
        }
    }

    /// Retrieve new data from the socket
    fn retrieve_from_socket(&mut self) -> Result<(), Error> {
        let mut buf = [0; 1024];
        let bytes = try!(self.socket.read(&mut buf));
        for b in buf[..bytes].iter() {
            self.buffer.push(*b);
        }
        Ok(())
    }

    /// Find where a message is located
    fn find_message(&self) -> Option<(usize, usize)> {
        let mut offset = 0;
        let mut message_len = 0;

        while offset < self.buffer.len() {
            // check for a number
            if self.buffer[offset] < 0x3A && self.buffer[offset] >= 0x30 {
                message_len *= 10;
                message_len += (self.buffer[offset] - 0x30) as usize;
                offset += 1;

            // skip newline and carriage return
            } else if self.buffer[offset] == 0xa || self.buffer[offset] == 0xd {
                offset += 1;
            } else {
                break;
            }
        }

        if message_len > 0 && self.buffer.len() >= offset + message_len {
            debug!("Found message in buffer starting at {} with length {}", offset, message_len);
            Some((offset, message_len))
        } else {
            debug!("Found no complete message in buffer");
            None
        }
    }

    fn make_message(&mut self, offset: usize, length: usize) -> Result<Message, Error> {
        let end = offset + length;
        assert!(self.buffer.len() >= end);

        // check the result of our conversion
        let json_try = match from_utf8(&self.buffer[offset..end]) {
            Ok(json_str) => Ok(Json::from_str(json_str)),
            Err(e) => Err(e),
        };

        // first make sure we have a correct internal state
        self.buffer = self.buffer[offset+length..].to_owned(); // iter().collect();

        let json = try!(try!(json_try));

        if is_event_json(&json) {
            let evt = try!(Event::from_json(&json));
            debug!("Valid event received: {}", json);
            Ok(Message::Event(evt))
        } else {
            let resp = try!(Response::from_json(&json));
            debug!("Valid response received: {}", json);
            Ok(Message::Response(resp))
        }
    }

    pub fn write(&mut self, request: Request) -> Result<(), Error> {
        let encoded = request.to_json().to_string();
        debug!("Sending message: {}", encoded);

        let bytes = encoded.as_bytes();
        try!(self.socket.write_all(format!("{}", bytes.len()).as_bytes()));
        try!(self.socket.write_all(bytes));
        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::colors::{BitDepth, ColorType};
use crate::error::PngError;
use crate::PngResult;
use byteorder::{BigEndian, ReadBytesExt};
use crc::crc32;
use indexmap::IndexSet;
use std::io::Cursor;

#[derive(Debug, Clone, Copy)]
/// Headers from the IHDR chunk of the image
pub struct IhdrData {
    /// The width of the image in pixels
    pub width: u32,
    /// The height of the image in pixels
    pub height: u32,
    /// The color type of the image
    pub color_type: ColorType,
    /// The bit depth of the image
    pub bit_depth: BitDepth,
    /// The compression method used for this image (0 for DEFLATE)
    pub compression: u8,
    /// The filter mode used for this image (currently only 0 is valid)
    pub filter: u8,
    /// The interlacing mode of the image (0 = None, 1 = Adam7)
    pub interlaced: u8,
}

impl IhdrData {
    /// Bits per pixel
    pub fn bpp(&self) -> u8 {
        self.bit_depth.as_u8() * self.color_type.channels_per_pixel()
    }

    /// Byte length of IDAT that is correct for this IHDR
    pub fn raw_data_size(&self) -> usize {
        let w = self.width as usize;
        let h = self.height as usize;
        let bpp = self.bpp();

        fn bitmap_size(bpp: u8, w: usize, h: usize) -> usize {
            (((w / 8) * bpp as usize) + ((w & 7) * bpp as usize + 7) / 8) * h
        }

        if self.interlaced == 0 {
            bitmap_size(bpp, w, h) + h
        } else {
            let mut size = bitmap_size(bpp, (w + 7) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
            if w > 4 {
                size += bitmap_size(bpp, (w + 3) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
            }
            size += bitmap_size(bpp, (w + 3) >> 2, (h + 3) >> 3) + ((h + 3) >> 3);
            if w > 2 {
                size += bitmap_size(bpp, (w + 1) >> 2, (h + 3) >> 2) + ((h + 3) >> 2);
            }
            size += bitmap_size(bpp, (w + 1) >> 1, (h + 1) >> 2) + ((h + 1) >> 2);
            if w > 1 {
                size += bitmap_size(bpp, w >> 1, (h + 1) >> 1) + ((h + 1) >> 1);
            }
            size + bitmap_size(bpp, w, h >> 1) + (h >> 1)
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
/// Options to use for performing operations on headers (such as stripping)
pub enum Headers {
    /// None
    None,
    /// Remove specific chunks
    Strip(Vec<String>),
    /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
    Safe,
    /// Remove all non-critical chunks except these
    Keep(IndexSet<String>),
    /// All non-critical headers
    All,
}

#[inline]
pub fn file_header_is_valid(bytes: &[u8]) -> bool {
    let expected_header: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];

    *bytes == expected_header
}

#[derive(Debug, Clone, Copy)]
pub struct RawHeader<'a> {
    pub name: [u8; 4],
    pub data: &'a [u8],
}

pub fn parse_next_header<'a>(
    byte_data: &'a [u8],
    byte_offset: &mut usize,
    fix_errors: bool,
) -> PngResult<Option<RawHeader<'a>>> {
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let length = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_start = *byte_offset;
    let chunk_name = byte_data
        .get(header_start..header_start + 4)
        .ok_or(PngError::TruncatedData)?;
    if chunk_name == b"IEND" {
        // End of data
        return Ok(None);
    }
    *byte_offset += 4;

    let data = byte_data
        .get(*byte_offset..*byte_offset + length as usize)
        .ok_or(PngError::TruncatedData)?;
    *byte_offset += length as usize;
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let crc = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_bytes = byte_data
        .get(header_start..header_start + 4 + length as usize)
        .ok_or(PngError::TruncatedData)?;
    if !fix_errors && crc32::checksum_ieee(header_bytes) != crc {
        return Err(PngError::new(&format!(
            "CRC Mismatch in {} header; May be recoverable by using --fix",
            String::from_utf8_lossy(chunk_name)
        )));
    }

    let mut name = [0u8; 4];
    name.copy_from_slice(chunk_name);
    Ok(Some(RawHeader { name, data }))
}

pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult<IhdrData> {
    let mut rdr = Cursor::new(&byte_data[0..8]);
    Ok(IhdrData {
        color_type: match byte_data[9] {
            0 => ColorType::Grayscale,
            2 => ColorType::RGB,
            3 => ColorType::Indexed,
            4 => ColorType::GrayscaleAlpha,
            6 => ColorType::RGBA,
            _ => return Err(PngError::new("Unexpected color type in header")),
        },
        bit_depth: match byte_data[8] {
            1 => BitDepth::One,
            2 => BitDepth::Two,
            4 => BitDepth::Four,
            8 => BitDepth::Eight,
            16 => BitDepth::Sixteen,
            _ => return Err(PngError::new("Unexpected bit depth in header")),
        },
        width: rdr.read_u32::<BigEndian>().unwrap(),
        height: rdr.read_u32::<BigEndian>().unwrap(),
        compression: byte_data[10],
        filter: byte_data[11],
        interlaced: byte_data[12],
    })
}
<commit_msg>Prevent panics in parse_ihdr_header (#392)<commit_after>use crate::colors::{BitDepth, ColorType};
use crate::error::PngError;
use crate::PngResult;
use byteorder::{BigEndian, ReadBytesExt};
use crc::crc32;
use indexmap::IndexSet;
use std::io::Cursor;

#[derive(Debug, Clone, Copy)]
/// Headers from the IHDR chunk of the image
pub struct IhdrData {
    /// The width of the image in pixels
    pub width: u32,
    /// The height of the image in pixels
    pub height: u32,
    /// The color type of the image
    pub color_type: ColorType,
    /// The bit depth of the image
    pub bit_depth: BitDepth,
    /// The compression method used for this image (0 for DEFLATE)
    pub compression: u8,
    /// The filter mode used for this image (currently only 0 is valid)
    pub filter: u8,
    /// The interlacing mode of the image (0 = None, 1 = Adam7)
    pub interlaced: u8,
}

impl IhdrData {
    /// Bits per pixel
    pub fn bpp(&self) -> u8 {
        self.bit_depth.as_u8() * self.color_type.channels_per_pixel()
    }

    /// Byte length of IDAT that is correct for this IHDR
    pub fn raw_data_size(&self) -> usize {
        let w = self.width as usize;
        let h = self.height as usize;
        let bpp = self.bpp();

        fn bitmap_size(bpp: u8, w: usize, h: usize) -> usize {
            (((w / 8) * bpp as usize) + ((w & 7) * bpp as usize + 7) / 8) * h
        }

        if self.interlaced == 0 {
            bitmap_size(bpp, w, h) + h
        } else {
            let mut size = bitmap_size(bpp, (w + 7) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
            if w > 4 {
                size += bitmap_size(bpp, (w + 3) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
            }
            size += bitmap_size(bpp, (w + 3) >> 2, (h + 3) >> 3) + ((h + 3) >> 3);
            if w > 2 {
                size += bitmap_size(bpp, (w + 1) >> 2, (h + 3) >> 2) + ((h + 3) >> 2);
            }
            size += bitmap_size(bpp, (w + 1) >> 1, (h + 1) >> 2) + ((h + 1) >> 2);
            if w > 1 {
                size += bitmap_size(bpp, w >> 1, (h + 1) >> 1) + ((h + 1) >> 1);
            }
            size + bitmap_size(bpp, w, h >> 1) + (h >> 1)
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
/// Options to use for performing operations on headers (such as stripping)
pub enum Headers {
    /// None
    None,
    /// Remove specific chunks
    Strip(Vec<String>),
    /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
    Safe,
    /// Remove all non-critical chunks except these
    Keep(IndexSet<String>),
    /// All non-critical headers
    All,
}

#[inline]
pub fn file_header_is_valid(bytes: &[u8]) -> bool {
    let expected_header: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];

    *bytes == expected_header
}

#[derive(Debug, Clone, Copy)]
pub struct RawHeader<'a> {
    pub name: [u8; 4],
    pub data: &'a [u8],
}

pub fn parse_next_header<'a>(
    byte_data: &'a [u8],
    byte_offset: &mut usize,
    fix_errors: bool,
) -> PngResult<Option<RawHeader<'a>>> {
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let length = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_start = *byte_offset;
    let chunk_name = byte_data
        .get(header_start..header_start + 4)
        .ok_or(PngError::TruncatedData)?;
    if chunk_name == b"IEND" {
        // End of data
        return Ok(None);
    }
    *byte_offset += 4;

    let data = byte_data
        .get(*byte_offset..*byte_offset + length as usize)
        .ok_or(PngError::TruncatedData)?;
    *byte_offset += length as usize;
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let crc = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_bytes = byte_data
        .get(header_start..header_start + 4 + length as usize)
        .ok_or(PngError::TruncatedData)?;
    if !fix_errors && crc32::checksum_ieee(header_bytes) != crc {
        return Err(PngError::new(&format!(
            "CRC Mismatch in {} header; May be recoverable by using --fix",
            String::from_utf8_lossy(chunk_name)
        )));
    }

    let mut name = [0u8; 4];
    name.copy_from_slice(chunk_name);
    Ok(Some(RawHeader { name, data }))
}

pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult<IhdrData> {
    // This eliminates bounds checks for the rest of the function
    let interlaced = byte_data.get(12).copied().ok_or(PngError::TruncatedData)?;
    let mut rdr = Cursor::new(&byte_data[0..8]);
    Ok(IhdrData {
        color_type: match byte_data[9] {
            0 => ColorType::Grayscale,
            2 => ColorType::RGB,
            3 => ColorType::Indexed,
            4 => ColorType::GrayscaleAlpha,
            6 => ColorType::RGBA,
            _ => return Err(PngError::new("Unexpected color type in header")),
        },
        bit_depth: match byte_data[8] {
            1 => BitDepth::One,
            2 => BitDepth::Two,
            4 => BitDepth::Four,
            8 => BitDepth::Eight,
            16 => BitDepth::Sixteen,
            _ => return Err(PngError::new("Unexpected bit depth in header")),
        },
        width: rdr.read_u32::<BigEndian>().unwrap(),
        height: rdr.read_u32::<BigEndian>().unwrap(),
        compression: byte_data[10],
        filter: byte_data[11],
        interlaced,
    })
}
<|endoftext|>