<commit_before>// Copyright (C) 2014 The 6502-rs Developers
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the names of the copyright holders nor the names of any
//    contributors may be used to endorse or promote products derived from this
//    software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

use log;

use std;

use address::{AddressDiff};
use instruction;
use instruction::{DecodedInstr};
use memory::Memory;
use registers::{ Registers, Status, StatusArgs };
use registers::{ PS_NEGATIVE, PS_OVERFLOW, PS_ZERO, PS_CARRY };

pub struct Machine {
    pub registers: Registers,
    pub memory:    Memory
}

impl Machine {
    pub fn new() -> Machine {
    	Machine{
    	    registers: Registers::new(),
    	    memory:    Memory::new()
    	}
    }

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

    pub fn fetch_next_and_decode(&mut self) -> Option<DecodedInstr> {
        let x: u8 = self.memory.get_byte(self.registers.program_counter);

        match instruction::OPCODES[x as uint] {
            Some((instr, am)) => {
                let extra_bytes = am.extra_bytes();
                let num_bytes = AddressDiff(1) + extra_bytes;

                let data_start = self.registers.program_counter
                               + AddressDiff(1);

                let slice = self.memory.get_slice(data_start, extra_bytes);
                let am_out = am.process(self, slice);

                // Increment program counter
                self.registers.program_counter =
                    self.registers.program_counter + num_bytes;

                Some((instr, am_out))
            }
            _ => None
        }
    }

    pub fn execute_instruction(&mut self, decoded_instr: DecodedInstr) {
        match decoded_instr {
            (instruction::ADC, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "add with carry immediate: {}", val);
                self.add_with_carry(val as i8);
            },
            (instruction::ADC, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr) as i8;
                log!(log::DEBUG, "add with carry. address: {}. value: {}",
                                 addr, val);
                self.add_with_carry(val);
            },

            (instruction::BMI, instruction::UseRelative(rel)) => {
                let addr = self.registers.program_counter + rel;
                log!(log::DEBUG, "branch if minus relative. address: {}", addr);
                self.branch_if_minus(addr);
            },

            (instruction::LDA, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load A immediate: {}", val);
                self.load_accumulator(val as i8);
            },
            (instruction::LDA, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load A. address: {}. value: {}", addr, val);
                self.load_accumulator(val as i8);
            },

            (instruction::LDX, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load X immediate: {}", val);
                self.load_x_register(val as i8);
            },
            (instruction::LDX, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load X. address: {}. value: {}", addr, val);
                self.load_x_register(val as i8);
            },

            (instruction::LDY, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load Y immediate: {}", val);
                self.load_y_register(val as i8);
            },
            (instruction::LDY, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load Y. address: {}. value: {}", addr, val);
                self.load_y_register(val as i8);
            },

            (instruction::NOP, _) => {
                log!(log::DEBUG, "nop instr");
            },
            (_, _) => {
                log!(log::DEBUG, "attempting to execute unimplemented \
                                  instruction");
            },
        };
    }

    pub fn run(&mut self) {
        loop {
            if let Some(decoded_instr) = self.fetch_next_and_decode() {
                self.execute_instruction(decoded_instr);
            } else {
                break
            }
        }
    }

    fn load_register_with_flags(register: &mut i8,
                                status: &mut Status,
                                value: i8) {
        *register = value;

        let is_zero = value == 0;
        let is_negative = value < 0;

        status.set_with_mask(
            PS_ZERO | PS_NEGATIVE,
            Status::new(StatusArgs { zero: is_zero,
                                     negative: is_negative,
                                     ..StatusArgs::none() } ));
    }

    pub fn load_x_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_x,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_y_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_y,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_accumulator(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.accumulator,
                                          &mut self.registers.status,
                                          value);
    }

    // TODO akeeton: Implement binary-coded decimal.
    pub fn add_with_carry(&mut self, value: i8) {
        let a_before: i8 = self.registers.accumulator;
        let c_before: i8 = self.registers.status.get_carry();
        let a_after: i8 = a_before + c_before + value;

        debug_assert_eq!(a_after as u8, a_before as u8 + c_before as u8
                                        + value as u8);

        let did_carry = (a_after as u8) < (a_before as u8);

        let did_overflow   =
        	   (a_before < 0 && value < 0 && a_after >= 0)
        	|| (a_before > 0 && value > 0 && a_after <= 0);

        let mask = PS_CARRY | PS_OVERFLOW;

        self.registers.status.set_with_mask(mask,
            Status::new(StatusArgs { carry: did_carry,
                                     overflow: did_overflow,
                                     ..StatusArgs::none() } ));

        self.load_accumulator(a_after);

        log!(log::DEBUG, "accumulator: {}", self.registers.accumulator);
    }

    pub fn branch_if_minus(&mut self, addr: Address) {
        if (self.registers.status.contains(PS_NEGATIVE)) {
            self.registers.program_counter = addr;
        }
    }
}

impl std::fmt::Show for Machine {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Machine Dump:\n\nAccumulator: {}",
               self.registers.accumulator)
    }
}

#[test]
fn add_with_carry_test() {
    let mut machine = Machine::new();

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-1);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-127);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),     true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),      true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.registers.status.remove(PS_CARRY);
    machine.add_with_carry(-128);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW),  true);
}

#[test]
fn branch_if_minus_test() {
    {
        let mut machine      = Machine::new();
        let registers_before = machine.registers;

        machine.branch_if_minus(Address(0xABCD));
        assert_eq!(machine.registers, registers_before);
        assert_eq!(machine.registers.program_counter, Address(0));
    }

    {
        let mut machine = Machine::new();

        machine.registers.status.set_with_mask(PS_NEGATIVE, PS_NEGATIVE);
        let registers_before = machine.registers;

        machine.branch_if_minus(Address(0xABCD));
        assert_eq!(machine.registers, registers_before);
        assert_eq!(machine.registers.program_counter, Address(0xABCD));
    }
}
<commit_msg>branch_fi_minus_test() passes.<commit_after>// Copyright (C) 2014 The 6502-rs Developers
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the names of the copyright holders nor the names of any
//    contributors may be used to endorse or promote products derived from this
//    software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

use log;

use std;

use address::{AddressDiff};
use instruction;
use instruction::{DecodedInstr};
use memory::Memory;
use registers::{ Registers, Status, StatusArgs };
use registers::{ PS_NEGATIVE, PS_OVERFLOW, PS_ZERO, PS_CARRY };

pub struct Machine {
    pub registers: Registers,
    pub memory:    Memory
}

impl Machine {
    pub fn new() -> Machine {
    	Machine{
    	    registers: Registers::new(),
    	    memory:    Memory::new()
    	}
    }

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

    pub fn fetch_next_and_decode(&mut self) -> Option<DecodedInstr> {
        let x: u8 = self.memory.get_byte(self.registers.program_counter);

        match instruction::OPCODES[x as uint] {
            Some((instr, am)) => {
                let extra_bytes = am.extra_bytes();
                let num_bytes = AddressDiff(1) + extra_bytes;

                let data_start = self.registers.program_counter
                               + AddressDiff(1);

                let slice = self.memory.get_slice(data_start, extra_bytes);
                let am_out = am.process(self, slice);

                // Increment program counter
                self.registers.program_counter =
                    self.registers.program_counter + num_bytes;

                Some((instr, am_out))
            }
            _ => None
        }
    }

    pub fn execute_instruction(&mut self, decoded_instr: DecodedInstr) {
        match decoded_instr {
            (instruction::ADC, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "add with carry immediate: {}", val);
                self.add_with_carry(val as i8);
            },
            (instruction::ADC, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr) as i8;
                log!(log::DEBUG, "add with carry. address: {}. value: {}",
                                 addr, val);
                self.add_with_carry(val);
            },

            (instruction::BMI, instruction::UseRelative(rel)) => {
                let addr = self.registers.program_counter + rel;
                log!(log::DEBUG, "branch if minus relative. address: {}", addr);
                self.branch_if_minus(addr);
            },

            (instruction::LDA, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load A immediate: {}", val);
                self.load_accumulator(val as i8);
            },
            (instruction::LDA, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load A. address: {}. value: {}", addr, val);
                self.load_accumulator(val as i8);
            },

