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

use std::fmt;
use std::os::raw::c_void;
use std::ptr;
use std::cmp::{Eq, PartialEq};

use stack::Stack;

// Requires cdecl calling convention on x86, which is the default for "C" blocks.
extern "C" {
    /// Creates a new `Context` ontop of some stack.
    ///
    /// # Arguments
    /// * `sp`   - A pointer to the bottom of the stack.
    /// * `size` - The size of the stack.
    /// * `f`    - A function to be invoked on the first call to jump_fcontext(this, _).
    #[inline(never)]
    fn make_fcontext(sp: *mut c_void, size: usize, f: ContextFn) -> *const c_void;

    /// Yields the execution to another `Context`.
    ///
    /// # Arguments
    /// * `to` - A pointer to the `Context` with whom we swap execution.
    /// * `p`  - An arbitrary argument that will be set as the `data` field
    ///          of the `Transfer` object passed to the other Context.
    #[inline(never)]
    fn jump_fcontext(to: *const c_void, p: usize) -> Transfer;

    /// Yields the execution to another `Context` and executes a function "ontop" of it's stack.
    ///
    /// # Arguments
    /// * `to` - A pointer to the `Context` with whom we swap execution.
    /// * `p`  - An arbitrary argument that will be set as the `data` field
    ///          of the `Transfer` object passed to the other Context.
    /// * `f`  - A function to be invoked on `to` before returning.
    #[inline(never)]
    fn ontop_fcontext(to: *const c_void, p: usize, f: ResumeOntopFn) -> Transfer;
}

/// Functions of this signature are used as the entry point for a new `Context`.
pub type ContextFn = extern "C" fn(t: Transfer) -> !;

/// Functions of this signature are used as the callback after resuming ontop of a `Context`.
pub type ResumeOntopFn = extern "C" fn(t: Transfer) -> Transfer;

/// A `Context` provides the capability of saving and restoring the current state of execution.
///
/// If we have 2 or more `Context` instances, we can thus easily "freeze" the
/// current state of execution and explicitely switch to another `Context`.
/// This `Context` is then resumed exactly where it left of and
/// can in turn "freeze" and switch to another `Context`.
///
/// # Examples
///
/// See [examples/basic.rs](https://github.com/zonyitoo/context-rs/blob/master/examples/basic.rs)
#[repr(C)]
pub struct Context(*const c_void);

// NOTE: Rustc is kinda dumb and introduces a overhead of up to 500% compared to the asm methods
//       if we don't explicitely inline them or use LTO (e.g.: 3ns/iter VS. 18ns/iter on i7 3770).
impl Context {
    /// Returns a null-pointer reference to a `Context` struct.
    ///
    /// This method is used in combination with `Transfer::empty()`.
    #[inline(always)]
    pub unsafe fn null_ref() -> Context {
        Context(ptr::null_mut())
    }

    /// Creates a new `Context` prepared to execute `f` at the beginning of `stack`
    /// and returns a reference to it.
    ///
    /// The passed method `f` is not executed until the first call to `resume()`.
    #[inline(always)]
    pub fn new(stack: &Stack, f: ContextFn) -> Context {
        Context(unsafe { make_fcontext(stack.top(), stack.len(), f) })
    }

    /// Yields the execution to another `Context`.
    ///
    /// The exact behaviour of this method is implementation defined, but the general mechanism is:
    /// The current state of execution is preserved somewhere and the previously saved state
    /// in the `Context` pointed to by `&self` is restored and executed next.
    ///
    /// This behaviour is similiar in spirit to regular function calls with the difference
    /// that the call to `resume()` only returns when someone resumes the caller in turn.
    #[inline(always)]
    pub fn resume(self, data: usize) -> Transfer {
        unsafe { jump_fcontext(self.0, data) }
    }

    /// Yields the execution to another `Context` and executes a function "ontop" of it's stack.
    ///
    /// This method works similiary to `resume()`.
    /// The difference is that the argument `f` is executed right before the targeted `Context`
    /// pointed to by `&self` is woken up and returns from it's call to `resume()`.
    /// The method `f` is now passed the `Transfer` struct which would normally be returned by
    /// `resume()` and is allowed to inspect and modify it. When `f` is done it
    /// has to return a `Transfer` struct which is then finally the one the `resume()`
    /// method returns in the targeted `Context`.
    ///
    /// This behaviour can be used to either execute additional code or map the `Transfer` struct
    /// to another one before it's returned, without the targeted `Context` giving it's consent.
    /// This behaviour can for instance be used to unwind the stack of an unfinished `Context`,
    /// by calling this method with a function that panics, or to deallocate the own stack,
    /// by deferring the actual deallocation until we are switched back to another, safe `Context`.
    #[inline(always)]
    pub fn resume_ontop(self, data: usize, f: ResumeOntopFn) -> Transfer {
        unsafe { ontop_fcontext(self.0, data, f) }
    }
}

impl fmt::Debug for Context {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:p}", self.0)
    }
}

impl PartialEq for Context {
    fn eq(&self, other: &Context) -> bool {
        self.0 == other.0
    }
}

impl Eq for Context {}

/// This is the return value by `Context::resume()` and `Context::resume_ontop()`.
#[repr(C)]
pub struct Transfer {
    /// The previously executed `Context` which yielded to resume the current one.
    pub context: Context,

    /// The `data` which was passed to `Context::resume()` or
    /// `Context::resume_ontop()` to resume the current `Context`.
    pub data: usize,
}

impl Transfer {
    /// Returns a new `Transfer` struct with the members set to their respective arguments.
    #[inline(always)]
    pub fn new(context: Context, data: usize) -> Transfer {
        Transfer {
            context: context,
            data: data,
        }
    }

    /// Returns a `Transfer` struct with the `context` member set to a null-pointer reference.
    ///
    /// This method can be used if there is a need to return a `Transfer`
    /// struct but there is no `Context` left to be returned.
    /// This is for instance the case if you use `resume_ontop()` to destroy
    /// the stack as can be seen in the deallocate_ontop() test.
    #[inline(always)]
    pub unsafe fn empty(data: usize) -> Transfer {
        Transfer {
            context: Context::null_ref(),
            data: data,
        }
    }
}

impl fmt::Debug for Transfer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "Transfer {{ context: {:p}, data: {:p} }}",
               self.context.0,
               self.data as *const c_void)
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use std::os::raw::c_void;

    use stack::ProtectedFixedSizeStack;
    use super::*;

    #[test]
    fn type_sizes() {
        assert_eq!(mem::size_of::<Context>(), mem::size_of::<usize>());
        assert_eq!(mem::size_of::<Context>(), mem::size_of::<*const c_void>());
    }

    #[test]
    fn number_generator() {
        extern "C" fn number_generator(mut t: Transfer) -> ! {
            for i in 0usize.. {
                assert_eq!(t.data, i);
                t = t.context.resume(i);
            }

            unreachable!();
        }

        let stack = ProtectedFixedSizeStack::default();
        let mut t = Transfer::new(Context::new(&stack, number_generator), 0);

        for i in 0..10usize {
            t = t.context.resume(i);
            assert_eq!(t.data, i);

            if t.data == 9 {
                break;
            }
        }
    }

    #[test]
    fn resume_ontop() {
        extern "C" fn resume(t: Transfer) -> ! {
            assert_eq!(t.data, 0);
            t.context.resume_ontop(1, resume_ontop);
            unreachable!();
        }

        extern "C" fn resume_ontop(mut t: Transfer) -> Transfer {
            assert_eq!(t.data, 1);
            t.data = 123;
            t
        }

        let stack = ProtectedFixedSizeStack::default();
        let mut t = Transfer::new(Context::new(&stack, resume), 0);

        t = t.context.resume(0);
        assert_eq!(t.data, 123);
    }
}
<commit_msg>let Context derive Debug, Eq, PartialEq<commit_after>// Copyright 2016 coroutine-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std::fmt;
use std::os::raw::c_void;
use std::ptr;

