<commit_before>extern crate semver;

use std::collections::HashMap;
use std::io::net::ip::IpAddr;

#[deriving(PartialEq, Show, Clone)]
pub enum Scheme {
    Http,
    Https
}

#[deriving(PartialEq, Show, Clone)]
pub enum Host<'a> {
    HostName(&'a str),
    HostIp(IpAddr)
}

#[deriving(PartialEq, Show, Clone)]
pub enum Method<'a> {
    Get,
    Post,
    Put,
    Patch,
    Delete,
    Other(&'a str)
}

pub trait Request {
    /// The version of HTTP being used
    fn http_version(&self) -> semver::Version;

    /// The version of the conduit spec being used
    fn conduit_version(&self) -> semver::Version;

    /// The request method, such as GET, POST, PUT, DELETE or PATCH
    fn method<'a>(&'a self) -> Method<'a>;

    /// The scheme part of the request URL
    fn scheme(&self) -> Scheme;

    /// The host part of the requested URL
    fn host<'a>(&'a self) -> Host<'a>;

    /// The initial part of the request URL's path that corresponds
    /// to a virtual root. This allows an application to have a
    /// virtual location that consumes part of the path.
    fn virtual_root<'a>(&'a self) -> Option<&'a str>;

    /// The remainder of the path.
    fn path<'a>(&'a self) -> &'a str;

    /// The portion of the request URL that follows the "?"
    fn query_string<'a>(&'a self) -> Option<&'a str>;

    /// The remote IP address of the client or the last proxy that
    /// sent the request.
    fn remote_ip(&self) -> IpAddr;

    /// The byte-size of the body, if any
    fn content_length(&self) -> Option<uint>;

    fn headers<'a>(&'a self) -> &'a mut Headers;

    /// A Reader for the body of the request
    fn body<'a>(&'a mut self) -> &'a mut Reader;
}

pub trait Headers {
    /// Find the value of a given header. Multi-line headers are represented
    /// as an array.
    fn find<'a>(&'a str) -> Option<&'a [&'a str]>;

    /// Returns true if a particular header exists
    fn has(&str) -> bool;

    /// Iterate over all of the available headers.
    fn iter<'a>(&'a self) -> &'a Iterator<(&'a str, &'a [&'a str])>;
}

pub struct Response {
    /// The status code as a tuple of the return code and status string
    pub status: (uint, &'static str),

    /// A Map of the headers
    pub headers: HashMap<String, String>,

    /// A Writer for body of the response
    pub body: Box<Reader + Send>
}

pub type Handler = fn(&mut Request) -> Response;
<commit_msg>typedef HeaderEntries<commit_after>extern crate semver;

use std::collections::HashMap;
use std::io::net::ip::IpAddr;

#[deriving(PartialEq, Show, Clone)]
pub enum Scheme {
    Http,
    Https
}

#[deriving(PartialEq, Show, Clone)]
pub enum Host<'a> {
    HostName(&'a str),
    HostIp(IpAddr)
}

#[deriving(PartialEq, Show, Clone)]
pub enum Method<'a> {
    Get,
    Post,
    Put,
    Patch,
    Delete,
    Other(&'a str)
}

pub trait Request {
    /// The version of HTTP being used
    fn http_version(&self) -> semver::Version;

    /// The version of the conduit spec being used
    fn conduit_version(&self) -> semver::Version;

    /// The request method, such as GET, POST, PUT, DELETE or PATCH
    fn method<'a>(&'a self) -> Method<'a>;

    /// The scheme part of the request URL
    fn scheme(&self) -> Scheme;

    /// The host part of the requested URL
    fn host<'a>(&'a self) -> Host<'a>;

    /// The initial part of the request URL's path that corresponds
    /// to a virtual root. This allows an application to have a
    /// virtual location that consumes part of the path.
    fn virtual_root<'a>(&'a self) -> Option<&'a str>;

    /// The remainder of the path.
    fn path<'a>(&'a self) -> &'a str;

    /// The portion of the request URL that follows the "?"
    fn query_string<'a>(&'a self) -> Option<&'a str>;

    /// The remote IP address of the client or the last proxy that
    /// sent the request.
    fn remote_ip(&self) -> IpAddr;

    /// The byte-size of the body, if any
    fn content_length(&self) -> Option<uint>;

    fn headers<'a>(&'a self) -> &'a mut Headers;

    /// A Reader for the body of the request
    fn body<'a>(&'a mut self) -> &'a mut Reader;
}

pub type HeaderEntries<'a> = &'a Iterator<(&'a str, &'a [&'a str])>;

pub trait Headers {
    /// Find the value of a given header. Multi-line headers are represented
    /// as an array.
    fn find<'a>(&'a str) -> Option<&'a [&'a str]>;

    /// Returns true if a particular header exists
    fn has(&str) -> bool;

    /// Iterate over all of the available headers.
    fn iter<'a>(&'a self) -> HeaderEntries<'a>;
}

pub struct Response {
    /// The status code as a tuple of the return code and status string
    pub status: (uint, &'static str),

    /// A Map of the headers
    pub headers: HashMap<String, String>,

    /// A Writer for body of the response
    pub body: Box<Reader + Send>
}

pub type Handler = fn(&mut Request) -> Response;
<|endoftext|>


<commit_before>use super::*;

#[derive(Debug, Clone)]
#[doc(hidden)]
pub struct Context {
    // TODO file from config should be in here
    config: RunningConfig,
    /// Periodically flushes dirty data. We keep this in an
    /// Arc separate from the PageCache below to separate
    /// "high-level" references from Db, Tree etc... from
    /// "low-level" references like background threads.
    /// When the last high-level reference is dropped, it
    /// should trigger all background threads to clean
    /// up synchronously.
    #[cfg(all(
        not(miri),
        any(
            windows,
            target_os = "linux",
            target_os = "macos",
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "openbsd",
            target_os = "netbsd",
        )
    ))]
    pub(crate) flusher: Arc<Mutex<Option<flusher::Flusher>>>,
    #[doc(hidden)]
    pub pagecache: Arc<PageCache>,
}

impl std::ops::Deref for Context {
    type Target = RunningConfig;