            (instruction::LDX, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load X immediate: {}", val);
                self.load_x_register(val as i8);
            },
            (instruction::LDX, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load X. address: {}. value: {}", addr, val);
                self.load_x_register(val as i8);
            },

            (instruction::LDY, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load Y immediate: {}", val);
                self.load_y_register(val as i8);
            },
            (instruction::LDY, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load Y. address: {}. value: {}", addr, val);
                self.load_y_register(val as i8);
            },

            (instruction::NOP, _) => {
                log!(log::DEBUG, "nop instr");
            },
            (_, _) => {
                log!(log::DEBUG, "attempting to execute unimplemented \
                                  instruction");
            },
        };
    }

    pub fn run(&mut self) {
        loop {
            if let Some(decoded_instr) = self.fetch_next_and_decode() {
                self.execute_instruction(decoded_instr);
            } else {
                break
            }
        }
    }

    fn load_register_with_flags(register: &mut i8,
                                status: &mut Status,
                                value: i8) {
        *register = value;

        let is_zero = value == 0;
        let is_negative = value < 0;

        status.set_with_mask(
            PS_ZERO | PS_NEGATIVE,
            Status::new(StatusArgs { zero: is_zero,
                                     negative: is_negative,
                                     ..StatusArgs::none() } ));
    }

    pub fn load_x_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_x,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_y_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_y,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_accumulator(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.accumulator,
                                          &mut self.registers.status,
                                          value);
    }

    // TODO akeeton: Implement binary-coded decimal.
    pub fn add_with_carry(&mut self, value: i8) {
        let a_before: i8 = self.registers.accumulator;
        let c_before: i8 = self.registers.status.get_carry();
        let a_after: i8 = a_before + c_before + value;

        debug_assert_eq!(a_after as u8, a_before as u8 + c_before as u8
                                        + value as u8);

        let did_carry = (a_after as u8) < (a_before as u8);

        let did_overflow   =
        	   (a_before < 0 && value < 0 && a_after >= 0)
        	|| (a_before > 0 && value > 0 && a_after <= 0);

        let mask = PS_CARRY | PS_OVERFLOW;

        self.registers.status.set_with_mask(mask,
            Status::new(StatusArgs { carry: did_carry,
                                     overflow: did_overflow,
                                     ..StatusArgs::none() } ));

        self.load_accumulator(a_after);

        log!(log::DEBUG, "accumulator: {}", self.registers.accumulator);
    }

    pub fn branch_if_minus(&mut self, addr: Address) {
        if (self.registers.status.contains(PS_NEGATIVE)) {
            self.registers.program_counter = addr;
        }
    }
}

impl std::fmt::Show for Machine {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Machine Dump:\n\nAccumulator: {}",
               self.registers.accumulator)
    }
}

#[test]
fn add_with_carry_test() {
    let mut machine = Machine::new();

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-1);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-127);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),     true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),      true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.registers.status.remove(PS_CARRY);
    machine.add_with_carry(-128);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW),  true);
}

#[test]
fn branch_if_minus_test() {
    {
        let mut machine      = Machine::new();
        let registers_before = machine.registers;

        machine.branch_if_minus(Address(0xABCD));
        assert_eq!(machine.registers, registers_before);
        assert_eq!(machine.registers.program_counter, Address(0));
    }

    {
        let mut machine = Machine::new();

        machine.registers.status.set_with_mask(PS_NEGATIVE, PS_NEGATIVE);
        let registers_before = machine.registers;

        machine.branch_if_minus(Address(0xABCD));
        assert_eq!(machine.registers.status, registers_before.status);
        assert_eq!(machine.registers.program_counter, Address(0xABCD));
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Untyped root container for a Cap'n Proto value.

use any_pointer;
use private::units::*;
use private::arena::{BuilderArena, ReaderArena, SegmentBuilder, SegmentReader};
use private::layout;
use traits::{FromPointerReader, FromPointerBuilder, SetPointerBuilder};
use {OutputSegments, Result, Word};

/// Options controlling how data is read.
#[derive(Clone, Copy)]
pub struct ReaderOptions {

    /// Limits how many total words of data are allowed to be traversed. Traversal is counted when
    /// a new struct or list builder is obtained, e.g. from a get() accessor. This means that calling
    /// the getter for the same sub-struct multiple times will cause it to be double-counted. Once
    /// the traversal limit is reached, an error will be reported.
    ///
    /// This limit exists for security reasons. It is possible for an attacker to construct a message
    /// in which multiple pointers point at the same location. This is technically invalid, but hard
    /// to detect. Using such a message, an attacker could cause a message which is small on the wire
    /// to appear much larger when actually traversed, possibly exhausting server resources leading to
    /// denial-of-service.
    ///
    /// It makes sense to set a traversal limit that is much larger than the underlying message.
    /// Together with sensible coding practices (e.g. trying to avoid calling sub-object getters
    /// multiple times, which is expensive anyway), this should provide adequate protection without
    /// inconvenience.
    pub traversal_limit_in_words : u64,

    /// Limits how deeply nested a message structure can be, e.g. structs containing other structs or
    /// lists of structs.
    ///
    /// Like the traversal limit, this limit exists for security reasons. Since it is common to use
    /// recursive code to traverse recursive data structures, an attacker could easily cause a stack
    /// overflow by sending a very-depply-nested (or even cyclic) message, without the message even
    /// being very large. The default limit of 64 is probably low enough to prevent any chance of
    /// stack overflow, yet high enough that it is never a problem in practice.
    pub nesting_limit : i32,
}

pub const DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversal_limit_in_words : 8 * 1024 * 1024, nesting_limit : 64 };

impl ReaderOptions {
    pub fn new() -> ReaderOptions { DEFAULT_READER_OPTIONS }

    pub fn nesting_limit<'a>(&'a mut self, value : i32) -> &'a mut ReaderOptions {
        self.nesting_limit = value;
        return self;
    }

    pub fn traversal_limit_in_words<'a>(&'a mut self, value : u64) -> &'a mut ReaderOptions {
        self.traversal_limit_in_words = value;
        return self;
    }
}

type SegmentId = u32;

/// An object that manages the buffers underlying a Cap'n Proto message reader.
pub trait ReaderSegments {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]>;
}

/// An array of segments.
pub struct SegmentArray<'a> {
    segments: &'a [&'a [Word]],
}

impl <'a> SegmentArray<'a> {
    pub fn new(segments: &'a [&'a [Word]]) -> SegmentArray<'a> {
        SegmentArray { segments: segments }
    }
}

impl <'b> ReaderSegments for SegmentArray<'b> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        self.segments.get(id as usize).map(|slice| *slice)
    }
}

/// A container used to read a message.
///
/// The underlying implemention uses the `ReaderSegments` as a trait object. However, we
/// need to include `S` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Reader<S> where S: ReaderSegments {
    arena: Box<ReaderArena>,
    segments: Box<S>,
    options: ReaderOptions,
}

unsafe impl <S> Send for Reader<S> where S: Send + ReaderSegments {}

impl <S> Reader<S> where S: ReaderSegments {
    pub fn new(segments: S, options: ReaderOptions) -> Reader<S> {
        let boxed_segments = Box::new(segments);
        let boxed_segments_ref = {
            let r :&S = &*boxed_segments;
            unsafe { ::std::mem::transmute(r as &ReaderSegments) }
        };
        let arena = ReaderArena::new(boxed_segments_ref, options);
        Reader { arena: arena, segments: boxed_segments, options: options }
    }

    fn get_root_internal(&self) -> Result<any_pointer::Reader> {
        unsafe {
            let segment : *const SegmentReader = &self.arena.segment0;

            let pointer_reader = try!(layout::PointerReader::get_root(
                segment, (*segment).get_start_ptr(), self.options.nesting_limit));

            Ok(any_pointer::Reader::new(pointer_reader))
        }
    }

    /// Gets the root of the message, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerReader<'a>>(&'a self) -> Result<T> {
        try!(self.get_root_internal()).get_as()
    }

    pub fn into_segments(self) -> S {
        *self.segments
    }
}