use stack::Stack;

// Requires cdecl calling convention on x86, which is the default for "C" blocks.
extern "C" {
    /// Creates a new `Context` ontop of some stack.
    ///
    /// # Arguments
    /// * `sp`   - A pointer to the bottom of the stack.
    /// * `size` - The size of the stack.
    /// * `f`    - A function to be invoked on the first call to jump_fcontext(this, _).
    #[inline(never)]
    fn make_fcontext(sp: *mut c_void, size: usize, f: ContextFn) -> *const c_void;

    /// Yields the execution to another `Context`.
    ///
    /// # Arguments
    /// * `to` - A pointer to the `Context` with whom we swap execution.
    /// * `p`  - An arbitrary argument that will be set as the `data` field
    ///          of the `Transfer` object passed to the other Context.
    #[inline(never)]
    fn jump_fcontext(to: *const c_void, p: usize) -> Transfer;

    /// Yields the execution to another `Context` and executes a function "ontop" of it's stack.
    ///
    /// # Arguments
    /// * `to` - A pointer to the `Context` with whom we swap execution.
    /// * `p`  - An arbitrary argument that will be set as the `data` field
    ///          of the `Transfer` object passed to the other Context.
    /// * `f`  - A function to be invoked on `to` before returning.
    #[inline(never)]
    fn ontop_fcontext(to: *const c_void, p: usize, f: ResumeOntopFn) -> Transfer;
}

/// Functions of this signature are used as the entry point for a new `Context`.
pub type ContextFn = extern "C" fn(t: Transfer) -> !;

/// Functions of this signature are used as the callback after resuming ontop of a `Context`.
pub type ResumeOntopFn = extern "C" fn(t: Transfer) -> Transfer;

/// A `Context` provides the capability of saving and restoring the current state of execution.
///
/// If we have 2 or more `Context` instances, we can thus easily "freeze" the
/// current state of execution and explicitely switch to another `Context`.
/// This `Context` is then resumed exactly where it left of and
/// can in turn "freeze" and switch to another `Context`.
///
/// # Examples
///
/// See [examples/basic.rs](https://github.com/zonyitoo/context-rs/blob/master/examples/basic.rs)
#[derive(Eq, PartialEq, Debug)]
#[repr(C)]
pub struct Context(*const c_void);

// NOTE: Rustc is kinda dumb and introduces a overhead of up to 500% compared to the asm methods
//       if we don't explicitely inline them or use LTO (e.g.: 3ns/iter VS. 18ns/iter on i7 3770).
impl Context {
    /// Returns a null-pointer reference to a `Context` struct.
    ///
    /// This method is used in combination with `Transfer::empty()`.
    #[inline(always)]
    pub unsafe fn null_ref() -> Context {
        Context(ptr::null_mut())
    }

    /// Creates a new `Context` prepared to execute `f` at the beginning of `stack`
    /// and returns a reference to it.
    ///
    /// The passed method `f` is not executed until the first call to `resume()`.
    #[inline(always)]
    pub fn new(stack: &Stack, f: ContextFn) -> Context {
        Context(unsafe { make_fcontext(stack.top(), stack.len(), f) })
    }

    /// Yields the execution to another `Context`.
    ///
    /// The exact behaviour of this method is implementation defined, but the general mechanism is:
    /// The current state of execution is preserved somewhere and the previously saved state
    /// in the `Context` pointed to by `&self` is restored and executed next.
    ///
    /// This behaviour is similiar in spirit to regular function calls with the difference
    /// that the call to `resume()` only returns when someone resumes the caller in turn.
    #[inline(always)]
    pub fn resume(self, data: usize) -> Transfer {
        unsafe { jump_fcontext(self.0, data) }
    }

    /// Yields the execution to another `Context` and executes a function "ontop" of it's stack.
    ///
    /// This method works similiary to `resume()`.
    /// The difference is that the argument `f` is executed right before the targeted `Context`
    /// pointed to by `&self` is woken up and returns from it's call to `resume()`.
    /// The method `f` is now passed the `Transfer` struct which would normally be returned by
    /// `resume()` and is allowed to inspect and modify it. When `f` is done it
    /// has to return a `Transfer` struct which is then finally the one the `resume()`
    /// method returns in the targeted `Context`.
    ///
    /// This behaviour can be used to either execute additional code or map the `Transfer` struct
    /// to another one before it's returned, without the targeted `Context` giving it's consent.
    /// This behaviour can for instance be used to unwind the stack of an unfinished `Context`,
    /// by calling this method with a function that panics, or to deallocate the own stack,
    /// by deferring the actual deallocation until we are switched back to another, safe `Context`.
    #[inline(always)]
    pub fn resume_ontop(self, data: usize, f: ResumeOntopFn) -> Transfer {
        unsafe { ontop_fcontext(self.0, data, f) }
    }
}

/// This is the return value by `Context::resume()` and `Context::resume_ontop()`.
#[repr(C)]
pub struct Transfer {
    /// The previously executed `Context` which yielded to resume the current one.
    pub context: Context,

    /// The `data` which was passed to `Context::resume()` or
    /// `Context::resume_ontop()` to resume the current `Context`.
    pub data: usize,
}

impl Transfer {
    /// Returns a new `Transfer` struct with the members set to their respective arguments.
    #[inline(always)]
    pub fn new(context: Context, data: usize) -> Transfer {
        Transfer {
            context: context,
            data: data,
        }
    }

    /// Returns a `Transfer` struct with the `context` member set to a null-pointer reference.
    ///
    /// This method can be used if there is a need to return a `Transfer`
    /// struct but there is no `Context` left to be returned.
    /// This is for instance the case if you use `resume_ontop()` to destroy
    /// the stack as can be seen in the deallocate_ontop() test.
    #[inline(always)]
    pub unsafe fn empty(data: usize) -> Transfer {
        Transfer {
            context: Context::null_ref(),
            data: data,
        }
    }
}

impl fmt::Debug for Transfer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "Transfer {{ context: {:?}, data: {:p} }}",
               self.context,
               self.data as *const c_void)
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use std::os::raw::c_void;

    use stack::ProtectedFixedSizeStack;
    use super::*;

    #[test]
    fn type_sizes() {
        assert_eq!(mem::size_of::<Context>(), mem::size_of::<usize>());
        assert_eq!(mem::size_of::<Context>(), mem::size_of::<*const c_void>());
    }

    #[test]
    fn number_generator() {
        extern "C" fn number_generator(mut t: Transfer) -> ! {
            for i in 0usize.. {
                assert_eq!(t.data, i);
                t = t.context.resume(i);
            }

            unreachable!();
        }

        let stack = ProtectedFixedSizeStack::default();
        let mut t = Transfer::new(Context::new(&stack, number_generator), 0);

        for i in 0..10usize {
            t = t.context.resume(i);
            assert_eq!(t.data, i);

            if t.data == 9 {
                break;
            }
        }
    }

    #[test]
    fn resume_ontop() {
        extern "C" fn resume(t: Transfer) -> ! {
            assert_eq!(t.data, 0);
            t.context.resume_ontop(1, resume_ontop);
            unreachable!();
        }

        extern "C" fn resume_ontop(mut t: Transfer) -> Transfer {
            assert_eq!(t.data, 1);
            t.data = 123;
            t
        }

        let stack = ProtectedFixedSizeStack::default();
        let mut t = Transfer::new(Context::new(&stack, resume), 0);

        t = t.context.resume(0);
        assert_eq!(t.data, 123);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::FuzzTargetImpl;
use language_e2e_tests::account_universe::{
    all_transactions_strategy, log_balance_strategy, run_and_assert_universe, AccountUniverseGen,
};
use libra_proptest_helpers::ValueGenerator;
use libra_types::transaction::SignedTransaction;
use proptest::{collection::vec, prelude::*, test_runner};

#[derive(Clone, Debug, Default)]
pub struct LanguageTransactionExecution;

impl FuzzTargetImpl for LanguageTransactionExecution {
    fn description(&self) -> &'static str {
        "Language execute randomly generated transactions"
    }

    fn fuzz(&self, data: &[u8]) {
        let passthrough_rng =
            test_runner::TestRng::from_seed(test_runner::RngAlgorithm::PassThrough, &data);

        let mut generator = ValueGenerator::new_with_rng(passthrough_rng);
        let txn_strategy = vec(all_transactions_strategy(0, 1_000_000), 1..40);

        let txns = generator.generate(txn_strategy);

        let universe_strategy =
            AccountUniverseGen::strategy(2..20, log_balance_strategy(10_000_000));

        let universe = generator.generate(universe_strategy);

        run_and_assert_universe(universe, txns).unwrap();
    }
}