    fn deref(&self) -> &RunningConfig {
        &self.config
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        #[cfg(all(
            not(miri),
            any(
                windows,
                target_os = "linux",
                target_os = "macos",
                target_os = "dragonfly",
                target_os = "freebsd",
                target_os = "openbsd",
                target_os = "netbsd",
            )
        ))]
        {
            if let Some(flusher) = self.flusher.lock().take() {
                drop(flusher)
            }
        }

        loop {
            match self.pagecache.flush() {
                Ok(0) => return,
                Ok(_) => continue,
                Err(e) => {
                    error!(
                        "failed to flush underlying \
                         pagecache during drop: {:?}",
                        e
                    );
                    return;
                }
            }
        }
    }
}

impl Context {
    pub(crate) fn start(config: RunningConfig) -> Result<Self> {
        trace!("starting context");

        let pagecache = Arc::new(PageCache::start(config.clone())?);

        Ok(Self {
            config,
            pagecache,
            #[cfg(all(
                not(miri),
                any(
                    windows,
                    target_os = "linux",
                    target_os = "macos",
                    target_os = "dragonfly",
                    target_os = "freebsd",
                    target_os = "openbsd",
                    target_os = "netbsd",
                )
            ))]
            flusher: Arc::new(parking_lot::Mutex::new(None)),
        })
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.pagecache.was_recovered()
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<u64> {
        let _cc = concurrency_control::read();
        self.pagecache.generate_id_inner()
    }

    pub(crate) fn pin_log(&self, guard: &Guard) -> Result<RecoveryGuard<'_>> {
        self.pagecache.pin_log(guard)
    }
}
<commit_msg>Get rid of buggy Context::drop implementation.<commit_after>use super::*;

#[derive(Debug, Clone)]
#[doc(hidden)]
pub struct Context {
    // TODO file from config should be in here
    config: RunningConfig,
    /// Periodically flushes dirty data. We keep this in an
    /// Arc separate from the PageCache below to separate
    /// "high-level" references from Db, Tree etc... from
    /// "low-level" references like background threads.
    /// When the last high-level reference is dropped, it
    /// should trigger all background threads to clean
    /// up synchronously.
    #[cfg(all(
        not(miri),
        any(
            windows,
            target_os = "linux",
            target_os = "macos",
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "openbsd",
            target_os = "netbsd",
        )
    ))]
    pub(crate) flusher: Arc<Mutex<Option<flusher::Flusher>>>,
    #[doc(hidden)]
    pub pagecache: Arc<PageCache>,
}

impl std::ops::Deref for Context {
    type Target = RunningConfig;

    fn deref(&self) -> &RunningConfig {
        &self.config
    }
}

impl Context {
    pub(crate) fn start(config: RunningConfig) -> Result<Self> {
        trace!("starting context");

        let pagecache = Arc::new(PageCache::start(config.clone())?);

        Ok(Self {
            config,
            pagecache,
            #[cfg(all(
                not(miri),
                any(
                    windows,
                    target_os = "linux",
                    target_os = "macos",
                    target_os = "dragonfly",
                    target_os = "freebsd",
                    target_os = "openbsd",
                    target_os = "netbsd",
                )
            ))]
            flusher: Arc::new(parking_lot::Mutex::new(None)),
        })
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.pagecache.was_recovered()
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<u64> {
        let _cc = concurrency_control::read();
        self.pagecache.generate_id_inner()
    }

    pub(crate) fn pin_log(&self, guard: &Guard) -> Result<RecoveryGuard<'_>> {
        self.pagecache.pin_log(guard)
    }
}
<|endoftext|>


<commit_before>use mio::Sender;

use handler::Notify;

pub trait AsyncContext<M:Send> {
    fn async_channel<'x>(&'x self) -> &'x Sender<Notify<M>>;
    fn async_add_machine(&mut self, m: M) -> Result<(), M> {
        use mio::NotifyError::*;
        match self.async_channel().send(Notify::NewMachine(m)) {
            Ok(()) => Ok(()),
            Err(Io(e)) => {
                // We would probably do something better here, but mio doesn't
                // give us a message. But anyway it's probably never happen
                panic!("Io error when sending notify");
            }
            Err(Full(Notify::NewMachine(m))) => Err(m),
            Err(Closed(_)) => {
                // It should never happen because we usually send from the
                // inside of a main loop
                panic!("Sending to closed channel. Main loop is already shut \
                    down");
            }
        }
    }
}
<commit_msg>Better error message<commit_after>use mio::Sender;

use handler::Notify;

pub trait AsyncContext<M:Send> {
    fn async_channel<'x>(&'x self) -> &'x Sender<Notify<M>>;
    fn async_add_machine(&mut self, m: M) -> Result<(), M> {
        use mio::NotifyError::*;
        match self.async_channel().send(Notify::NewMachine(m)) {
            Ok(()) => Ok(()),
            Err(Io(e)) => {
                // We would probably do something better here, but mio doesn't
                // give us a message. But anyway it's probably never happen
                panic!("Io error when sending notify: {}", e);
            }
            Err(Full(Notify::NewMachine(m))) => Err(m),
            Err(Closed(_)) => {
                // It should never happen because we usually send from the
                // inside of a main loop
                panic!("Sending to closed channel. Main loop is already shut \
                    down");
            }
        }
    }
}
<|endoftext|>


<commit_before>#![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() {
      create_project_file(path)
    } else {
      println!("Project already exists.");
    }
}

fn create_project_file(path: &Path) {
    match File::create(path).write(TEMPLATE.as_bytes()) {
        Ok(()) => (), // succeeded
        Err(e) => println!("Failed to create project {}", path.filename()),
    }
//    Command::new("vim").arg(format!("{}", path.display())).detached();
}

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]
// TODO: Fix this test so it verifies something better.
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>Make these crates as test only.<commit_after>#![allow(experimental)]

use std::io::{File,fs};
//use std::io::process::Command;
use std::path::posix::Path;
use std::os::{homedir};
#[cfg(test)] use std::rand::random;
#[cfg(test)] 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() {
      create_project_file(path)
    } else {
      println!("Project already exists.");
    }
}