/// An object that allocates memory for a Cap'n Proto message as it is being built.
pub unsafe trait Allocator {
    /// Allocates memory for a new segment, returning a pointer to the start of the segment
    /// and a u32 indicating the length of the segment.
    ///
    /// UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    /// for the lifetime of the object and doesn't overlap with other allocated memory.
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32);

    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
}

/// A container used to build a message.
///
/// The underlying implemention uses the `Allocator` as a trait object. However, we
/// need to include `A` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Builder<A> where A: Allocator {
    arena: Box<BuilderArena>,
    allocator: Box<A>,
}

unsafe impl <A> Send for Builder<A> where A: Send + Allocator {}

impl <A> Builder<A> where A: Allocator {
    pub fn new(allocator: A) -> Builder<A> {
        let mut boxed_allocator = Box::new(allocator);
        let boxed_allocator_ref = {
            let r: &mut A = &mut *boxed_allocator;
            unsafe { ::std::mem::transmute(r as &mut Allocator) }
        };
        let arena = BuilderArena::new(boxed_allocator_ref);
        Builder { arena: arena, allocator: boxed_allocator }
    }

    fn get_root_internal<'a>(&mut self) -> any_pointer::Builder<'a> {
        let root_segment: *mut SegmentBuilder = &mut self.arena.segment0;

        if self.arena.segment0.current_size() == 0 {
            match self.arena.segment0.allocate(WORDS_PER_POINTER as u32) {
                None => {panic!("could not allocate root pointer") }
                Some(location) => {
                    assert!(location == self.arena.segment0.get_ptr_unchecked(0),
                            "First allocated word of new segment was not at offset 0");

                    any_pointer::Builder::new(layout::PointerBuilder::get_root(root_segment, location))

                }
            }
        } else {
            any_pointer::Builder::new(
                layout::PointerBuilder::get_root(root_segment,
                                                 self.arena.segment0.get_ptr_unchecked(0)))
        }

    }

    /// Initializes the root as a value of the given type.
    pub fn init_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> T {
        self.get_root_internal().init_as()
    }

    /// Gets the root, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> Result<T> {
        self.get_root_internal().get_as()
    }

    /// Sets the root to a deep copy of the given value.
    pub fn set_root<To, From : SetPointerBuilder<To>>(&mut self, value : From) -> Result<()> {
        self.get_root_internal().set_as(value)
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        self.arena.get_segments_for_output()
    }
}

impl <A> Drop for Builder<A> where A: Allocator {
    fn drop(&mut self) {
        self.allocator.pre_drop(self.arena.segment0.current_size());
    }
}

pub struct HeapAllocator {
    owned_memory : Vec<Vec<Word>>,
    next_size: u32,
    allocation_strategy: AllocationStrategy,
}

#[derive(Clone, Copy)]
pub enum AllocationStrategy {
    FixedSize,
    GrowHeuristically
}

pub const SUGGESTED_FIRST_SEGMENT_WORDS : u32 = 1024;
pub const SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = AllocationStrategy::GrowHeuristically;

impl HeapAllocator {
    pub fn new() -> HeapAllocator {
        HeapAllocator { owned_memory: Vec::new(),
                        next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
                        allocation_strategy: SUGGESTED_ALLOCATION_STRATEGY }
    }

    pub fn first_segment_words(mut self, value: u32) -> HeapAllocator {
        self.next_size = value;
        self
    }

    pub fn allocation_strategy(mut self, value : AllocationStrategy) -> HeapAllocator {
        self.allocation_strategy = value;
        self
    }
}

unsafe impl Allocator for HeapAllocator {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        let size = ::std::cmp::max(minimum_size, self.next_size);
        let mut new_words = Word::allocate_zeroed_vec(size as usize);
        let ptr = new_words.as_mut_ptr();
        self.owned_memory.push(new_words);

        match self.allocation_strategy {
            AllocationStrategy::GrowHeuristically => { self.next_size += size; }
            _ => { }
        }
        (ptr, size as u32)
    }
}

impl Builder<HeapAllocator> {
    pub fn new_default() -> Builder<HeapAllocator> {
        Builder::new(HeapAllocator::new())
    }
}

pub struct ScratchSpace<'a> {
    slice: &'a mut [Word],
    in_use: bool,
}

impl <'a> ScratchSpace<'a> {
    pub fn new(slice: &'a mut [Word]) -> ScratchSpace<'a> {
        ScratchSpace { slice: slice, in_use: false }
    }
}

pub struct ScratchSpaceHeapAllocator<'a, 'b: 'a> {
    scratch_space: &'a mut ScratchSpace<'b>,
    allocator: HeapAllocator,
}

impl <'a, 'b: 'a> ScratchSpaceHeapAllocator<'a, 'b> {
    pub fn new(scratch_space: &'a mut ScratchSpace<'b>) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: scratch_space,
                                    allocator: HeapAllocator::new()}
    }

    pub fn second_segment_words(mut self, value: u32) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.first_segment_words(value) }

    }

    pub fn allocation_strategy(mut self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.allocation_strategy(value) }
    }

}

unsafe impl <'a, 'b: 'a> Allocator for ScratchSpaceHeapAllocator<'a, 'b> {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        if !self.scratch_space.in_use {
            self.scratch_space.in_use = true;
            (self.scratch_space.slice.as_mut_ptr(), self.scratch_space.slice.len() as u32)
        } else {
            self.allocator.allocate_segment(minimum_size)
        }
    }

    fn pre_drop(&mut self, segment0_currently_allocated: u32) {
        let ptr = self.scratch_space.slice.as_mut_ptr();
        unsafe {
            ::std::ptr::write_bytes(ptr, 0u8, segment0_currently_allocated as usize);
        }
        self.scratch_space.in_use = false;
    }
}
<commit_msg>impl Default for ReaderOptions<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Untyped root container for a Cap'n Proto value.

use any_pointer;
use private::units::*;
use private::arena::{BuilderArena, ReaderArena, SegmentBuilder, SegmentReader};
use private::layout;
use traits::{FromPointerReader, FromPointerBuilder, SetPointerBuilder};
use {OutputSegments, Result, Word};

/// Options controlling how data is read.
#[derive(Clone, Copy)]
pub struct ReaderOptions {

    /// Limits how many total words of data are allowed to be traversed. Traversal is counted when
    /// a new struct or list builder is obtained, e.g. from a get() accessor. This means that calling
    /// the getter for the same sub-struct multiple times will cause it to be double-counted. Once
    /// the traversal limit is reached, an error will be reported.
    ///
    /// This limit exists for security reasons. It is possible for an attacker to construct a message
    /// in which multiple pointers point at the same location. This is technically invalid, but hard
    /// to detect. Using such a message, an attacker could cause a message which is small on the wire
    /// to appear much larger when actually traversed, possibly exhausting server resources leading to
    /// denial-of-service.
    ///
    /// It makes sense to set a traversal limit that is much larger than the underlying message.
    /// Together with sensible coding practices (e.g. trying to avoid calling sub-object getters
    /// multiple times, which is expensive anyway), this should provide adequate protection without
    /// inconvenience.
    pub traversal_limit_in_words : u64,

    /// Limits how deeply nested a message structure can be, e.g. structs containing other structs or
    /// lists of structs.
    ///
    /// Like the traversal limit, this limit exists for security reasons. Since it is common to use
    /// recursive code to traverse recursive data structures, an attacker could easily cause a stack
    /// overflow by sending a very-depply-nested (or even cyclic) message, without the message even
    /// being very large. The default limit of 64 is probably low enough to prevent any chance of
    /// stack overflow, yet high enough that it is never a problem in practice.
    pub nesting_limit : i32,
}

pub const DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversal_limit_in_words : 8 * 1024 * 1024, nesting_limit : 64 };

impl Default for ReaderOptions {
    fn default() -> ReaderOptions {
        ReaderOptions { traversal_limit_in_words : 8 * 1024 * 1024, nesting_limit : 64 }
    }
}

impl ReaderOptions {
    pub fn new() -> ReaderOptions { DEFAULT_READER_OPTIONS }