#[derive(Clone, Debug, Default)]
pub struct SignedTransactionTarget;

impl FuzzTargetImpl for SignedTransactionTarget {
    fn description(&self) -> &'static str {
        "SignedTransaction (LCS deserializer)"
    }

    fn generate(&self, _idx: usize, gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        let value = gen.generate(any_with::<SignedTransaction>(()));
        Some(lcs::to_bytes(&value).expect("serialization should work"))
    }

    fn fuzz(&self, data: &[u8]) {
        let _: Result<SignedTransaction, _> = lcs::from_bytes(&data);
    }
}
<commit_msg>[fuzzing] disallow 0 amount txn<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::FuzzTargetImpl;
use language_e2e_tests::account_universe::{
    all_transactions_strategy, log_balance_strategy, run_and_assert_universe, AccountUniverseGen,
};
use libra_proptest_helpers::ValueGenerator;
use libra_types::transaction::SignedTransaction;
use proptest::{collection::vec, prelude::*, test_runner};

#[derive(Clone, Debug, Default)]
pub struct LanguageTransactionExecution;

impl FuzzTargetImpl for LanguageTransactionExecution {
    fn description(&self) -> &'static str {
        "Language execute randomly generated transactions"
    }

    fn fuzz(&self, data: &[u8]) {
        let passthrough_rng =
            test_runner::TestRng::from_seed(test_runner::RngAlgorithm::PassThrough, &data);

        let mut generator = ValueGenerator::new_with_rng(passthrough_rng);
        let txn_strategy = vec(all_transactions_strategy(1, 1_000_000), 1..40);

        let txns = generator.generate(txn_strategy);

        let universe_strategy =
            AccountUniverseGen::strategy(2..20, log_balance_strategy(10_000_000));

        let universe = generator.generate(universe_strategy);

        run_and_assert_universe(universe, txns).unwrap();
    }
}

#[derive(Clone, Debug, Default)]
pub struct SignedTransactionTarget;

impl FuzzTargetImpl for SignedTransactionTarget {
    fn description(&self) -> &'static str {
        "SignedTransaction (LCS deserializer)"
    }

    fn generate(&self, _idx: usize, gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        let value = gen.generate(any_with::<SignedTransaction>(()));
        Some(lcs::to_bytes(&value).expect("serialization should work"))
    }

    fn fuzz(&self, data: &[u8]) {
        let _: Result<SignedTransaction, _> = lcs::from_bytes(&data);
    }
}
<|endoftext|>


<commit_before>enum Register8 {
    A,
    F,
    B,
    C,
    D,
    E,
    H,
    L
}

enum Register16 {
    AF,
    BC,
    DE,
    HL,
    SP,
    PC
}

struct Cpu {

    /// Accumulator register
    a: u8,
    /// Flags register
    f: u8,
    b: u8,
    c: u8,
    d: u8,
    e: u8,
    h: u8,
    l: u8,
    sp: u16,
    pc: u16

}

macro_rules! high_byte {
    ($x:expr) => { ($x >> 8) as u8 };
}

macro_rules! low_byte {
    ($x:expr) => { $x as u8 };
}

macro_rules! join_bytes {
    ($a:expr, $b:expr) => ({
        let mut join: u16 = $b as u16;
        join |= ($a as u16) << 8;
        join
    })
}

impl Cpu {

    /// Init a new Cpu instance
    fn new() -> Cpu {
        Cpu {
            a: 0x0,
            f: 0x0,
            b: 0x0,
            c: 0x0,
            d: 0x0,
            e: 0x0,
            h: 0x0,
            l: 0x0,
            sp: 0x00,
            pc: 0x00
        }
    }

    /// Get a 8bit value from register
    fn get_8(&self, reg: Register8) -> u8 {
        match reg {
            Register8::A => self.a,
            Register8::F => self.f,
            Register8::B => self.b,
            Register8::C => self.c,
            Register8::D => self.d,
            Register8::E => self.e,
            Register8::H => self.h,
            Register8::L => self.l,
        }
    }

    /// Set a 8bit value to a register
    fn set_8(&mut self, reg: Register8, v: u8) {
        match reg {
            Register8::A => self.a = v,
            Register8::F => self.f = v,
            Register8::B => self.b = v,
            Register8::C => self.c = v,
            Register8::D => self.d = v,
            Register8::E => self.e = v,
            Register8::H => self.h = v,
            Register8::L => self.l = v,
        };
    }

    /// Get a 16bit value from a register
    fn get_16(&self, reg: Register16) -> u16 {
        match reg {
            Register16::AF => { join_bytes!(self.a, self.f) }
            Register16::BC => { join_bytes!(self.b, self.c) }
            Register16::DE => { join_bytes!(self.d, self.e) }
            Register16::HL => { join_bytes!(self.h, self.l) }
            Register16::SP => self.sp,
            Register16::PC => self.pc
        }
    }

    /// Set a 16bit value to a register
    fn set_16(&mut self, reg: Register16, v: u16) {
        match reg {
            Register16::AF => {
                self.a = high_byte!(v);
                self.f = low_byte!(v);
            }
            Register16::BC => {
                self.b = high_byte!(v);
                self.c = low_byte!(v);
            }
            Register16::DE => {
                self.d = high_byte!(v);
                self.e = low_byte!(v);
            }
            Register16::HL => {
                self.h = high_byte!(v);
                self.l = low_byte!(v);
            }
            Register16::SP => self.sp = v,
            Register16::PC => self.pc = v
        }
    }

}

#[test]
fn set_get() {
    let mut cpu = Cpu::new();

    cpu.set_8(Register8::A, 50);
    cpu.set_8(Register8::F, 51);
    cpu.set_8(Register8::B, 52);
    cpu.set_8(Register8::C, 53);
    cpu.set_8(Register8::D, 54);
    cpu.set_8(Register8::E, 55);
    cpu.set_8(Register8::H, 56);
    cpu.set_8(Register8::L, 57);
    assert_eq!(cpu.get_8(Register8::A), 50);
    assert_eq!(cpu.get_8(Register8::F), 51);
    assert_eq!(cpu.get_8(Register8::B), 52);
    assert_eq!(cpu.get_8(Register8::C), 53);
    assert_eq!(cpu.get_8(Register8::D), 54);
    assert_eq!(cpu.get_8(Register8::E), 55);
    assert_eq!(cpu.get_8(Register8::H), 56);
    assert_eq!(cpu.get_8(Register8::L), 57);

    cpu.set_16(Register16::AF, 50_000);
    cpu.set_16(Register16::BC, 51_000);
    cpu.set_16(Register16::DE, 52_000);
    cpu.set_16(Register16::HL, 53_000);
    cpu.set_16(Register16::SP, 54_000);
    cpu.set_16(Register16::PC, 55_000);
    assert_eq!(cpu.get_16(Register16::AF), 50_000);
    assert_eq!(cpu.get_16(Register16::BC), 51_000);
    assert_eq!(cpu.get_16(Register16::DE), 52_000);
    assert_eq!(cpu.get_16(Register16::HL), 53_000);
    assert_eq!(cpu.get_16(Register16::SP), 54_000);
    assert_eq!(cpu.get_16(Register16::PC), 55_000);
}
<commit_msg>Flag implemented<commit_after>enum Register8 {
    A, F,
    B, C,
    D, E,
    H, L
}