fn create_project_file(path: &Path) {
    match File::create(path).write(TEMPLATE.as_bytes()) {
        Ok(()) => (), // succeeded
        Err(e) => println!("Failed to create project {}", path.filename()),
    }
//    Command::new("vim").arg(format!("{}", path.display())).detached();
}

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]
// TODO: Fix this test so it verifies something better.
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 std::intrinsics::TypeId;
use std::collections::HashMap;
use std::collections::hash_map::{ Occupied, Vacant };

// Stores the current pointers for concrete types.
local_data_key!(key_current: HashMap<TypeId, uint>)

/// Puts back the previous current pointer.
pub struct CurrentGuard<'a, T: 'a> {
    _val: &'a mut T,
    old_ptr: Option<uint>
}

impl<'a, T: 'static> CurrentGuard<'a, T> {
    /// Creates a new current guard.
    pub fn new(val: &mut T) -> CurrentGuard<T> {
        let id = TypeId::of::<T>();
        let ptr = val as *mut T as uint;
        let current = key_current.replace(None);
        let mut current = match current {
            None => HashMap::new(),
            Some(current) => current
        };
        let old_ptr = match current.entry(id) {
            Occupied(mut entry) => Some(entry.set(ptr)),
            Vacant(entry) => {
                entry.set(ptr);
                None
            }
        };
        key_current.replace(Some(current));
        CurrentGuard { old_ptr: old_ptr, _val: val }
    }
}

#[unsafe_destructor]
impl<'a, T: 'static> Drop for CurrentGuard<'a, T> {
    fn drop(&mut self) {
        let id = TypeId::of::<T>();
        let mut current = key_current.replace(None).unwrap();
        match self.old_ptr {
            None => {
                current.remove(&id);
                key_current.replace(Some(current));
                return;
            }
            Some(old_ptr) => {
                match current.entry(id) {
                    Occupied(mut entry) => { entry.set(old_ptr); }
                    Vacant(entry) => { entry.set(old_ptr); }
                };
            }
        };
        key_current.replace(Some(current));
    }
}

/// The current value of a type.
pub struct Current<T>;

impl<T: 'static> Current<T> {
    /// Gets mutable reference to current object.
    /// Requires mutable reference to prevent access to globals in safe code,
    /// and to prevent mutable borrows of same value in scope.
    /// Is unsafe because returned reference inherits lifetime from argument.
    pub unsafe fn current(&mut self) -> Option<&mut T> {
        use std::mem::transmute;
        let id = TypeId::of::<T>();
        let current = match key_current.replace(None) {
            None => { return None; }
            Some(current) => current
        };
        let ptr = match current.get(&id) {
            None => { return None; }
            Some(x) => *x
        };
        key_current.replace(Some(current));
        Some(transmute(ptr as *mut T))
    }

    /// Unwraps mutable reference to current object,
    /// but with nicer error message.
    pub unsafe fn current_unwrap(&mut self) -> &mut T {
        match self.current() {
            None => {
                use std::intrinsics::get_tydesc;
                let name = (*get_tydesc::<T>()).name;
                panic!("No current `{}` is set", name);
            }
            Some(x) => x
        }
    }
}

impl<T: 'static> Deref<T> for Current<T> {
    #[inline(always)]
    fn deref<'a>(&'a self) -> &'a T {
        use std::mem::transmute;
        unsafe {
            // Current does not contain anything,
            // so it is safe to transmute to mutable.
            transmute::<_, &'a mut Current<T>>(self).current_unwrap()
        }
    }
}

impl<T: 'static> DerefMut<T> for Current<T> {
    #[inline(always)]
    fn deref_mut<'a>(&'a mut self) -> &'a mut T {
        unsafe { self.current_unwrap() }
    }
}
<commit_msg>Fixed failure on unwinding due to current object not set<commit_after>
use std::intrinsics::TypeId;
use std::collections::HashMap;
use std::collections::hash_map::{ Occupied, Vacant };

// Stores the current pointers for concrete types.
local_data_key!(key_current: HashMap<TypeId, uint>)

/// Puts back the previous current pointer.
pub struct CurrentGuard<'a, T: 'a> {
    _val: &'a mut T,
    old_ptr: Option<uint>
}

impl<'a, T: 'static> CurrentGuard<'a, T> {
    /// Creates a new current guard.
    pub fn new(val: &mut T) -> CurrentGuard<T> {
        let id = TypeId::of::<T>();
        let ptr = val as *mut T as uint;
        let current = key_current.replace(None);
        let mut current = match current {
            None => HashMap::new(),
            Some(current) => current
        };
        let old_ptr = match current.entry(id) {
            Occupied(mut entry) => Some(entry.set(ptr)),
            Vacant(entry) => {
                entry.set(ptr);
                None
            }
        };
        key_current.replace(Some(current));
        CurrentGuard { old_ptr: old_ptr, _val: val }
    }
}

#[unsafe_destructor]
impl<'a, T: 'static> Drop for CurrentGuard<'a, T> {
    fn drop(&mut self) {
        let id = TypeId::of::<T>();
        let mut current = key_current.replace(None).unwrap();
        match self.old_ptr {
            None => {
                current.remove(&id);
                key_current.replace(Some(current));
                return;
            }
            Some(old_ptr) => {
                match current.entry(id) {
                    Occupied(mut entry) => { entry.set(old_ptr); }
                    Vacant(entry) => { entry.set(old_ptr); }
                };
            }
        };
        key_current.replace(Some(current));
    }
}

/// The current value of a type.
pub struct Current<T>;

impl<T: 'static> Current<T> {
    /// Gets mutable reference to current object.
    /// Requires mutable reference to prevent access to globals in safe code,
    /// and to prevent mutable borrows of same value in scope.
    /// Is unsafe because returned reference inherits lifetime from argument.
    pub unsafe fn current(&mut self) -> Option<&mut T> {
        use std::mem::transmute;
        let id = TypeId::of::<T>();
        let current = match key_current.replace(None) {
            None => { return None; }
            Some(current) => current
        };
        let ptr: Option<uint> = current.get(&id).map(|id| *id);
        key_current.replace(Some(current));
        let ptr = match ptr { None => { return None; } Some(x) => x };
        Some(transmute(ptr as *mut T))
    }

    /// Unwraps mutable reference to current object,
    /// but with nicer error message.
    pub unsafe fn current_unwrap(&mut self) -> &mut T {
        match self.current() {
            None => {
                use std::intrinsics::get_tydesc;
                let name = (*get_tydesc::<T>()).name;
                panic!("No current `{}` is set", name);
            }
            Some(x) => x
        }
    }
}