    pub fn nesting_limit<'a>(&'a mut self, value : i32) -> &'a mut ReaderOptions {
        self.nesting_limit = value;
        return self;
    }

    pub fn traversal_limit_in_words<'a>(&'a mut self, value : u64) -> &'a mut ReaderOptions {
        self.traversal_limit_in_words = value;
        return self;
    }
}

type SegmentId = u32;

/// An object that manages the buffers underlying a Cap'n Proto message reader.
pub trait ReaderSegments {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]>;
}

/// An array of segments.
pub struct SegmentArray<'a> {
    segments: &'a [&'a [Word]],
}

impl <'a> SegmentArray<'a> {
    pub fn new(segments: &'a [&'a [Word]]) -> SegmentArray<'a> {
        SegmentArray { segments: segments }
    }
}

impl <'b> ReaderSegments for SegmentArray<'b> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        self.segments.get(id as usize).map(|slice| *slice)
    }
}

/// A container used to read a message.
///
/// The underlying implemention uses the `ReaderSegments` as a trait object. However, we
/// need to include `S` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Reader<S> where S: ReaderSegments {
    arena: Box<ReaderArena>,
    segments: Box<S>,
    options: ReaderOptions,
}

unsafe impl <S> Send for Reader<S> where S: Send + ReaderSegments {}

impl <S> Reader<S> where S: ReaderSegments {
    pub fn new(segments: S, options: ReaderOptions) -> Reader<S> {
        let boxed_segments = Box::new(segments);
        let boxed_segments_ref = {
            let r :&S = &*boxed_segments;
            unsafe { ::std::mem::transmute(r as &ReaderSegments) }
        };
        let arena = ReaderArena::new(boxed_segments_ref, options);
        Reader { arena: arena, segments: boxed_segments, options: options }
    }

    fn get_root_internal(&self) -> Result<any_pointer::Reader> {
        unsafe {
            let segment : *const SegmentReader = &self.arena.segment0;

            let pointer_reader = try!(layout::PointerReader::get_root(
                segment, (*segment).get_start_ptr(), self.options.nesting_limit));

            Ok(any_pointer::Reader::new(pointer_reader))
        }
    }

    /// Gets the root of the message, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerReader<'a>>(&'a self) -> Result<T> {
        try!(self.get_root_internal()).get_as()
    }

    pub fn into_segments(self) -> S {
        *self.segments
    }
}

/// An object that allocates memory for a Cap'n Proto message as it is being built.
pub unsafe trait Allocator {
    /// Allocates memory for a new segment, returning a pointer to the start of the segment
    /// and a u32 indicating the length of the segment.
    ///
    /// UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    /// for the lifetime of the object and doesn't overlap with other allocated memory.
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32);

    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
}

/// A container used to build a message.
///
/// The underlying implemention uses the `Allocator` as a trait object. However, we
/// need to include `A` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Builder<A> where A: Allocator {
    arena: Box<BuilderArena>,
    allocator: Box<A>,
}

unsafe impl <A> Send for Builder<A> where A: Send + Allocator {}

impl <A> Builder<A> where A: Allocator {
    pub fn new(allocator: A) -> Builder<A> {
        let mut boxed_allocator = Box::new(allocator);
        let boxed_allocator_ref = {
            let r: &mut A = &mut *boxed_allocator;
            unsafe { ::std::mem::transmute(r as &mut Allocator) }
        };
        let arena = BuilderArena::new(boxed_allocator_ref);
        Builder { arena: arena, allocator: boxed_allocator }
    }

    fn get_root_internal<'a>(&mut self) -> any_pointer::Builder<'a> {
        let root_segment: *mut SegmentBuilder = &mut self.arena.segment0;

        if self.arena.segment0.current_size() == 0 {
            match self.arena.segment0.allocate(WORDS_PER_POINTER as u32) {
                None => {panic!("could not allocate root pointer") }
                Some(location) => {
                    assert!(location == self.arena.segment0.get_ptr_unchecked(0),
                            "First allocated word of new segment was not at offset 0");

                    any_pointer::Builder::new(layout::PointerBuilder::get_root(root_segment, location))

                }
            }
        } else {
            any_pointer::Builder::new(
                layout::PointerBuilder::get_root(root_segment,
                                                 self.arena.segment0.get_ptr_unchecked(0)))
        }

    }

    /// Initializes the root as a value of the given type.
    pub fn init_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> T {
        self.get_root_internal().init_as()
    }

    /// Gets the root, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> Result<T> {
        self.get_root_internal().get_as()
    }

    /// Sets the root to a deep copy of the given value.
    pub fn set_root<To, From : SetPointerBuilder<To>>(&mut self, value : From) -> Result<()> {
        self.get_root_internal().set_as(value)
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        self.arena.get_segments_for_output()
    }
}

impl <A> Drop for Builder<A> where A: Allocator {
    fn drop(&mut self) {
        self.allocator.pre_drop(self.arena.segment0.current_size());
    }
}

pub struct HeapAllocator {
    owned_memory : Vec<Vec<Word>>,
    next_size: u32,
    allocation_strategy: AllocationStrategy,
}

#[derive(Clone, Copy)]
pub enum AllocationStrategy {
    FixedSize,
    GrowHeuristically
}

pub const SUGGESTED_FIRST_SEGMENT_WORDS : u32 = 1024;
pub const SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = AllocationStrategy::GrowHeuristically;

impl HeapAllocator {
    pub fn new() -> HeapAllocator {
        HeapAllocator { owned_memory: Vec::new(),
                        next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
                        allocation_strategy: SUGGESTED_ALLOCATION_STRATEGY }
    }

    pub fn first_segment_words(mut self, value: u32) -> HeapAllocator {
        self.next_size = value;
        self
    }

    pub fn allocation_strategy(mut self, value : AllocationStrategy) -> HeapAllocator {
        self.allocation_strategy = value;
        self
    }
}

unsafe impl Allocator for HeapAllocator {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        let size = ::std::cmp::max(minimum_size, self.next_size);
        let mut new_words = Word::allocate_zeroed_vec(size as usize);
        let ptr = new_words.as_mut_ptr();
        self.owned_memory.push(new_words);

        match self.allocation_strategy {
            AllocationStrategy::GrowHeuristically => { self.next_size += size; }
            _ => { }
        }
        (ptr, size as u32)
    }
}

impl Builder<HeapAllocator> {
    pub fn new_default() -> Builder<HeapAllocator> {
        Builder::new(HeapAllocator::new())
    }
}

pub struct ScratchSpace<'a> {
    slice: &'a mut [Word],
    in_use: bool,
}

impl <'a> ScratchSpace<'a> {
    pub fn new(slice: &'a mut [Word]) -> ScratchSpace<'a> {
        ScratchSpace { slice: slice, in_use: false }
    }
}

pub struct ScratchSpaceHeapAllocator<'a, 'b: 'a> {
    scratch_space: &'a mut ScratchSpace<'b>,
    allocator: HeapAllocator,
}

impl <'a, 'b: 'a> ScratchSpaceHeapAllocator<'a, 'b> {
    pub fn new(scratch_space: &'a mut ScratchSpace<'b>) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: scratch_space,
                                    allocator: HeapAllocator::new()}
    }

    pub fn second_segment_words(mut self, value: u32) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.first_segment_words(value) }

    }

    pub fn allocation_strategy(mut self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.allocation_strategy(value) }
    }

}

unsafe impl <'a, 'b: 'a> Allocator for ScratchSpaceHeapAllocator<'a, 'b> {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        if !self.scratch_space.in_use {
            self.scratch_space.in_use = true;
            (self.scratch_space.slice.as_mut_ptr(), self.scratch_space.slice.len() as u32)
        } else {
            self.allocator.allocate_segment(minimum_size)
        }
    }

    fn pre_drop(&mut self, segment0_currently_allocated: u32) {
        let ptr = self.scratch_space.slice.as_mut_ptr();
        unsafe {
            ::std::ptr::write_bytes(ptr, 0u8, segment0_currently_allocated as usize);
        }
        self.scratch_space.in_use = false;
    }
}
<|endoftext|>


<commit_before>use std::str::FromStr;
use std::num::{FromPrimitive, ToPrimitive};
use rustc_serialize::json::{Json, ToJson};