enum Register16 {
    AF, BC,
    DE, HL,
    SP, PC
}

enum Flag {
    Z, N,
    H, C
}

struct Cpu {

    /// Accumulator register
    a: u8,
    /// Flags register
    f: u8,
    b: u8,
    c: u8,
    d: u8,
    e: u8,
    h: u8,
    l: u8,
    sp: u16,
    pc: u16

}

macro_rules! high_byte {
    ($x:expr) => { ($x >> 8) as u8 };
}

macro_rules! low_byte {
    ($x:expr) => { $x as u8 };
}

macro_rules! join_bytes {
    ($a:expr, $b:expr) => ({
        let mut join: u16 = $b as u16;
        join |= ($a as u16) << 8;
        join
    })
}

impl Cpu {

    /// Init a new Cpu instance
    fn new() -> Cpu {
        Cpu {
            a: 0x0,
            f: 0x0,
            b: 0x0,
            c: 0x0,
            d: 0x0,
            e: 0x0,
            h: 0x0,
            l: 0x0,
            sp: 0x00,
            pc: 0x00
        }
    }

    /// Get a 8bit value from register
    fn get_8(&self, reg: Register8) -> u8 {
        match reg {
            Register8::A => self.a,
            Register8::F => self.f,
            Register8::B => self.b,
            Register8::C => self.c,
            Register8::D => self.d,
            Register8::E => self.e,
            Register8::H => self.h,
            Register8::L => self.l,
        }
    }

    /// Set a 8bit value to a register
    fn set_8(&mut self, reg: Register8, v: u8) {
        match reg {
            Register8::A => self.a = v,
            Register8::F => self.f = v,
            Register8::B => self.b = v,
            Register8::C => self.c = v,
            Register8::D => self.d = v,
            Register8::E => self.e = v,
            Register8::H => self.h = v,
            Register8::L => self.l = v,
        };
    }

    /// Get a 16bit value from a register
    fn get_16(&self, reg: Register16) -> u16 {
        match reg {
            Register16::AF => { join_bytes!(self.a, self.f) }
            Register16::BC => { join_bytes!(self.b, self.c) }
            Register16::DE => { join_bytes!(self.d, self.e) }
            Register16::HL => { join_bytes!(self.h, self.l) }
            Register16::SP => self.sp,
            Register16::PC => self.pc
        }
    }

    /// Set a 16bit value to a register
    fn set_16(&mut self, reg: Register16, v: u16) {
        match reg {
            Register16::AF => {
                self.a = high_byte!(v);
                self.f = low_byte!(v);
            }
            Register16::BC => {
                self.b = high_byte!(v);
                self.c = low_byte!(v);
            }
            Register16::DE => {
                self.d = high_byte!(v);
                self.e = low_byte!(v);
            }
            Register16::HL => {
                self.h = high_byte!(v);
                self.l = low_byte!(v);
            }
            Register16::SP => self.sp = v,
            Register16::PC => self.pc = v
        }
    }

    /// Set a flags
    fn set_flag(&mut self, flag: Flag, set: bool) {
        let f = self.get_8(Register8::F);
        match (flag, set) {
            (Flag::Z, true) => { self.set_8(Register8::F, (f | 0b10000000)) }
            (Flag::N, true) => { self.set_8(Register8::F, (f | 0b01000000)) }
            (Flag::H, true) => { self.set_8(Register8::F, (f | 0b00100000)) }
            (Flag::C, true) => { self.set_8(Register8::F, (f | 0b00010000)) }
            (Flag::Z, false) => { self.set_8(Register8::F, (f & 0b01111111)) }
            (Flag::N, false) => { self.set_8(Register8::F, (f & 0b10111111)) }
            (Flag::H, false) => { self.set_8(Register8::F, (f & 0b11011111)) }
            (Flag::C, false) => { self.set_8(Register8::F, (f & 0b11101111)) }
        }
    }

}

#[test]
fn set_get() {
    let mut cpu = Cpu::new();

    cpu.set_8(Register8::A, 50);
    cpu.set_8(Register8::F, 51);
    cpu.set_8(Register8::B, 52);
    cpu.set_8(Register8::C, 53);
    cpu.set_8(Register8::D, 54);
    cpu.set_8(Register8::E, 55);
    cpu.set_8(Register8::H, 56);
    cpu.set_8(Register8::L, 57);
    assert_eq!(cpu.get_8(Register8::A), 50);
    assert_eq!(cpu.get_8(Register8::F), 51);
    assert_eq!(cpu.get_8(Register8::B), 52);
    assert_eq!(cpu.get_8(Register8::C), 53);
    assert_eq!(cpu.get_8(Register8::D), 54);
    assert_eq!(cpu.get_8(Register8::E), 55);
    assert_eq!(cpu.get_8(Register8::H), 56);
    assert_eq!(cpu.get_8(Register8::L), 57);

    cpu.set_16(Register16::AF, 50_000);
    cpu.set_16(Register16::BC, 51_000);
    cpu.set_16(Register16::DE, 52_000);
    cpu.set_16(Register16::HL, 53_000);
    cpu.set_16(Register16::SP, 54_000);
    cpu.set_16(Register16::PC, 55_000);
    assert_eq!(cpu.get_16(Register16::AF), 50_000);
    assert_eq!(cpu.get_16(Register16::BC), 51_000);
    assert_eq!(cpu.get_16(Register16::DE), 52_000);
    assert_eq!(cpu.get_16(Register16::HL), 53_000);
    assert_eq!(cpu.get_16(Register16::SP), 54_000);
    assert_eq!(cpu.get_16(Register16::PC), 55_000);
}

#[test]
fn flags() {
    let mut cpu = Cpu::new();

    cpu.set_flag(Flag::Z, true);
    assert_eq!(cpu.get_8(Register8::F), 0b10000000);
    cpu.set_flag(Flag::N, true);
    assert_eq!(cpu.get_8(Register8::F), 0b11000000);
    cpu.set_flag(Flag::H, true);
    assert_eq!(cpu.get_8(Register8::F), 0b11100000);
    cpu.set_flag(Flag::C, true);
    assert_eq!(cpu.get_8(Register8::F), 0b11110000);

    cpu.set_flag(Flag::Z, false);
    assert_eq!(cpu.get_8(Register8::F), 0b01110000);
    cpu.set_flag(Flag::N, false);
    assert_eq!(cpu.get_8(Register8::F), 0b00110000);
    cpu.set_flag(Flag::H, false);
    assert_eq!(cpu.get_8(Register8::F), 0b00010000);
    cpu.set_flag(Flag::C, false);
    assert_eq!(cpu.get_8(Register8::F), 0b00000000);
}
<|endoftext|>


<commit_before>#![allow(experimental)]

use std::io::{File,fs};
use std::path::posix::Path;
use std::os::{homedir};
use std::rand::random;
use std::finally::Finally;

static TEMPLATE: &'static str = include_str!("creator/template.yml");
static DEFAULT_MUXED_DIR: &'static str = "muxed";

pub fn new(name: &str) {
    let muxed_dir = if muxed_dir_exists(&DEFAULT_MUXED_DIR.to_string()) {
      Path::new(format!("{}/.{}/", homedir_string(), &DEFAULT_MUXED_DIR.to_string()))
    } else {
      create_muxed_dir(&DEFAULT_MUXED_DIR.to_string())
    };

    let path = &Path::new(format!("{}/{}", muxed_dir.display(), name));
    if !path.exists() {
      File::create(&Path::new(path));
    } else {
      println!("Project already exists.");
    }
}