impl<T: 'static> Deref<T> for Current<T> {
    #[inline(always)]
    fn deref<'a>(&'a self) -> &'a T {
        use std::mem::transmute;
        unsafe {
            // Current does not contain anything,
            // so it is safe to transmute to mutable.
            transmute::<_, &'a mut Current<T>>(self).current_unwrap()
        }
    }
}

impl<T: 'static> DerefMut<T> for Current<T> {
    #[inline(always)]
    fn deref_mut<'a>(&'a mut self) -> &'a mut T {
        unsafe { self.current_unwrap() }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R> where R: Read {
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source: source,
            remaining_chunks_size: None,
        }
    }
}

impl<R> Read for Decoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let remaining_chunks_size = match self.remaining_chunks_size {
            Some(c) => c,
            None => {
                // first possibility: we are not in a chunk, so we'll attempt to determine
                // the chunks size
                let mut chunk_size = Vec::new();

                loop {
                    let byte = match self.source.by_ref().bytes().next() {
                        Some(b) => try!(b),
                        None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    };

                    if byte == b'\r' {
                        break;
                    }

                    chunk_size.push(byte);
                }

                match self.source.by_ref().bytes().next() {
                    Some(Ok(b'\n')) => (),
                    _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                }

                let chunk_size = match String::from_utf8(chunk_size) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                // if the chunk size is 0, we are at EOF
                if chunk_size == 0 {
                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\r')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\n')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    return Ok(0);
                }

                // now that we now the current chunk size, calling ourselves recursively
                self.remaining_chunks_size = Some(chunk_size);
                return self.read(buf);
            }
        };

        // second possibility: we continue reading from a chunk
        if buf.len() < remaining_chunks_size {
            let read = try!(self.source.read(buf));
            self.remaining_chunks_size = Some(remaining_chunks_size - read);
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= remaining_chunks_size);

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));

        self.remaining_chunks_size = if read == remaining_chunks_size {
            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\r')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            None
        } else {
            Some(remaining_chunks_size - read)
        };

        return Ok(read);
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<commit_msg>Add test for zero length<commit_after>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R> where R: Read {
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source: source,
            remaining_chunks_size: None,
        }
    }
}

impl<R> Read for Decoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let remaining_chunks_size = match self.remaining_chunks_size {
            Some(c) => c,
            None => {
                // first possibility: we are not in a chunk, so we'll attempt to determine
                // the chunks size
                let mut chunk_size = Vec::new();

                loop {
                    let byte = match self.source.by_ref().bytes().next() {
                        Some(b) => try!(b),
                        None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    };

                    if byte == b'\r' {
                        break;
                    }

                    chunk_size.push(byte);
                }

                match self.source.by_ref().bytes().next() {
                    Some(Ok(b'\n')) => (),
                    _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                }

                let chunk_size = match String::from_utf8(chunk_size) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                // if the chunk size is 0, we are at EOF
                if chunk_size == 0 {
                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\r')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\n')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    return Ok(0);
                }

                // now that we now the current chunk size, calling ourselves recursively
                self.remaining_chunks_size = Some(chunk_size);
                return self.read(buf);
            }
        };

        // second possibility: we continue reading from a chunk
        if buf.len() < remaining_chunks_size {
            let read = try!(self.source.read(buf));
            self.remaining_chunks_size = Some(remaining_chunks_size - read);
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= remaining_chunks_size);

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));

        self.remaining_chunks_size = if read == remaining_chunks_size {
            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\r')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            None
        } else {
            Some(remaining_chunks_size - read)
        };

        return Ok(read);
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn test_decode_zero_length() {
        let mut decoder = Decoder::new(b"0\r\n\r\n" as &[u8]);

        let mut decoded = String::new();
        decoder.read_to_string(&mut decoded).unwrap();

        assert_eq!(decoded, "");
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<|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);

        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;
        }
    }
}

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(_)) => Err(ErrorKind::MoreDataNeeded.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(..));
                            /*
                        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);
                        }
                        */
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }

                    Ok((seek, Event::MoreDataNeeded))
                }
                 Err(Err::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: 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);
        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>Reset the state when parsing the headers<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;
        }
    }
}

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(_)) => Err(ErrorKind::MoreDataNeeded.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(..));
                            /*
                        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);
                        }
                        */
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }

                    Ok((seek, Event::MoreDataNeeded))
                }
                 Err(Err::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: 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);
        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 std::char;
use std::cmp::{max, min};
use std::collections::VecDeque;
use num::{Complex, Float};
use rustty;
use rustty::{Attr, Color, Terminal, Cell, CellAccessor, HasSize};
use rustty::ui::{Alignable, Widget, VerticalAlign};
use itertools::{Itertools, EitherOrBoth};

pub struct Canvas {
    term: Terminal,
    spectrum: Widget,
    waterfall: Widget,
    history: VecDeque<Vec<f32>>,
}

impl Canvas {
    pub fn new() -> Result<Self, rustty::Error> {
        let term = try!(Terminal::new());
        let (cols, rows) = term.size();

        let spectrum_height = rows / 2;
        let waterfall_height = if rows % 2 != 0 { rows / 2 } else { rows / 2 + 1 };

        let mut spectrum = Widget::new(cols, spectrum_height);
        spectrum.valign(&term, VerticalAlign::Top, 0);

        let mut waterfall = Widget::new(cols, waterfall_height);
        waterfall.valign(&term, VerticalAlign::Bottom, 0);

        Ok(Canvas {
            term: term,
            spectrum: spectrum,
            waterfall: waterfall,
            history: VecDeque::with_capacity(waterfall_height * 2),
        })
    }

    /// Adds a spectrum to the history and draws it on the waterfall
    /// and the spectrum view.
    pub fn add_spectrum(&mut self, spec: Vec<Complex<f32>>) {
        draw_spectrum(&mut self.spectrum, &spec);

        let normalized = normalize_spectrum(&spec);
        self.history.push_front(normalized);

        draw_waterfall(&mut self.waterfall, &self.history);

        self.spectrum.draw_into(&mut self.term);
        self.waterfall.draw_into(&mut self.term);
        self.term.swap_buffers().unwrap();
    }