bitflags! {
    #[derive(Show)] flags WSHeader: u16 {
        // Main structure, mask with & to get header parts
        const WS_FIN     = 0b1000000000000000, // final flag
        const WS_RSV     = 0b0111000000000000, // reserved
        const WS_OPCODE  = 0b0000111100000000, // opcode
        const WS_MASK    = 0b0000000010000000, // mask flag
        const WS_LEN     = 0b0000000001111111, // length

        // Opcodes, check for equality after masking with WS_OPCODE
        const WS_OPCONT  = 0b0000000000000000,
        const WS_OPTEXT  = 0b0000000100000000,
        const WS_OPBIN   = 0b0000001000000000,
        const WS_OPTERM  = 0b0000100000000000,
        const WS_OPPING  = 0b0000100100000000,
        const WS_OPPONG  = 0b0000101000000000,

        // Bits reserved for extensions, check for equality after &-ing with OP_RSV
        const WS_RSV1 = 0b0100000000000000,
        const WS_RSV2 = 0b0010000000000000,
        const WS_RSV3 = 0b0001000000000000,

        // Opcodes reserved for extensions, check for equality after &-ing with OP_OPCODE
        const WS_OPEXT1 = 0b0000001100000000,
        const WS_OPEXT2 = 0b0000010000000000,
        const WS_OPEXT3 = 0b0000010100000000,
        const WS_OPEXT4 = 0b0000011000000000,
        const WS_OPEXT5 = 0b0000011100000000,

        const WS_OPCTL1 = 0b0000101100000000,
        const WS_OPCTL2 = 0b0000110000000000,
        const WS_OPCTL3 = 0b0000110100000000,
        const WS_OPCTL4 = 0b0000111000000000,
        const WS_OPCTL5 = 0b0000111100000000,

        // Helper masks
        const WS_OPCTRL  = 0b0000100000000000, // if matches with &, this is a control code
        const WS_LEN16   = 0b0000000001111110, // if &WS_LEN equals this, it has 16-bit length
        const WS_LEN64   = 0b0000000001111111, // if &WS_LEN equals this, it has 32-bit length
    }
}

// TODO: use this instead of u16
#[derive(Copy, Show)]
pub enum WSStatusCode {
    NoError, // = 1000,
    GoneAway, // = 1001,
    ProtocolError, // = 1002,
    UnsupportedData, // = 1003,

    // RESERVERD = 1004,
    NoCode, // = 1005, // reserved
    Aborted, // = 1006, // reserved

    InvalidData, // = 1007,
    ClientError, // = 1008,
    TooLargeData, // = 1009,
    ExtensionMissing, // = 1010,
    ServerError, // = 1011,

    TlsError, // = 1015 // reserved

    // Also:
    // 0-999 - cannot be used,
    // 1000-2999 - reserved for protocol,
    ProtocolCode(u16),
    // 3000-3999 - reserved for apps, issued by IANA,
    ApplicationCode(u16),
    // 4000-4999 - for private use
    OtherCode(u16)
}

impl ToPrimitive for WSStatusCode {
    fn to_u64(&self) -> Option<u64> {
        match *self {
            WSStatusCode::NoError => Some(1000u64),
            WSStatusCode::GoneAway => Some(1001u64),
            WSStatusCode::ProtocolError => Some(1002u64),
            WSStatusCode::UnsupportedData => Some(1003u64),

            WSStatusCode::NoCode => Some(1005u64), // reserved
            WSStatusCode::Aborted => Some(1006u64), // reserved

            WSStatusCode::InvalidData => Some(1007u64),
            WSStatusCode::ClientError => Some(1008u64),
            WSStatusCode::TooLargeData => Some(1009u64),
            WSStatusCode::ExtensionMissing => Some(1010u64),
            WSStatusCode::ServerError => Some(1011u64),

            WSStatusCode::TlsError => Some(1015u64), // reserved

            WSStatusCode::ProtocolCode(code) if 1000 <= code && code <= 2999 => Some(code as u64),
            WSStatusCode::ApplicationCode(code) if 3000 <= code && code <= 3999 => Some(code as u64),
            WSStatusCode::OtherCode(code) if 4000 <= code && code <= 4999 => Some(code as u64),
            _ => None
        }
    }
    fn to_i64(&self) -> Option<i64> {
        self.to_u64().map(|v| v as i64)
    }
}

impl FromPrimitive for WSStatusCode {
    fn from_u64(n: u64) -> Option<WSStatusCode> {
        match n {
            1000u64 => Some(WSStatusCode::NoError),
            1001u64 => Some(WSStatusCode::GoneAway),
            1002u64 => Some(WSStatusCode::ProtocolError),
            1003u64 => Some(WSStatusCode::UnsupportedData),

            1005u64 => Some(WSStatusCode::NoCode), // reserved
            1006u64 => Some(WSStatusCode::Aborted), // reserved

            1007u64 => Some(WSStatusCode::InvalidData),
            1008u64 => Some(WSStatusCode::ClientError),
            1009u64 => Some(WSStatusCode::TooLargeData),
            1010u64 => Some(WSStatusCode::ExtensionMissing),
            1011u64 => Some(WSStatusCode::ServerError),

            1015u64 => Some(WSStatusCode::TlsError), // reserved

            code if 1000 <= code && code <= 2999 => Some(WSStatusCode::ProtocolCode(code as u16)),
            code if 3000 <= code && code <= 3999 => Some(WSStatusCode::ApplicationCode(code as u16)),
            code if 4000 <= code && code <= 4999 => Some(WSStatusCode::OtherCode(code as u16)),
            _ => None
        }
    }
    fn from_i64(n: i64) -> Option<WSStatusCode> {
        if n < 0i64 {
            None
        } else {
            FromPrimitive::from_u64(n as u64)
        }
    }
}

// TODO
// pub struct WSMessage<T=Vec<u8>> {
//     ...
//     data: T
// }
// pub struct WSDataWithStatus<T=Vec<u8>> {
//     pub status: WSStatusCode,
//     pub payload: Vec
// }
#[derive(Show)]
pub struct WSMessage {
    pub header: WSHeader,
    pub data: Vec<u8>,
    pub status: Option<WSStatusCode>
}

impl WSMessage {
    pub fn to_string(&self) -> String {
        String::from_utf8_lossy(self.data[]).into_owned()
    }

    pub fn push(&mut self, msg: WSMessage) {
        self.data.push_all(msg.data[]);
    }

    #[inline] pub fn text(data: &str) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPTEXT,
            data: data.as_bytes().to_vec(),
            status: None
        }
    }

    #[inline] pub fn ext(extn: u8, data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WSHeader::from_bits_truncate((extn as u16) << 8),
            data: data.to_vec(),
            status: None
        }
    }

    #[inline] pub fn binary(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPBIN,
            data: data.to_vec(),
            status: None
        }
    }

    pub fn first(mut self) -> WSMessage {
        self.header.remove(WS_FIN);
        self
    }

    pub fn more(mut self) -> WSMessage {
        self.header.remove(WS_FIN | WS_OPCODE);
        self.header.insert(WS_OPCONT);
        self
    }

    pub fn last(mut self) -> WSMessage {
        self.header.remove(WS_OPCODE);
        self.header.insert(WS_FIN | WS_OPCONT);
        self
    }

    pub fn is_first(&self) -> bool {
        !self.is_final() && !self.is_cont()
    }

    pub fn is_last(&self) -> bool {
        self.is_final() && self.is_cont()
    }

    pub fn is_more(&self) -> bool {
        !self.is_final() && self.is_cont()
    }

    pub fn is_whole(&self) -> bool {
        self.is_final() && !self.is_cont()
    }

    #[inline] pub fn is_final(&self) -> bool {
        self.header.contains(WS_FIN)
    }

    #[inline] pub fn opcode(&self) -> WSHeader {
        self.header & WS_OPCODE
    }

    pub fn mask(mut self) -> WSMessage {
        self.header.insert(WS_MASK);
        self
    }

    pub fn unmask(mut self) -> WSMessage {
        self.header.remove(WS_MASK);
        self
    }

    pub fn is_masked(&self) -> bool {
        self.header.contains(WS_MASK)
    }

    #[inline] pub fn close(status: WSStatusCode, data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPTERM,
            data: data.to_vec(),
            status: Some(status)
        }
    }

    #[inline] pub fn ping(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPPING,
            data: data.to_vec(),
            status: None
        }
    }

    #[inline] pub fn pong(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPPONG,
            data: data.to_vec(),
            status: None
        }
    }

    pub fn is_control(&self) -> bool {
        self.header.contains(WS_OPCTRL)
    }

    pub fn rsv(mut self, n: u8) -> WSMessage {
        self.header.insert(WSHeader::from_bits_truncate(((n & 0x7u8) as u16) << 12));
        self
    }

    pub fn is_rsv(&self, n: u8) -> bool {
        (self.header.bits() >> 12) as u8 & 0x7u8 == n
    }

    #[inline] pub fn is_text(&self) -> bool { self.opcode() == WS_OPTEXT }
    #[inline] pub fn is_binary(&self) -> bool { self.opcode() == WS_OPBIN }
    #[inline] pub fn is_ext(&self, n: u8) -> bool { ((self.opcode().bits() >> 8) as u8) == n }
    #[inline] pub fn is_ping(&self) -> bool { self.opcode() == WS_OPPING }
    #[inline] pub fn is_pong(&self) -> bool { self.opcode() == WS_OPPONG }
    #[inline] pub fn is_close(&self) -> bool { self.opcode() == WS_OPTERM }
    #[inline] pub fn is_cont(&self) -> bool { self.opcode() == WS_OPCONT }

    pub fn split(self, maxlen: uint) -> WSFragmentedMessage {
        WSFragmentedMessage {
            size: self.data.len() + if self.status.is_none() { 0 } else { 2 },
            original: self,
            maxsize: maxlen,
            pos: 0
        }
    }
}