fn create_muxed_dir(name: &String) -> Path {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    fs::mkdir(path, ::std::io::UserRWX);
    path.clone()
}

fn muxed_dir_exists(name: &String) -> bool {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    path.exists()
}

fn homedir_string() -> String {
    let home_unwrap = homedir().unwrap();
    format!("{}", home_unwrap.display())
}

#[test]
fn muxed_dir_exists_returns_false() {
  let dir = format!("test_dir_{}", random::<f64>());
  assert!(!muxed_dir_exists(&dir));
}

#[test]
fn muxed_dir_exists_returns_true() {
  let dir = format!("test_dir_{}", random::<f64>());
  create_muxed_dir(&dir);
  assert!(muxed_dir_exists(&dir));

  let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
  fs::rmdir_recursive(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let name        = format!("test_project_{}", random::<f64>());
    let dir         = format!("test_dir_{}", random::<f64>());
    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    fs::rmdir_recursive(muxed_path);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    assert!(path.exists());
    fs::unlink(path);
}

#[test]
#[should_fail]
fn new_doesnt_overwrite_existing_file() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    (|| {
        new(name.as_slice());
    }).finally(|| {
        fs::unlink(path);
    })
}
<commit_msg>Attempt to execute an opening of the users text editor.<commit_after>#![allow(experimental)]

use std::io::{File,fs};
use std::io::process::Command;
use std::path::posix::Path;
use std::os::{homedir};
use std::rand::random;
use std::finally::Finally;

static TEMPLATE: &'static str = include_str!("creator/template.yml");
static DEFAULT_MUXED_DIR: &'static str = "muxed";

pub fn new(name: &str) {
    let muxed_dir = if muxed_dir_exists(&DEFAULT_MUXED_DIR.to_string()) {
      Path::new(format!("{}/.{}/", homedir_string(), &DEFAULT_MUXED_DIR.to_string()))
    } else {
      create_muxed_dir(&DEFAULT_MUXED_DIR.to_string())
    };

    let path = &Path::new(format!("{}/{}", muxed_dir.display(), name));
    if !path.exists() {
      File::create(path);

      let mut process = match Command::new("$EDITOR").arg(format!("{}", path.display())).spawn() {
        Ok(p) => p,
        Err(e) => fail!("failed to execute process: {}", e),
      };
    } else {
      println!("Project already exists.");
    }
}

fn create_muxed_dir(name: &String) -> Path {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    fs::mkdir(path, ::std::io::UserRWX);
    path.clone()
}

fn muxed_dir_exists(name: &String) -> bool {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    path.exists()
}

fn homedir_string() -> String {
    let home_unwrap = homedir().unwrap();
    format!("{}", home_unwrap.display())
}

#[test]
fn muxed_dir_exists_returns_false() {
  let dir = format!("test_dir_{}", random::<f64>());
  assert!(!muxed_dir_exists(&dir));
}

#[test]
fn muxed_dir_exists_returns_true() {
  let dir = format!("test_dir_{}", random::<f64>());
  create_muxed_dir(&dir);
  assert!(muxed_dir_exists(&dir));

  let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
  fs::rmdir_recursive(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let name        = format!("test_project_{}", random::<f64>());
    let dir         = format!("test_dir_{}", random::<f64>());
    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir.as_slice()));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    fs::rmdir_recursive(muxed_path);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    assert!(path.exists());
    fs::unlink(path);
}

#[test]
#[should_fail]
fn new_doesnt_overwrite_existing_file() {
    let name = format!("test_project_{}", random::<f64>());
    let path = &Path::new(format!("{}/.muxed/{}", homedir_string(), name));
    new(name.as_slice());
    (|| {
        new(name.as_slice());
    }).finally(|| {
        fs::unlink(path);
    })
}
<|endoftext|>


<commit_before>use itertools::Itertools;
use ordered_float::OrderedFloat;
use rayon::prelude::*;

use super::CharIndex;
use super::enigma::Enigma;

lazy_static! {
    static ref QGRAMS: Vec<f64> = {
        let f = include_str!("data/quadgrams.txt");
        let mut qgrams = vec![0.0; 456_976];  // 26^4 = 456,976

        for line in f.lines() {
            let line: Vec<_> = line.split(' ').collect();
            let qgram: &str = line[0];
            let count: f64 = line[1].parse().unwrap();

            let index = qgram.chars().fold(0, |acc, c| 26 * acc + c.index());
            qgrams[index] = count.ln();
        }

        qgrams
    };

    static ref ALPHAS: Vec<String> = {
        // Construct the vector ["AAA", "AAB", ..., "ZZY", "ZZZ"]
        iproduct!(b'A'..b'[', b'A'..b'[', b'A'..b'[')
            .map(|(a, b, c)| String::from_utf8(vec![a, b, c]).unwrap())
            .collect()
    };

    static ref ROTORS: Vec<String> = {
        // Create all permutations of rotors 1 through 5
        iproduct!(b'1'..b'6', b'1'..b'6', b'1'..b'6')
            .filter(|&(r1, r2, r3)| (r1 != r2) && (r1 != r3) && (r2 != r3))
            .map(|(r1, r2, r3)| String::from_utf8(vec![r1, r2, r3]).unwrap())
            .collect()
    };
}

/// Decrypts the given string by iterating through multiple possible Enigma
/// configurations, returning the tuple `(plaintext, rotor, key, ring)`
/// corresponding to the most probable plaintext.
pub fn decrypt(msg: &str) -> (String, String, String, String) {
    // Try all rotor/key permutations (60*26^3 == 1,054,560 decryptions).
    let (rotor, key) = iproduct!(ROTORS.iter(), ALPHAS.iter())
        .collect::<Vec<_>>()
        .into_par_iter()
        .max_by_key(|&(rotor, key)| {
            let mut enigma = Enigma::new(rotor, key, "AAA", 'B', "");
            let plaintext = enigma.encrypt(msg);
            OrderedFloat(qgram_score(&plaintext))
        }).unwrap();

    // Use the best rotor/key settings found previously, and iterate through
    // ring settings for the latter two rotors (the ring setting for the slow
    // rotor doesn't matter, so fix it to 'A'), shifting the key setting to
    // match the corresponding ring setting (26^2 == 676 decryptions).
    let (k1, k2, k3) = key.chars().map(|c| c.index()).next_tuple().unwrap();

    let (_, msg, key, ring) = iproduct!(0..26, 0..26).map(|(r2, r3)| {
        let key = &ALPHAS[k1 * 676 + ((k2 + r2) % 26) * 26 + (k3 + r3) % 26];
        let ring = &ALPHAS[r2 * 26 + r3];

        let mut enigma = Enigma::new(rotor, key, ring, 'B', "");
        let plaintext = enigma.encrypt(msg);
        let score = qgram_score(&plaintext);
        (OrderedFloat(score), plaintext, key, ring)
    }).max().unwrap();

    (msg, rotor.clone(), key.clone(), ring.clone())
}

/// Strips all non-alphabetic characters from the given message string and
/// returns the sum of the log-probabilities of each quadgram substring.
pub fn qgram_score(msg: &str) -> f64 {
    let char_indices: Vec<usize> = msg.chars()
        .filter(|&c| c.is_alphabetic())
        .map(|c| c.index())
        .collect();

    if char_indices.len() < 4 {
        panic!("Message must contain 4 or more alphabetic characters.");
    }

    char_indices.windows(4)
        .map(|w| w.iter().fold(0, |acc, x| 26 * acc + x))
        .map(|i| QGRAMS[i])
        .sum()
}

#[cfg(test)]
mod tests {
    use super::{QGRAMS, qgram_score};

    macro_rules! assert_approx_eq {
        ($a:expr, $b:expr) => {
            let (a, b) = (&$a, &$b);
            assert!((*a - *b).abs() < 1.0e-6, "{} is not approximately equal to {}", *a, *b);
        }
    }