    pub fn get_term(&mut self) -> &mut Terminal {
        &mut self.term
    }
}

fn draw_waterfall<T: CellAccessor + HasSize>(canvas: &mut T, spectra: &VecDeque<Vec<f32>>) {
    let (cols, rows) = canvas.size();
    for (row, mut specs) in (0..rows).zip(&spectra.iter().chunks_lazy(2)) {
        if let Some(upper_heights) = specs.next().map(|vec| vec.iter()) {
            match specs.next().map(|vec| vec.iter()) {
                Some(lower_heights) => {
                    for (c, heights) in (0..cols).zip(upper_heights.zip_longest(lower_heights)) {
                        let (u, l) = match heights {
                            EitherOrBoth::Both(&upper, &lower) => (upper, lower),
                            EitherOrBoth::Left(&upper) => (upper, 0.0),
                            EitherOrBoth::Right(&lower) => (0.0, lower),
                        };
                        *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(u, l);
                    }
                },
                None => {
                    for (c, u) in (0..cols).zip(upper_heights) {
                        *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(*u, 0.0);
                    }
                }
            }
        }
    }
}

fn spectrum_heights_to_waterfall_cell(upper: f32, lower: f32) -> Cell {
    Cell::new('▀',
              Color::Byte(color_mapping(upper)),
              Color::Byte(color_mapping(lower)),
              Attr::Default)
}

fn color_mapping(f: f32) -> u8 {
    let lower = 16.0;
    let upper = 231.0;
    let mapped = f * (upper - lower) + lower;
    if mapped < lower {
        lower as u8
    } else if mapped > upper {
        upper as u8
    } else {
        mapped as u8
    }
}

fn normalize_spectrum(spec: &[Complex<f32>]) -> Vec<f32> {
    // FFT shift
    let (first_half, last_half) = spec.split_at((spec.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());

    // normalize and take the log
    shifted_spec.map(Complex::norm).map(Float::log10).collect()
}

// indexing is from the top of the cell
fn pixel_nums_to_braille(p1: Option<u8>, p2: Option<u8>) -> char {
    let pixel_map = [[0x01, 0x08],
                     [0x02, 0x10],
                     [0x04, 0x20],
                     [0x40, 0x80]];

    let mut c = 0;
    if let Some(p) = p1 {
        for i in p..4 {
            c |= pixel_map[i as usize][0];
        }
    }

    if let Some(p) = p2 {
        for i in p..4 {
            c |= pixel_map[i as usize][1];
        }
    }

    char::from_u32((0x2800 + c) as u32).unwrap()
}

fn fft_shift<T: Clone>(spec: &mut [T]) {
    let spec_copy = spec.to_owned();

    let (first_half, last_half) = spec_copy.split_at((spec_copy.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());
    for (x, y) in spec.iter_mut().zip(shifted_spec) {
        *x = y.clone();
    }
}

fn spectrum_to_bin_heights(spec: &[Complex<f32>], dest: &mut [f32]) {
    //TODO should be plotting in log scale

    // subsample
    let mut last_idx = -1isize;
    for (i, x) in spec.iter().map(Complex::norm).enumerate() {
        if (i * dest.len() / spec.len()) as isize > last_idx {
            last_idx += 1;
            dest[last_idx as usize] = x;
        }
    }

    //TODO unnecessary allocation
    fft_shift(dest);
}

fn char_to_cell(c: char) -> Cell {
    let mut cell = Cell::with_char(c);
    cell.set_attrs(Attr::Bold);
    cell
}

fn draw_pixel_pair<T>(canvas: &mut T, col_idx: usize, p1: usize, p2: usize)
    where T: CellAccessor + HasSize
{
    let (_, rows) = canvas.size();
    let max_pixel_height = 4 * rows;

    // clamp heights
    let p1 = if p1 >= max_pixel_height { max_pixel_height - 1} else { p1 };
    let p2 = if p2 >= max_pixel_height { max_pixel_height - 1} else { p2 };

    // Reverse it, since the terminal indexing is from the top
    let p1 = max_pixel_height - p1 - 1;
    let p2 = max_pixel_height - p2 - 1;

    // cell indices
    let c1 = p1 / 4;
    let c2 = p2 / 4;

    // Fill in full height cells.
    let full_cell_char = pixel_nums_to_braille(Some(0), Some(0));
    for row_idx in max(c1, c2)..rows {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(full_cell_char);
    }

    let left_fill_cell_char = pixel_nums_to_braille(Some(0), None);
    for row_idx in min(c1, c2)..c2 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(left_fill_cell_char);
    }

    let right_fill_cell_char = pixel_nums_to_braille(None, Some(0));
    for row_idx in min(c1, c2)..c1 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(right_fill_cell_char);
    }

    // Now fill in partial height cells.
    if c1 == c2 {
        // top pixels are in the same cell
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some((p2 % 4) as u8)));
    } else if c1 > c2 {
        // right pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some(0)));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(None, Some((p2 % 4) as u8)));
    } else {
        // left pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), None));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some(0), Some((p2 % 4) as u8)));
    }
}

fn draw_spectrum<T: CellAccessor + HasSize>(canvas: &mut T, spec: &[Complex<f32>]) {
    canvas.clear(Cell::default());
    let (num_cols, num_rows) = canvas.size();
    let pixel_height = num_rows * 4;
    let pixel_width = num_cols * 2;
    // TODO what should this value be?
    let max_height = 500.0;

    let mut bins = vec![0.0; pixel_width];
    spectrum_to_bin_heights(spec, &mut bins[..]);

    for col_idx in 0..num_cols {
        // height in float between 0 and 1.
        let h1 = bins[col_idx * 2] / max_height;
        let h2 = bins[col_idx * 2 + 1] / max_height;

        // The "pixel" height of each point.
        let p1 = (h1 * pixel_height as f32).floor() as usize;
        let p2 = (h2 * pixel_height as f32).floor() as usize;

        draw_pixel_pair(canvas, col_idx, p1, p2);
    }
}

#[cfg(test)]
mod tests {
    use super::{pixel_nums_to_braille, fft_shift, draw_pixel_pair};
    use rustty::Terminal;