pub struct WSFragmentedMessage {
    original: WSMessage,
    maxsize: uint,
    pos: uint,
    size: uint
}

impl Iterator for WSFragmentedMessage {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        if self.size == 0 {
            None
        } else if self.size <= self.maxsize { // last
            self.size = 0;
            Some(WSMessage {
                header: self.original.header | WS_FIN,
                status: self.original.status,
                data: self.original.data[self.pos..].to_vec()
            })
        } else if self.pos == 0 { // first
            let maxsize = self.maxsize - if self.original.status.is_none() { 0 } else { 2 };
            let result = Some(WSMessage {
                header: self.original.header - WS_FIN,
                status: self.original.status,
                data: self.original.data[..maxsize].to_vec()
            });
            self.original.header.remove(WS_FIN | WS_OPCODE);
            self.original.status = None;
            self.pos = maxsize;
            self.size -= self.maxsize;
            result
        } else { // middle
            let pos = self.pos;
            self.pos += self.maxsize;
            self.size -= self.maxsize;
            Some(WSMessage {
                header: self.original.header,
                status: None,
                data: self.original.data[pos..pos+self.maxsize].to_vec()
            })
        }
    }
}

impl ToJson for WSMessage {
    fn to_json(&self) -> Json {
        self.to_string()[].parse::<Json>().unwrap()
    }
}

impl FromStr for WSMessage {
    #[inline] fn from_str(s: &str) -> Option<WSMessage> {
        Some(WSMessage::text(s))
    }
}
<commit_msg>truncate custom opcode to 4 bits<commit_after>use std::str::FromStr;
use std::num::{FromPrimitive, ToPrimitive};
use rustc_serialize::json::{Json, ToJson};

bitflags! {
    #[derive(Show)] flags WSHeader: u16 {
        // Main structure, mask with & to get header parts
        const WS_FIN     = 0b1000000000000000, // final flag
        const WS_RSV     = 0b0111000000000000, // reserved
        const WS_OPCODE  = 0b0000111100000000, // opcode
        const WS_MASK    = 0b0000000010000000, // mask flag
        const WS_LEN     = 0b0000000001111111, // length

        // Opcodes, check for equality after masking with WS_OPCODE
        const WS_OPCONT  = 0b0000000000000000,
        const WS_OPTEXT  = 0b0000000100000000,
        const WS_OPBIN   = 0b0000001000000000,
        const WS_OPTERM  = 0b0000100000000000,
        const WS_OPPING  = 0b0000100100000000,
        const WS_OPPONG  = 0b0000101000000000,

        // Bits reserved for extensions, check for equality after &-ing with OP_RSV
        const WS_RSV1 = 0b0100000000000000,
        const WS_RSV2 = 0b0010000000000000,
        const WS_RSV3 = 0b0001000000000000,

        // Opcodes reserved for extensions, check for equality after &-ing with OP_OPCODE
        const WS_OPEXT1 = 0b0000001100000000,
        const WS_OPEXT2 = 0b0000010000000000,
        const WS_OPEXT3 = 0b0000010100000000,
        const WS_OPEXT4 = 0b0000011000000000,
        const WS_OPEXT5 = 0b0000011100000000,

        const WS_OPCTL1 = 0b0000101100000000,
        const WS_OPCTL2 = 0b0000110000000000,
        const WS_OPCTL3 = 0b0000110100000000,
        const WS_OPCTL4 = 0b0000111000000000,
        const WS_OPCTL5 = 0b0000111100000000,

        // Helper masks
        const WS_OPCTRL  = 0b0000100000000000, // if matches with &, this is a control code
        const WS_LEN16   = 0b0000000001111110, // if &WS_LEN equals this, it has 16-bit length
        const WS_LEN64   = 0b0000000001111111, // if &WS_LEN equals this, it has 32-bit length
    }
}

// TODO: use this instead of u16
#[derive(Copy, Show)]
pub enum WSStatusCode {
    NoError, // = 1000,
    GoneAway, // = 1001,
    ProtocolError, // = 1002,
    UnsupportedData, // = 1003,

    // RESERVERD = 1004,
    NoCode, // = 1005, // reserved
    Aborted, // = 1006, // reserved

    InvalidData, // = 1007,
    ClientError, // = 1008,
    TooLargeData, // = 1009,
    ExtensionMissing, // = 1010,
    ServerError, // = 1011,

    TlsError, // = 1015 // reserved

    // Also:
    // 0-999 - cannot be used,
    // 1000-2999 - reserved for protocol,
    ProtocolCode(u16),
    // 3000-3999 - reserved for apps, issued by IANA,
    ApplicationCode(u16),
    // 4000-4999 - for private use
    OtherCode(u16)
}

impl ToPrimitive for WSStatusCode {
    fn to_u64(&self) -> Option<u64> {
        match *self {
            WSStatusCode::NoError => Some(1000u64),
            WSStatusCode::GoneAway => Some(1001u64),
            WSStatusCode::ProtocolError => Some(1002u64),
            WSStatusCode::UnsupportedData => Some(1003u64),

            WSStatusCode::NoCode => Some(1005u64), // reserved
            WSStatusCode::Aborted => Some(1006u64), // reserved

            WSStatusCode::InvalidData => Some(1007u64),
            WSStatusCode::ClientError => Some(1008u64),
            WSStatusCode::TooLargeData => Some(1009u64),
            WSStatusCode::ExtensionMissing => Some(1010u64),
            WSStatusCode::ServerError => Some(1011u64),

            WSStatusCode::TlsError => Some(1015u64), // reserved

            WSStatusCode::ProtocolCode(code) if 1000 <= code && code <= 2999 => Some(code as u64),
            WSStatusCode::ApplicationCode(code) if 3000 <= code && code <= 3999 => Some(code as u64),
            WSStatusCode::OtherCode(code) if 4000 <= code && code <= 4999 => Some(code as u64),
            _ => None
        }
    }
    fn to_i64(&self) -> Option<i64> {
        self.to_u64().map(|v| v as i64)
    }
}