    #[test]
    fn qgram_estimates() {
        assert_approx_eq!(QGRAMS[0], 8.81060879);
        assert_approx_eq!(qgram_score("THE QUICK BROWN FOX"), 149.80102862);
    }

    #[test]
    fn sensible_qgram_scores() {
        assert!(qgram_score("AN ENGLISH PHRASE") > qgram_score("ESARHP HSILGNE NA"));
    }

    #[test]
    #[should_panic]
    fn invalid_qgram_check() {
        qgram_score("ABC");
    }
}
<commit_msg>Revert decryption optimization<commit_after>use ordered_float::OrderedFloat;
use rayon::prelude::*;

use super::CharIndex;
use super::enigma::Enigma;

lazy_static! {
    static ref QGRAMS: Vec<f64> = {
        let f = include_str!("data/quadgrams.txt");
        let mut qgrams = vec![0.0; 456_976];  // 26^4 = 456,976

        for line in f.lines() {
            let line: Vec<_> = line.split(' ').collect();
            let qgram: &str = line[0];
            let count: f64 = line[1].parse().unwrap();

            let index = qgram.chars().fold(0, |acc, c| 26 * acc + c.index());
            qgrams[index] = count.ln();
        }

        qgrams
    };

    static ref ALPHAS: Vec<String> = {
        // Construct the vector ["AAA", "AAB", ..., "ZZY", "ZZZ"]
        iproduct!(b'A'..b'[', b'A'..b'[', b'A'..b'[')
            .map(|(a, b, c)| String::from_utf8(vec![a, b, c]).unwrap())
            .collect()
    };

    static ref ROTORS: Vec<String> = {
        // Create all permutations of rotors 1 through 5
        iproduct!(b'1'..b'6', b'1'..b'6', b'1'..b'6')
            .filter(|&(r1, r2, r3)| (r1 != r2) && (r1 != r3) && (r2 != r3))
            .map(|(r1, r2, r3)| String::from_utf8(vec![r1, r2, r3]).unwrap())
            .collect()
    };
}

/// Decrypts the given string by iterating through multiple possible Enigma
/// configurations, returning the tuple `(plaintext, rotor, key, ring)`
/// corresponding to the most probable plaintext.
pub fn decrypt(msg: &str) -> (String, String, String, String) {
    // Try all rotor/key permutations (60*26^3 == 1,054,560 decryptions)
    let (rotor, key) = iproduct!(ROTORS.iter(), ALPHAS.iter())
        .collect::<Vec<_>>()
        .into_par_iter()
        .max_by_key(|&(rotor, key)| {
            let mut enigma = Enigma::new(rotor, key, "AAA", 'B', "");
            let plaintext = enigma.encrypt(msg);
            OrderedFloat(qgram_score(&plaintext))
        }).unwrap();

    // Keep the best rotor configuration found previously, and use the same
    // key setting for the first (slow) rotor. The ring setting for the first
    // rotor doesn't matter, so fix it to 'A', and try all key/ring settings
    // for the remaining two rotors (26^4 == 456,976 decryptions)
    let key_offset = key.chars().nth(0).unwrap().index() * 676;

    let (_, msg, key, ring) = iproduct!(key_offset..(key_offset + 676), 0..676)
        .collect::<Vec<_>>()
        .par_iter()
        .map(|&(key_index, ring_index)| {
            let key = &ALPHAS[key_index];
            let ring = &ALPHAS[ring_index];

            let mut enigma = Enigma::new(rotor, key, ring, 'B', "");
            let plaintext = enigma.encrypt(msg);
            let score = qgram_score(&plaintext);
            (OrderedFloat(score), plaintext, key, ring)
        }).max().unwrap();

    (msg, rotor.clone(), key.clone(), ring.clone())
}

/// Strips all non-alphabetic characters from the given message string and
/// returns the sum of the log-probabilities of each quadgram substring.
pub fn qgram_score(msg: &str) -> f64 {
    let char_indices: Vec<usize> = msg.chars()
        .filter(|&c| c.is_alphabetic())
        .map(|c| c.index())
        .collect();

    if char_indices.len() < 4 {
        panic!("Message must contain 4 or more alphabetic characters.");
    }

    char_indices.windows(4)
        .map(|w| w.iter().fold(0, |acc, x| 26 * acc + x))
        .map(|i| QGRAMS[i])
        .sum()
}

#[cfg(test)]
mod tests {
    use super::{QGRAMS, qgram_score};

    macro_rules! assert_approx_eq {
        ($a:expr, $b:expr) => {
            let (a, b) = (&$a, &$b);
            assert!((*a - *b).abs() < 1.0e-6, "{} is not approximately equal to {}", *a, *b);
        }
    }

    #[test]
    fn qgram_estimates() {
        assert_approx_eq!(QGRAMS[0], 8.81060879);
        assert_approx_eq!(qgram_score("THE QUICK BROWN FOX"), 149.80102862);
    }

    #[test]
    fn sensible_qgram_scores() {
        assert!(qgram_score("AN ENGLISH PHRASE") > qgram_score("ESARHP HSILGNE NA"));
    }

    #[test]
    #[should_panic]
    fn invalid_qgram_check() {
        qgram_score("ABC");
    }
}
<|endoftext|>


<commit_before>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_data::timeinfo::TimeInfo;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::{Demuxer, Event};
use av_format::demuxer::{Descr, Descriptor};
use av_format::demuxer::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational64;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, Err, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

use nom::Needed;

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| d as u64);
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| track_to_stream(tr)).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            },
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => size,
                    _ => 1024,
                };
                Err(ErrorKind::MoreDataNeeded(sz).into())
            },
            e => {
                println!("error reading headers: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }

    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                },
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(ErrorKind::MoreDataNeeded(size).into())
                },
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9" => Some(CodecID::VP9),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

pub fn track_to_stream(t: &TrackEntry) -> Stream {
    Stream {
        id: t.track_uid as usize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::from_integer(1),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                //println!("parsing simple block: {:?}", block);
                let packet = Packet {
                    data: i.into(),
                    t: TimeInfo {
                        pts: Some(block.timecode as i64),
                        dts: None,
                        duration: None,
                        timebase: None,
                    },
                    pos: None,
                    stream_index: block.track_number as isize,
                    is_key: block.keyframe,
                    is_corrupted: false,
                };

                v.push(Event::NewPacket(packet));
            } else {
                println!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe<'a>(&'a self) -> &'a Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    }
};

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::Context;
    use av_format::buffer::*;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => println!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => println!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    println!("Error at size {}", n);
                    panic!();
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(Box::new(MkvDemuxer::new()),
                                       Box::new(AccReader::new(Cursor::new(webm))));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        panic!();
    }
}
<commit_msg>Report the stream timebase<commit_after>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_data::timeinfo::TimeInfo;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::{Demuxer, Event};
use av_format::demuxer::{Descr, Descriptor};
use av_format::demuxer::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational64;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, Err, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

use nom::Needed;

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| d as u64);
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| {
                        track_to_stream(self.info.as_ref().unwrap(), tr)
                    }).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            },
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => size,
                    _ => 1024,
                };
                Err(ErrorKind::MoreDataNeeded(sz).into())
            },
            e => {
                println!("error reading headers: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }

    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                },
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(ErrorKind::MoreDataNeeded(size).into())
                },
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9" => Some(CodecID::VP9),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