    #[test]
    fn test_fft_shift_dc() {
        let len = 9;
        let mut spec = vec![0; len];
        spec[0] = 1;
        fft_shift(&mut spec[..]);
        assert_eq!(spec[len / 2], 1);
    }

    #[test]
    fn test_fft_shift_even() {
        let mut before: Vec<usize> = (0..10).collect();
        let after = vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
        fft_shift(&mut before[..]);
        assert_eq!(before, after);
    }

    #[test]
    fn test_pixel_nums() {
        assert_eq!(pixel_nums_to_braille(Some(0), Some(0)), '⣿');
        assert_eq!(pixel_nums_to_braille(Some(1), Some(2)), '⣦');
        assert_eq!(pixel_nums_to_braille(None, Some(3)), '⢀');
        assert_eq!(pixel_nums_to_braille(Some(2), None), '⡄');
        assert_eq!(pixel_nums_to_braille(None, None), '⠀');
    }

    #[test]
    fn test_draw_pixel_pair() {
        let mut term = Terminal::new().unwrap();

        // Test drawing with the same top cell
        draw_pixel_pair(&mut term, 0, 4, 6);
        assert_eq!(term[(0, term.rows() - 3)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣰');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        // Test drawing with the top pixel in each column being in
        // different cells
        draw_pixel_pair(&mut term, 0, 4, 8);
        assert_eq!(term[(0, term.rows() - 4)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 3)].ch(), '⢀');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣸');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        draw_pixel_pair(&mut term, 1, 13, 2);
        assert_eq!(term[(1, term.rows() - 5)].ch(), ' ');
        assert_eq!(term[(1, term.rows() - 4)].ch(), '⡄');
        assert_eq!(term[(1, term.rows() - 3)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 2)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 1)].ch(), '⣷');
        term.clear().unwrap();
    }
}
<commit_msg>Plot waterfall in grayscale with correct scaling.<commit_after>use std::char;
use std::cmp::{max, min};
use std::collections::VecDeque;
use num::{Complex, Float};
use rustty;
use rustty::{Attr, Color, Terminal, Cell, CellAccessor, HasSize};
use rustty::ui::{Alignable, Widget, VerticalAlign};
use itertools::{Itertools, EitherOrBoth};

pub struct Canvas {
    term: Terminal,
    spectrum: Widget,
    waterfall: Widget,
    history: VecDeque<Vec<f32>>,
}

impl Canvas {
    pub fn new() -> Result<Self, rustty::Error> {
        let term = try!(Terminal::new());
        let (cols, rows) = term.size();

        let spectrum_height = rows / 2;
        let waterfall_height = if rows % 2 != 0 { rows / 2 } else { rows / 2 + 1 };

        let mut spectrum = Widget::new(cols, spectrum_height);
        spectrum.valign(&term, VerticalAlign::Top, 0);

        let mut waterfall = Widget::new(cols, waterfall_height);
        waterfall.valign(&term, VerticalAlign::Bottom, 0);

        Ok(Canvas {
            term: term,
            spectrum: spectrum,
            waterfall: waterfall,
            history: VecDeque::with_capacity(waterfall_height * 2),
        })
    }

    /// Adds a spectrum to the history and draws it on the waterfall
    /// and the spectrum view.
    pub fn add_spectrum(&mut self, spec: Vec<Complex<f32>>) {
        draw_spectrum(&mut self.spectrum, &spec);

        let normalized = normalize_spectrum(&spec, 26.0);
        let averaged = normalized.chunks(2).map(|v| (v[0] + v[1]) / 2.0).collect();
        self.history.push_front(averaged);

        draw_waterfall(&mut self.waterfall, &self.history);

        self.spectrum.draw_into(&mut self.term);
        self.waterfall.draw_into(&mut self.term);
        self.term.swap_buffers().unwrap();
    }

    pub fn get_term(&mut self) -> &mut Terminal {
        &mut self.term
    }
}

fn draw_waterfall<T: CellAccessor + HasSize>(canvas: &mut T, spectra: &VecDeque<Vec<f32>>) {
    let (cols, rows) = canvas.size();
    for (row, mut specs) in (0..rows).zip(&spectra.iter().chunks_lazy(2)) {
        if let Some(upper_heights) = specs.next().map(|vec| vec.iter()) {
            match specs.next().map(|vec| vec.iter()) {
                Some(lower_heights) => {
                    for (c, heights) in (0..cols).zip(upper_heights.zip_longest(lower_heights)) {
                        let (u, l) = match heights {
                            EitherOrBoth::Both(&upper, &lower) => (upper, lower),
                            EitherOrBoth::Left(&upper) => (upper, 0.0),
                            EitherOrBoth::Right(&lower) => (0.0, lower),
                        };
                        *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(u, l);
                    }
                },
                None => {
                    for (c, u) in (0..cols).zip(upper_heights) {
                        *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(*u, 0.0);
                    }
                }
            }
        }
    }
}

fn spectrum_heights_to_waterfall_cell(upper: f32, lower: f32) -> Cell {
    Cell::new('▀',
              Color::Byte(color_mapping(upper)),
              Color::Byte(color_mapping(lower)),
              Attr::Default)
}

/// Assumes `f` is between 0 and 1. Anything outside of this range
/// will be clamped.
fn color_mapping(f: f32) -> u8 {
    //let lower = 16.0;
    //let upper = 231.0;
    let lower = 232.0;
    let upper = 255.0;
    let mapped = f * (upper - lower) + lower;
    if mapped < lower {
        lower as u8
    } else if mapped > upper {
        upper as u8
    } else {
        mapped as u8
    }
}