impl FromPrimitive for WSStatusCode {
    fn from_u64(n: u64) -> Option<WSStatusCode> {
        match n {
            1000u64 => Some(WSStatusCode::NoError),
            1001u64 => Some(WSStatusCode::GoneAway),
            1002u64 => Some(WSStatusCode::ProtocolError),
            1003u64 => Some(WSStatusCode::UnsupportedData),

            1005u64 => Some(WSStatusCode::NoCode), // reserved
            1006u64 => Some(WSStatusCode::Aborted), // reserved

            1007u64 => Some(WSStatusCode::InvalidData),
            1008u64 => Some(WSStatusCode::ClientError),
            1009u64 => Some(WSStatusCode::TooLargeData),
            1010u64 => Some(WSStatusCode::ExtensionMissing),
            1011u64 => Some(WSStatusCode::ServerError),

            1015u64 => Some(WSStatusCode::TlsError), // reserved

            code if 1000 <= code && code <= 2999 => Some(WSStatusCode::ProtocolCode(code as u16)),
            code if 3000 <= code && code <= 3999 => Some(WSStatusCode::ApplicationCode(code as u16)),
            code if 4000 <= code && code <= 4999 => Some(WSStatusCode::OtherCode(code as u16)),
            _ => None
        }
    }
    fn from_i64(n: i64) -> Option<WSStatusCode> {
        if n < 0i64 {
            None
        } else {
            FromPrimitive::from_u64(n as u64)
        }
    }
}

// TODO
// pub struct WSMessage<T=Vec<u8>> {
//     ...
//     data: T
// }
// pub struct WSDataWithStatus<T=Vec<u8>> {
//     pub status: WSStatusCode,
//     pub payload: Vec
// }
#[derive(Show)]
pub struct WSMessage {
    pub header: WSHeader,
    pub data: Vec<u8>,
    pub status: Option<WSStatusCode>
}

impl WSMessage {
    pub fn to_string(&self) -> String {
        String::from_utf8_lossy(self.data[]).into_owned()
    }

    pub fn push(&mut self, msg: WSMessage) {
        self.data.push_all(msg.data[]);
    }

    #[inline] pub fn text(data: &str) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPTEXT,
            data: data.as_bytes().to_vec(),
            status: None
        }
    }

    #[inline] pub fn ext(extn: u8, data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WSHeader::from_bits_truncate(((extn & 0x0f) as u16) << 8),
            data: data.to_vec(),
            status: None
        }
    }

    #[inline] pub fn binary(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPBIN,
            data: data.to_vec(),
            status: None
        }
    }

    pub fn first(mut self) -> WSMessage {
        self.header.remove(WS_FIN);
        self
    }

    pub fn more(mut self) -> WSMessage {
        self.header.remove(WS_FIN | WS_OPCODE);
        self.header.insert(WS_OPCONT);
        self
    }

    pub fn last(mut self) -> WSMessage {
        self.header.remove(WS_OPCODE);
        self.header.insert(WS_FIN | WS_OPCONT);
        self
    }

    pub fn is_first(&self) -> bool {
        !self.is_final() && !self.is_cont()
    }

    pub fn is_last(&self) -> bool {
        self.is_final() && self.is_cont()
    }

    pub fn is_more(&self) -> bool {
        !self.is_final() && self.is_cont()
    }

    pub fn is_whole(&self) -> bool {
        self.is_final() && !self.is_cont()
    }

    #[inline] pub fn is_final(&self) -> bool {
        self.header.contains(WS_FIN)
    }

    #[inline] pub fn opcode(&self) -> WSHeader {
        self.header & WS_OPCODE
    }

    pub fn mask(mut self) -> WSMessage {
        self.header.insert(WS_MASK);
        self
    }

    pub fn unmask(mut self) -> WSMessage {
        self.header.remove(WS_MASK);
        self
    }

    pub fn is_masked(&self) -> bool {
        self.header.contains(WS_MASK)
    }

    #[inline] pub fn close(status: WSStatusCode, data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPTERM,
            data: data.to_vec(),
            status: Some(status)
        }
    }

    #[inline] pub fn ping(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPPING,
            data: data.to_vec(),
            status: None
        }
    }

    #[inline] pub fn pong(data: &[u8]) -> WSMessage {
        WSMessage {
            header: WS_FIN | WS_OPPONG,
            data: data.to_vec(),
            status: None
        }
    }

    pub fn is_control(&self) -> bool {
        self.header.contains(WS_OPCTRL)
    }

    pub fn rsv(mut self, n: u8) -> WSMessage {
        self.header.insert(WSHeader::from_bits_truncate(((n & 0x7u8) as u16) << 12));
        self
    }

    pub fn is_rsv(&self, n: u8) -> bool {
        (self.header.bits() >> 12) as u8 & 0x7u8 == n
    }

    #[inline] pub fn is_text(&self) -> bool { self.opcode() == WS_OPTEXT }
    #[inline] pub fn is_binary(&self) -> bool { self.opcode() == WS_OPBIN }
    #[inline] pub fn is_ext(&self, n: u8) -> bool { ((self.opcode().bits() >> 8) as u8) == (n & 0x0f) }
    #[inline] pub fn is_ping(&self) -> bool { self.opcode() == WS_OPPING }
    #[inline] pub fn is_pong(&self) -> bool { self.opcode() == WS_OPPONG }
    #[inline] pub fn is_close(&self) -> bool { self.opcode() == WS_OPTERM }
    #[inline] pub fn is_cont(&self) -> bool { self.opcode() == WS_OPCONT }

    pub fn split(self, maxlen: uint) -> WSFragmentedMessage {
        WSFragmentedMessage {
            size: self.data.len() + if self.status.is_none() { 0 } else { 2 },
            original: self,
            maxsize: maxlen,
            pos: 0
        }
    }
}

pub struct WSFragmentedMessage {
    original: WSMessage,
    maxsize: uint,
    pos: uint,
    size: uint
}

impl Iterator for WSFragmentedMessage {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        if self.size == 0 {
            None
        } else if self.size <= self.maxsize { // last
            self.size = 0;
            Some(WSMessage {
                header: self.original.header | WS_FIN,
                status: self.original.status,
                data: self.original.data[self.pos..].to_vec()
            })
        } else if self.pos == 0 { // first
            let maxsize = self.maxsize - if self.original.status.is_none() { 0 } else { 2 };
            let result = Some(WSMessage {
                header: self.original.header - WS_FIN,
                status: self.original.status,
                data: self.original.data[..maxsize].to_vec()
            });
            self.original.header.remove(WS_FIN | WS_OPCODE);
            self.original.status = None;
            self.pos = maxsize;
            self.size -= self.maxsize;
            result
        } else { // middle
            let pos = self.pos;
            self.pos += self.maxsize;
            self.size -= self.maxsize;
            Some(WSMessage {
                header: self.original.header,
                status: None,
                data: self.original.data[pos..pos+self.maxsize].to_vec()
            })
        }
    }
}

impl ToJson for WSMessage {
    fn to_json(&self) -> Json {
        self.to_string()[].parse::<Json>().unwrap()
    }
}

impl FromStr for WSMessage {
    #[inline] fn from_str(s: &str) -> Option<WSMessage> {
        Some(WSMessage::text(s))
    }
}
<|endoftext|>


<commit_before>//! Functions for creating and modifying messages to send across the message bus.
use std::ops::DerefMut;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};
use demarshal::{demarshal,DemarshalError};

#[derive(Debug,Default,PartialEq)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

pub const HEADER_FIELD_INVALID : u8     = 0;
pub const HEADER_FIELD_PATH: u8         = 1;
pub const HEADER_FIELD_INTERFACE: u8    = 2;
pub const HEADER_FIELD_MEMBER: u8       = 3;
pub const HEADER_FIELD_ERROR_NAME: u8   = 4;
pub const HEADER_FIELD_REPLY_SERIAL: u8 = 5;
pub const HEADER_FIELD_DESTINATION: u8  = 6;
pub const HEADER_FIELD_SENDER: u8       = 7;
pub const HEADER_FIELD_SIGNATURE: u8    = 8;

pub const FLAGS_NO_REPLY_EXPECTED : u8  = 1;

#[derive(Debug)]
pub struct HeaderField (
    pub u8,
    pub Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_owned()
    }
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: Vec<HeaderField>,
    pub body: Vec<u8>
}

impl Marshal for Message {
    fn dbus_encode (&self, buf: &mut Vec<u8>) -> usize {
        let endian = if self.big_endian { 'B' as u8 } else { 'l' as u8 };
        endian.dbus_encode(buf);
        self.message_type.0.dbus_encode(buf);
        self.flags.dbus_encode(buf);
        self.version.dbus_encode(buf);
        let len : u32 = self.body.len() as u32;
        len.dbus_encode(buf);
        self.serial.dbus_encode(buf);
        self.headers.dbus_encode(buf);
        pad_to_multiple(buf, 8);
        0
    }