// TODO: make sure the timecode_scale isn't 0
pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = if let Some(ts) = t.track_timecode_scale  {
        (ts * info.timecode_scale as f64) as i64
    } else {
        info.timecode_scale as i64
    };

    Stream {
        id: t.track_uid as usize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                //println!("parsing simple block: {:?}", block);
                let packet = Packet {
                    data: i.into(),
                    t: TimeInfo {
                        pts: Some(block.timecode as i64),
                        dts: None,
                        duration: None,
                        timebase: None,
                    },
                    pos: None,
                    stream_index: block.track_number as isize,
                    is_key: block.keyframe,
                    is_corrupted: false,
                };

                v.push(Event::NewPacket(packet));
            } else {
                println!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe<'a>(&'a self) -> &'a Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    }
};

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::Context;
    use av_format::buffer::*;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => println!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => println!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    println!("Error at size {}", n);
                    panic!();
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(Box::new(MkvDemuxer::new()),
                                       Box::new(AccReader::new(Cursor::new(webm))));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        println!("DEMUXER CONTEXT event: {:?}", context.read_event());
        panic!();
    }
}
<|endoftext|>


<commit_before>use av_format::error::*;
use std::io::SeekFrom;
//use av_data::packet::Packet;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational32;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    //pub clusters: Vec<Cluster>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            //clusters: Vec::new(),
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return IResult::Done(input, ());
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            IResult::Done(i, _) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| {
                    d as u64
                });
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| track_to_stream(tr)).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            IResult::Incomplete(_) => Err(ErrorKind::MoreDataNeeded.into()),
            e => {
              println!("error reading headers: {:?}", e);
              Err(ErrorKind::InvalidData.into())
            },
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        match segment_element(buf.data()) {
            IResult::Done(i, element) => {
                let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                match element {
                    SegmentElement::Cluster(c) => {
                        //self.clusters.push(c);
                        println!("got cluster element at timecode: {}", c.timecode);
                        for block in c.simple_block.iter() {
                          println!("got simple block of size {}", block.len());
                          if let IResult::Done(_,o) = simple_block(block) {
                            println!("parsing simple block: {:?}", o);
                          } else {
                            println!("error parsing simple block");
                          }
                          self.blockstream.extend(*block);
                        }

                        for block_group in c.block_group.iter() {
                          println!("got block group of size {}", block_group.block.len());
                          self.blockstream.extend(block_group.block);
                        }
                    }
                    el => {
                        println!("got element: {:#?}", el);
                    }
                }

                Ok((seek, Event::MoreDataNeeded))
            }
            IResult::Incomplete(_) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
            e => {
                println!("parsing issue: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9"  => Some(CodecID::VP9),
        _ => None
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width : video.pixel_width as usize,
            height : video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format : None
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None
    }
}

pub fn track_to_stream(t: &TrackEntry) -> Stream {
    Stream {
    id: t.track_uid as usize,
    index: t.track_number as usize,
    start: None,
    duration: t.default_duration,
    timebase : Rational32::from_integer(1),
    // TODO: Extend CodecParams and fill it with the remaining information
    params : CodecParams {
        extradata: t.codec_private.clone(),
        bit_rate: 0,
        codec_id: track_entry_codec_id(t),
        kind: track_entry_media_kind(t)
    },
  }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::{IResult, Offset};
    use av_format::demuxer::context::DemuxerContext;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            IResult::Done(i, _) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(webm)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<commit_msg>generate packets from a Cluster<commit_after>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational32;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    //pub clusters: Vec<Cluster>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            //clusters: Vec::new(),
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return IResult::Done(input, ());
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return IResult::Error(nom::ErrorKind::Custom(1));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            IResult::Done(i, _) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| {
                    d as u64
                });
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| track_to_stream(tr)).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            IResult::Incomplete(_) => Err(ErrorKind::MoreDataNeeded.into()),
            e => {
              println!("error reading headers: {:?}", e);
              Err(ErrorKind::InvalidData.into())
            },
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        match segment_element(buf.data()) {
            IResult::Done(i, element) => {
                let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                match element {
                    SegmentElement::Cluster(c) => {
                        //self.clusters.push(c);
                        println!("got cluster element at timecode: {}", c.timecode);
                        for block in c.simple_block.iter() {
                          println!("got simple block of size {}", block.len());
                          if let IResult::Done(_,o) = simple_block(block) {
                            println!("parsing simple block: {:?}", o);
                          } else {
                            println!("error parsing simple block");
                          }
                          self.blockstream.extend(*block);
                        }

                        for block_group in c.block_group.iter() {
                          println!("got block group of size {}", block_group.block.len());
                          self.blockstream.extend(block_group.block);
                        }
                    }
                    el => {
                        println!("got element: {:#?}", el);
                    }
                }

                Ok((seek, Event::MoreDataNeeded))
            }
            IResult::Incomplete(_) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
            e => {
                println!("parsing issue: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9"  => Some(CodecID::VP9),
        _ => None
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width : video.pixel_width as usize,
            height : video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format : None
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None
    }
}

pub fn track_to_stream(t: &TrackEntry) -> Stream {
    Stream {
    id: t.track_uid as usize,
    index: t.track_number as usize,
    start: None,
    duration: t.default_duration,
    timebase : Rational32::from_integer(1),
    // TODO: Extend CodecParams and fill it with the remaining information
    params : CodecParams {
        extradata: t.codec_private.clone(),
        bit_rate: 0,
        codec_id: track_entry_codec_id(t),
        kind: track_entry_media_kind(t)
    },
  }
}

impl<'a> Cluster<'a> {
  pub fn generate_packet(&mut self) -> Option<Event> {
    if self.simple_block.len() > 0 {
      let block_data = self.simple_block.remove(0);

      if let IResult::Done(i, block) = simple_block(block_data) {
        println!("parsing simple block: {:?}", block);
        let packet = Packet {
          data:         i.into(),
          pts:          Some(block.timecode as i64),
          dts:          None,
          pos:          None,
          stream_index: block.track_number as isize,
          is_key:       block.keyframe,
          is_corrupted: false,
        };

        Some(Event::NewPacket(packet))
      } else {
        println!("error parsing simple block");
        None
      }
    } else {
      None
    }
  }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::{IResult, Offset};
    use av_format::demuxer::context::DemuxerContext;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            IResult::Done(i, _) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(webm)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<|endoftext|>


<commit_before>// Copyright(C) 2016 Chris Liebert

extern crate glium;
extern crate nalgebra;
extern crate quick3d;

use glium::glutin;
use glium::glutin::{ElementState,Event,MouseButton,VirtualKeyCode};
use glium::DisplayBuild;

use nalgebra::Matrix4;
use std::io::Error;

use quick3d::common::{Camera, Mesh, Scene};
use quick3d::loader::DBLoader;
use quick3d::renderer;

fn main() {
    let screen_width = 1200;
    let screen_height = 800;
    let db_file: &str = "test.db";

    let dbloader: DBLoader = DBLoader::new(db_file);
    let shader_dbloader: DBLoader = DBLoader::new("shaders.db");
    let scene: Scene = dbloader.load_scene();

    let display = glutin::WindowBuilder::new()
        //.resizable()
        //.with_vsync()
        .with_depth_buffer(24)
        .with_title("Rust Window")
        .with_gl_debug_flag(true)
        .with_visibility(false) // Window is shown when scene finishes loading.
        .with_dimensions(screen_width, screen_height)
        .build_glium()
        .unwrap();

    let camera: Camera = Camera::new(screen_width as f32, screen_height as f32);
    let renderer = renderer::Renderer::new(&display, scene);

    let shader_program = renderer.create_shader_program("default", &shader_dbloader, &display);

    // Show the window once the data is loaded
    let window = match display.get_window() {
        Some(x) =>  {
        	x.show();
        	x
        }
        None => {
            panic!("Error retrieving window");
        }
    };

    // The torus will be movable in the scene if it is found
    let torus: Result<&Mesh, Error> = renderer.get_mesh("Torus");

    let mut torus_x = 0.0f32;
    let mut torus_y = 0.0f32;
    let mut torus_vertical_speed = 0.0f32;
    let mut torus_horizontal_speed = 0.0f32;
    let mut left_button_pressed = false;

    let mut mouse_last_x: i32 = 0;
    let mut mouse_last_y: i32 = 0;

    let mut _mouse_dx: i32 = 0;
    let mut _mouse_dy: i32 = 0;

    let mut _camera_forward_speed = 0.0f32;
    let mut _camera_left_speed = 0.0f32;

    let mut _w_state = false;
    let mut _camera_forward_speed = 0.0;

    let screen_center_x: i32 = (screen_width / 2) as i32;
	let screen_center_y: i32 = (screen_height / 2) as i32;

    'running: loop {
    	camera.update();
        renderer.render(&display, &shader_program, &camera);

        // Check for events
        for event in display.poll_events() {
            match event {
                Event::Closed => break 'running,
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::Escape)) => {
                    break 'running;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::W)) => {
			        _camera_forward_speed = 1.0;
		        }
		        Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::W)) => {
		            _camera_forward_speed = 0.0;
		        }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::A)) => {
                    camera.move_left(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::S)) => {
                    camera.move_backward(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::D)) => {
                    camera.move_right(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::MouseInput(ElementState::Pressed, MouseButton::Left) => {
                    left_button_pressed = true;
                }
                Event::MouseInput(ElementState::Released, MouseButton::Left) => {
                    left_button_pressed = false;
                }
                Event::MouseMoved(x, y) => {
                	_mouse_dx = mouse_last_x - x;
                	_mouse_dy = mouse_last_y - y;
                	if left_button_pressed {
                		// Rotate the camera if the left mouse button is pressed
                		camera.aim(_mouse_dx as f64, _mouse_dy as f64);

                		if x + 10 >= screen_width as i32 || x <= 10 {
		                	let _ = window.set_cursor_position(screen_center_x, y);
                		} else if y + 10 >= screen_height as i32 || y <= 10 {
		                	let _ = window.set_cursor_position(x, screen_center_y);
                		}
                	}
                	mouse_last_x = x;
                	mouse_last_y = y;
                }

	            _ => (),
            }
        }

		camera.move_forward(_camera_forward_speed * 0.01);

		// Move the torus (if found) based on changes from keyboard input
        match torus {
        	Ok(torus) => {
		        // Get existing matrix
        		let mut matrix: Matrix4<f32> = *torus.matrix.borrow();
		        torus_x += torus_horizontal_speed;
		        torus_y += torus_vertical_speed;
		        matrix.m14 = torus_x;
		        matrix.m24 = torus_y;
		        // Mutate the matrix
		        *torus.matrix.borrow_mut() = matrix;
        	},
        	Err(_) => {}
        }
    }
}
<commit_msg>updated rust example to try to use a specific shader version<commit_after>// Copyright(C) 2016 Chris Liebert