fn normalize_spectrum(spec: &[Complex<f32>], max_db: f32) -> Vec<f32> {
    // FFT shift
    let (first_half, last_half) = spec.split_at((spec.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());

    // normalize and take the log
    shifted_spec.map(Complex::norm)
                .map(Float::log10)
                .map(|x| 10.0 * x)
                .map(|x| x / max_db)
                .collect()
}

// indexing is from the top of the cell
fn pixel_nums_to_braille(p1: Option<u8>, p2: Option<u8>) -> char {
    let pixel_map = [[0x01, 0x08],
                     [0x02, 0x10],
                     [0x04, 0x20],
                     [0x40, 0x80]];

    let mut c = 0;
    if let Some(p) = p1 {
        for i in p..4 {
            c |= pixel_map[i as usize][0];
        }
    }

    if let Some(p) = p2 {
        for i in p..4 {
            c |= pixel_map[i as usize][1];
        }
    }

    char::from_u32((0x2800 + c) as u32).unwrap()
}

fn fft_shift<T: Clone>(spec: &mut [T]) {
    let spec_copy = spec.to_owned();

    let (first_half, last_half) = spec_copy.split_at((spec_copy.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());
    for (x, y) in spec.iter_mut().zip(shifted_spec) {
        *x = y.clone();
    }
}

fn spectrum_to_bin_heights(spec: &[Complex<f32>], dest: &mut [f32]) {
    //TODO should be plotting in log scale

    // subsample
    let mut last_idx = -1isize;
    for (i, x) in spec.iter().map(Complex::norm).enumerate() {
        if (i * dest.len() / spec.len()) as isize > last_idx {
            last_idx += 1;
            dest[last_idx as usize] = x;
        }
    }

    //TODO unnecessary allocation
    fft_shift(dest);
}

fn char_to_cell(c: char) -> Cell {
    let mut cell = Cell::with_char(c);
    cell.set_attrs(Attr::Bold);
    cell
}

fn draw_pixel_pair<T>(canvas: &mut T, col_idx: usize, p1: usize, p2: usize)
    where T: CellAccessor + HasSize
{
    let (_, rows) = canvas.size();
    let max_pixel_height = 4 * rows;

    // clamp heights
    let p1 = if p1 >= max_pixel_height { max_pixel_height - 1} else { p1 };
    let p2 = if p2 >= max_pixel_height { max_pixel_height - 1} else { p2 };

    // Reverse it, since the terminal indexing is from the top
    let p1 = max_pixel_height - p1 - 1;
    let p2 = max_pixel_height - p2 - 1;

    // cell indices
    let c1 = p1 / 4;
    let c2 = p2 / 4;

    // Fill in full height cells.
    let full_cell_char = pixel_nums_to_braille(Some(0), Some(0));
    for row_idx in max(c1, c2)..rows {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(full_cell_char);
    }

    let left_fill_cell_char = pixel_nums_to_braille(Some(0), None);
    for row_idx in min(c1, c2)..c2 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(left_fill_cell_char);
    }

    let right_fill_cell_char = pixel_nums_to_braille(None, Some(0));
    for row_idx in min(c1, c2)..c1 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(right_fill_cell_char);
    }

    // Now fill in partial height cells.
    if c1 == c2 {
        // top pixels are in the same cell
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some((p2 % 4) as u8)));
    } else if c1 > c2 {
        // right pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some(0)));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(None, Some((p2 % 4) as u8)));
    } else {
        // left pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), None));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some(0), Some((p2 % 4) as u8)));
    }
}

fn draw_spectrum<T: CellAccessor + HasSize>(canvas: &mut T, spec: &[Complex<f32>]) {
    canvas.clear(Cell::default());
    let (num_cols, num_rows) = canvas.size();
    let pixel_height = num_rows * 4;
    let pixel_width = num_cols * 2;
    // TODO what should this value be?
    let max_height = 500.0;

    let mut bins = vec![0.0; pixel_width];
    spectrum_to_bin_heights(spec, &mut bins[..]);

    for col_idx in 0..num_cols {
        // height in float between 0 and 1.
        let h1 = bins[col_idx * 2] / max_height;
        let h2 = bins[col_idx * 2 + 1] / max_height;

        // The "pixel" height of each point.
        let p1 = (h1 * pixel_height as f32).floor() as usize;
        let p2 = (h2 * pixel_height as f32).floor() as usize;

        draw_pixel_pair(canvas, col_idx, p1, p2);
    }
}

#[cfg(test)]
mod tests {
    use super::{pixel_nums_to_braille, fft_shift, draw_pixel_pair};
    use rustty::Terminal;

    #[test]
    fn test_fft_shift_dc() {
        let len = 9;
        let mut spec = vec![0; len];
        spec[0] = 1;
        fft_shift(&mut spec[..]);
        assert_eq!(spec[len / 2], 1);
    }

    #[test]
    fn test_fft_shift_even() {
        let mut before: Vec<usize> = (0..10).collect();
        let after = vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
        fft_shift(&mut before[..]);
        assert_eq!(before, after);
    }

    #[test]
    fn test_pixel_nums() {
        assert_eq!(pixel_nums_to_braille(Some(0), Some(0)), '⣿');
        assert_eq!(pixel_nums_to_braille(Some(1), Some(2)), '⣦');
        assert_eq!(pixel_nums_to_braille(None, Some(3)), '⢀');
        assert_eq!(pixel_nums_to_braille(Some(2), None), '⡄');
        assert_eq!(pixel_nums_to_braille(None, None), '⠀');
    }

    #[test]
    fn test_draw_pixel_pair() {
        let mut term = Terminal::new().unwrap();

        // Test drawing with the same top cell
        draw_pixel_pair(&mut term, 0, 4, 6);
        assert_eq!(term[(0, term.rows() - 3)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣰');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        // Test drawing with the top pixel in each column being in
        // different cells
        draw_pixel_pair(&mut term, 0, 4, 8);
        assert_eq!(term[(0, term.rows() - 4)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 3)].ch(), '⢀');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣸');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        draw_pixel_pair(&mut term, 1, 13, 2);
        assert_eq!(term[(1, term.rows() - 5)].ch(), ' ');
        assert_eq!(term[(1, term.rows() - 4)].ch(), '⡄');
        assert_eq!(term[(1, term.rows() - 3)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 2)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 1)].ch(), '⣷');
        term.clear().unwrap();
    }
}
<|endoftext|>