    fn get_type (&self) -> String {
        panic!("Don't do that.")
    }
}

/// Create a Message for a D-Bus method call.  Once a Message is created, arguments
/// can be added with Message.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_CALL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_DESTINATION,
                 Variant::new(Value::from(dest), "s"))
     .add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(iface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(method), "s"))
}

/// Create a Message for a D-Bus method return.  Once created, return values can be added
/// with Message.add_arg
pub fn create_method_return(reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_RETURN,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
}

/// Create a Message for a D-Bus error.  Once created, return values can be added
/// with Message.add_arg
pub fn create_error(error_name: &str, reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_ERROR,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
     .add_header(HEADER_FIELD_ERROR_NAME,
                 Variant::new(Value::from(error_name), "s"))
}

/// Create a Message for a D-Bus signal.  Once created, return values can be added
/// with Message.add_arg
pub fn create_signal(path: &str, interface: &str, member: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_SIGNAL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(interface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(member), "s"))
}

impl Message {
    /// Add the given argument to the Message.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> Message {
        if let None = self.get_header(HEADER_FIELD_SIGNATURE) {
            let value = Value::BasicValue(BasicValue::Signature(Signature("".to_owned())));
            let variant = Variant::new(value, "g");
            self = self.add_header(HEADER_FIELD_SIGNATURE, variant);
        };
        {
            let b : &mut Box<Value> = &mut self.get_header(HEADER_FIELD_SIGNATURE).unwrap().object;
            let val : &mut Value = b.deref_mut();
            match *val {
                Value::BasicValue(BasicValue::Signature(ref mut s)) => s.0.push_str(&arg.get_type()),
                _ => panic!("Garbage in signature field")
            };
        }
        arg.dbus_encode(&mut self.body);
        self
    }

    pub fn get_header(&mut self, name: u8) -> Option<&mut Variant> {
        match self.headers.iter().position(|x| { x.0 == name }) {
            Some(idx) => Some(&mut self.headers[idx].1),
            _ => None
        }
    }

    pub fn add_header(mut self, name: u8, val: Variant) -> Message {
        self.headers.push(HeaderField (name, val));
        self
    }

    /// Get the sequence of Values from out of a Message.  Returns None if the message doesn't have
    /// a body.
    pub fn get_body(&mut self) -> Result<Option<Vec<Value>>,DemarshalError> {
        if self.body.len() == 0 {
            return Ok(None);
        }

        // Get the signature out of the headers
        let v = match self.headers.iter().position(|x| { x.0 == HEADER_FIELD_SIGNATURE }) {
            Some(idx) => &self.headers[idx].1,
            None => return Ok(None)
        };

        let sigval = match *v.object {
            Value::BasicValue(BasicValue::Signature(ref x)) => x,
            _ => return Ok(None)
        };

        let mut sig = "(".to_owned() + &sigval.0 + ")";
        let mut offset = 0;
        match try!(demarshal(&mut self.body, &mut offset, &mut sig)) {
            Value::Struct(x) => Ok(Some(x.objects)),
            x => panic!("Didn't get a struct: {:?}", x)
        }
    }
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<commit_msg>clippy: address len_zero warning<commit_after>//! Functions for creating and modifying messages to send across the message bus.
use std::ops::DerefMut;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};
use demarshal::{demarshal,DemarshalError};

#[derive(Debug,Default,PartialEq)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

pub const HEADER_FIELD_INVALID : u8     = 0;
pub const HEADER_FIELD_PATH: u8         = 1;
pub const HEADER_FIELD_INTERFACE: u8    = 2;
pub const HEADER_FIELD_MEMBER: u8       = 3;
pub const HEADER_FIELD_ERROR_NAME: u8   = 4;
pub const HEADER_FIELD_REPLY_SERIAL: u8 = 5;
pub const HEADER_FIELD_DESTINATION: u8  = 6;
pub const HEADER_FIELD_SENDER: u8       = 7;
pub const HEADER_FIELD_SIGNATURE: u8    = 8;

pub const FLAGS_NO_REPLY_EXPECTED : u8  = 1;

#[derive(Debug)]
pub struct HeaderField (
    pub u8,
    pub Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_owned()
    }
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: Vec<HeaderField>,
    pub body: Vec<u8>
}

impl Marshal for Message {
    fn dbus_encode (&self, buf: &mut Vec<u8>) -> usize {
        let endian = if self.big_endian { 'B' as u8 } else { 'l' as u8 };
        endian.dbus_encode(buf);
        self.message_type.0.dbus_encode(buf);
        self.flags.dbus_encode(buf);
        self.version.dbus_encode(buf);
        let len : u32 = self.body.len() as u32;
        len.dbus_encode(buf);
        self.serial.dbus_encode(buf);
        self.headers.dbus_encode(buf);
        pad_to_multiple(buf, 8);
        0
    }

    fn get_type (&self) -> String {
        panic!("Don't do that.")
    }
}

/// Create a Message for a D-Bus method call.  Once a Message is created, arguments
/// can be added with Message.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_CALL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_DESTINATION,
                 Variant::new(Value::from(dest), "s"))
     .add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(iface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(method), "s"))
}

/// Create a Message for a D-Bus method return.  Once created, return values can be added
/// with Message.add_arg
pub fn create_method_return(reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_RETURN,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
}

/// Create a Message for a D-Bus error.  Once created, return values can be added
/// with Message.add_arg
pub fn create_error(error_name: &str, reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_ERROR,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
     .add_header(HEADER_FIELD_ERROR_NAME,
                 Variant::new(Value::from(error_name), "s"))
}

/// Create a Message for a D-Bus signal.  Once created, return values can be added
/// with Message.add_arg
pub fn create_signal(path: &str, interface: &str, member: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_SIGNAL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),
    }.add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(interface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(member), "s"))
}

impl Message {
    /// Add the given argument to the Message.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> Message {
        if let None = self.get_header(HEADER_FIELD_SIGNATURE) {
            let value = Value::BasicValue(BasicValue::Signature(Signature("".to_owned())));
            let variant = Variant::new(value, "g");
            self = self.add_header(HEADER_FIELD_SIGNATURE, variant);
        };
        {
            let b : &mut Box<Value> = &mut self.get_header(HEADER_FIELD_SIGNATURE).unwrap().object;
            let val : &mut Value = b.deref_mut();
            match *val {
                Value::BasicValue(BasicValue::Signature(ref mut s)) => s.0.push_str(&arg.get_type()),
                _ => panic!("Garbage in signature field")
            };
        }
        arg.dbus_encode(&mut self.body);
        self
    }

    pub fn get_header(&mut self, name: u8) -> Option<&mut Variant> {
        match self.headers.iter().position(|x| { x.0 == name }) {
            Some(idx) => Some(&mut self.headers[idx].1),
            _ => None
        }
    }

    pub fn add_header(mut self, name: u8, val: Variant) -> Message {
        self.headers.push(HeaderField (name, val));
        self
    }

    /// Get the sequence of Values from out of a Message.  Returns None if the message doesn't have
    /// a body.
    pub fn get_body(&mut self) -> Result<Option<Vec<Value>>,DemarshalError> {
        if self.body.is_empty() {
            return Ok(None);
        }

        // Get the signature out of the headers
        let v = match self.headers.iter().position(|x| { x.0 == HEADER_FIELD_SIGNATURE }) {
            Some(idx) => &self.headers[idx].1,
            None => return Ok(None)
        };

        let sigval = match *v.object {
            Value::BasicValue(BasicValue::Signature(ref x)) => x,
            _ => return Ok(None)
        };

        let mut sig = "(".to_owned() + &sigval.0 + ")";
        let mut offset = 0;
        match try!(demarshal(&mut self.body, &mut offset, &mut sig)) {
            Value::Struct(x) => Ok(Some(x.objects)),
            x => panic!("Didn't get a struct: {:?}", x)
        }
    }
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<|endoftext|>