extern crate glium;
extern crate nalgebra;
extern crate quick3d;

use glium::glutin;
use glium::glutin::{ElementState,Event,MouseButton,VirtualKeyCode};
use glium::DisplayBuild;

use nalgebra::Matrix4;
use std::io::Error;

use quick3d::common::{Camera, Mesh, Scene};
use quick3d::loader::DBLoader;
use quick3d::renderer;

fn main() {
    let screen_width = 1200;
    let screen_height = 800;
    let db_file: &str = "test.db";

    let dbloader: DBLoader = DBLoader::new(db_file);
    let shader_dbloader: DBLoader = DBLoader::new("shaders.db");
    let scene: Scene = dbloader.load_scene();

    let display = glutin::WindowBuilder::new()
        //.resizable()
        //.with_vsync()
        .with_depth_buffer(24)
        .with_title("Rust Window")
        .with_gl_debug_flag(true)
        .with_visibility(false) // Window is shown when scene finishes loading.
        .with_dimensions(screen_width, screen_height)
        .build_glium()
        .unwrap();

    let camera: Camera = Camera::new(screen_width as f32, screen_height as f32);
    let renderer = renderer::Renderer::new(&display, scene);

	// Attempt to load GLSL version 330 if it is supported
	let desired_glsl_version = glium::Version(glium::Api::Gl, 3, 3);
	let shader_name = "default";
    let shader_program = match renderer.create_shader_program_with_version(&shader_name, &shader_dbloader, &desired_glsl_version, &display) {
    	Ok(p) => p,
    	Err(e) =>  {
    		println!("Unable to load {:?}: {}", desired_glsl_version, e);
    		renderer.create_shader_program(&shader_name, &shader_dbloader, &display)
    	}
    };

    // Show the window once the data is loaded
    let window = match display.get_window() {
        Some(x) =>  {
        	x.show();
        	x
        }
        None => {
            panic!("Error retrieving window");
        }
    };

    // The torus will be movable in the scene if it is found
    let torus: Result<&Mesh, Error> = renderer.get_mesh("Torus");

    let mut torus_x = 0.0f32;
    let mut torus_y = 0.0f32;
    let mut torus_vertical_speed = 0.0f32;
    let mut torus_horizontal_speed = 0.0f32;
    let mut left_button_pressed = false;

    let mut mouse_last_x: i32 = 0;
    let mut mouse_last_y: i32 = 0;

    let mut _mouse_dx: i32 = 0;
    let mut _mouse_dy: i32 = 0;

    let mut _camera_forward_speed = 0.0f32;
    let mut _camera_left_speed = 0.0f32;

    let mut _w_state = false;
    let mut _camera_forward_speed = 0.0;

    let screen_center_x: i32 = (screen_width / 2) as i32;
	let screen_center_y: i32 = (screen_height / 2) as i32;

    'running: loop {
    	camera.update();
        renderer.render(&display, &shader_program, &camera);

        // Check for events
        for event in display.poll_events() {
            match event {
                Event::Closed => break 'running,
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::Escape)) => {
                    break 'running;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::W)) => {
			        _camera_forward_speed = 1.0;
		        }
		        Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::W)) => {
		            _camera_forward_speed = 0.0;
		        }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::A)) => {
                    camera.move_left(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::S)) => {
                    camera.move_backward(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::D)) => {
                    camera.move_right(1.0);
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::I)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::J)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = -0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::K)) => {
                    torus_vertical_speed = 0.0f32;
                }
                Event::KeyboardInput(ElementState::Pressed, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.001f32;
                }
                Event::KeyboardInput(ElementState::Released, _, Some(VirtualKeyCode::L)) => {
                    torus_horizontal_speed = 0.0f32;
                }
                Event::MouseInput(ElementState::Pressed, MouseButton::Left) => {
                    left_button_pressed = true;
                }
                Event::MouseInput(ElementState::Released, MouseButton::Left) => {
                    left_button_pressed = false;
                }
                Event::MouseMoved(x, y) => {
                	_mouse_dx = mouse_last_x - x;
                	_mouse_dy = mouse_last_y - y;
                	if left_button_pressed {
                		// Rotate the camera if the left mouse button is pressed
                		camera.aim(_mouse_dx as f64, _mouse_dy as f64);

                		if x + 10 >= screen_width as i32 || x <= 10 {
		                	let _ = window.set_cursor_position(screen_center_x, y);
                		} else if y + 10 >= screen_height as i32 || y <= 10 {
		                	let _ = window.set_cursor_position(x, screen_center_y);
                		}
                	}
                	mouse_last_x = x;
                	mouse_last_y = y;
                }

	            _ => (),
            }
        }

		camera.move_forward(_camera_forward_speed * 0.01);

		// Move the torus (if found) based on changes from keyboard input
        match torus {
        	Ok(torus) => {
		        // Get existing matrix
        		let mut matrix: Matrix4<f32> = *torus.matrix.borrow();
		        torus_x += torus_horizontal_speed;
		        torus_y += torus_vertical_speed;
		        matrix.m14 = torus_x;
		        matrix.m24 = torus_y;
		        // Mutate the matrix
		        *torus.matrix.borrow_mut() = matrix;
        	},
        	Err(_) => {}
        }
    }
}
<|endoftext|>