<commit_before>mod select_dsl;
mod cache_limit_dsl;
mod dump_dsl;
mod database_unmap_dsl;
mod load_dsl;
mod status_dsl;
mod delete_dsl;
mod io_flush_dsl;
mod log_level_dsl;
mod log_put_dsl;
mod log_reopen_dsl;
mod table_create_dsl;
mod table_list_dsl;
mod table_rename_dsl;
mod table_remove_dsl;
mod column_copy_dsl;
mod column_create_dsl;
mod column_list_dsl;
mod column_rename_dsl;
mod column_remove_dsl;
mod reindex_dsl;
mod request_cancel_dsl;
mod schema_dsl;
mod shutdown_dsl;
mod object_exist_dsl;
mod object_inspect_dsl;
mod object_remove_dsl;
mod normalize_dsl;
mod normalizer_list_dsl;
mod plugin_register_dsl;
mod plugin_unregister_dsl;
mod thread_limit_dsl;
mod tokenize_dsl;
mod tokenizer_list_dsl;
mod truncate_dsl;
#[cfg(feature="sharding")]
mod logical_select_dsl;
#[cfg(feature="sharding")]
mod logical_shard_list_dsl;
#[cfg(feature="sharding")]
mod logical_table_remove_dsl;

pub use self::select_dsl::select;
pub use self::cache_limit_dsl::cache_limit;
pub use self::dump_dsl::dump;
pub use self::database_unmap_dsl::database_unmap;
pub use self::load_dsl::load;
pub use self::status_dsl::status;
pub use self::delete_dsl::delete;
pub use self::io_flush_dsl::io_flush;
pub use self::log_level_dsl::log_level;
pub use self::log_put_dsl::log_put;
pub use self::log_reopen_dsl::log_reopen;
pub use self::table_create_dsl::table_create;
pub use self::table_list_dsl::table_list;
pub use self::table_rename_dsl::table_rename;
pub use self::table_remove_dsl::table_remove;
pub use self::column_copy_dsl::column_copy;
pub use self::column_create_dsl::column_create;
pub use self::column_list_dsl::column_list;
pub use self::column_rename_dsl::column_rename;
pub use self::column_remove_dsl::column_remove;
pub use self::reindex_dsl::reindex;
pub use self::request_cancel_dsl::request_cancel;
pub use self::schema_dsl::schema;
pub use self::shutdown_dsl::shutdown;
pub use self::object_exist_dsl::object_exist;
pub use self::object_inspect_dsl::object_inspect;
pub use self::object_remove_dsl::object_remove;
pub use self::normalize_dsl::normalize;
pub use self::normalizer_list_dsl::normalizer_list;
pub use self::plugin_register_dsl::plugin_register;
pub use self::plugin_unregister_dsl::plugin_unregister;
pub use self::thread_limit_dsl::thread_limit;
pub use self::tokenize_dsl::tokenize;
pub use self::tokenizer_list_dsl::tokenizer_list;
pub use self::truncate_dsl::truncate;
#[cfg(feature="sharding")]
pub use self::logical_select_dsl::logical_select;
#[cfg(feature="sharding")]
pub use self::logical_shard_list_dsl::logical_shard_list;
#[cfg(feature="sharding")]
pub use self::logical_table_remove_dsl::logical_table_remove;
<commit_msg>dsl: Replace `self::` with `dsl::`<commit_after>mod select_dsl;
mod cache_limit_dsl;
mod dump_dsl;
mod database_unmap_dsl;
mod load_dsl;
mod status_dsl;
mod delete_dsl;
mod io_flush_dsl;
mod log_level_dsl;
mod log_put_dsl;
mod log_reopen_dsl;
mod table_create_dsl;
mod table_list_dsl;
mod table_rename_dsl;
mod table_remove_dsl;
mod column_copy_dsl;
mod column_create_dsl;
mod column_list_dsl;
mod column_rename_dsl;
mod column_remove_dsl;
mod reindex_dsl;
mod request_cancel_dsl;
mod schema_dsl;
mod shutdown_dsl;
mod object_exist_dsl;
mod object_inspect_dsl;
mod object_remove_dsl;
mod normalize_dsl;
mod normalizer_list_dsl;
mod plugin_register_dsl;
mod plugin_unregister_dsl;
mod thread_limit_dsl;
mod tokenize_dsl;
mod tokenizer_list_dsl;
mod truncate_dsl;
#[cfg(feature="sharding")]
mod logical_select_dsl;
#[cfg(feature="sharding")]
mod logical_shard_list_dsl;
#[cfg(feature="sharding")]
mod logical_table_remove_dsl;

pub use dsl::select_dsl::select;
pub use dsl::cache_limit_dsl::cache_limit;
pub use dsl::dump_dsl::dump;
pub use dsl::database_unmap_dsl::database_unmap;
pub use dsl::load_dsl::load;
pub use dsl::status_dsl::status;
pub use dsl::delete_dsl::delete;
pub use dsl::io_flush_dsl::io_flush;
pub use dsl::log_level_dsl::log_level;
pub use dsl::log_put_dsl::log_put;
pub use dsl::log_reopen_dsl::log_reopen;
pub use dsl::table_create_dsl::table_create;
pub use dsl::table_list_dsl::table_list;
pub use dsl::table_rename_dsl::table_rename;
pub use dsl::table_remove_dsl::table_remove;
pub use dsl::column_copy_dsl::column_copy;
pub use dsl::column_create_dsl::column_create;
pub use dsl::column_list_dsl::column_list;
pub use dsl::column_rename_dsl::column_rename;
pub use dsl::column_remove_dsl::column_remove;
pub use dsl::reindex_dsl::reindex;
pub use dsl::request_cancel_dsl::request_cancel;
pub use dsl::schema_dsl::schema;
pub use dsl::shutdown_dsl::shutdown;
pub use dsl::object_exist_dsl::object_exist;
pub use dsl::object_inspect_dsl::object_inspect;
pub use dsl::object_remove_dsl::object_remove;
pub use dsl::normalize_dsl::normalize;
pub use dsl::normalizer_list_dsl::normalizer_list;
pub use dsl::plugin_register_dsl::plugin_register;
pub use dsl::plugin_unregister_dsl::plugin_unregister;
pub use dsl::thread_limit_dsl::thread_limit;
pub use dsl::tokenize_dsl::tokenize;
pub use dsl::tokenizer_list_dsl::tokenizer_list;
pub use dsl::truncate_dsl::truncate;
#[cfg(feature="sharding")]
pub use dsl::logical_select_dsl::logical_select;
#[cfg(feature="sharding")]
pub use dsl::logical_shard_list_dsl::logical_shard_list;
#[cfg(feature="sharding")]
pub use dsl::logical_table_remove_dsl::logical_table_remove;
<|endoftext|>