<commit_before>// Copyright 2019, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use ::ffi::{self, cairo_status_t};
use ::{Status, Surface, UserDataKey};

use libc::{c_void, c_double, c_uchar, c_uint};
use std::any::Any;
use std::cell::{Cell, RefCell, RefMut};
use std::io;
use std::panic::AssertUnwindSafe;
use std::rc::Rc;

macro_rules! for_stream_constructors {
    ($constructor_ffi: ident) => {
        /// Takes full ownership of the output stream,
        /// which is not allowed to borrow any lifetime shorter than `'static`.
        ///
        /// Because the underlying `cairo_surface_t` is reference-counted,
        /// a lifetime parameter in a Rust wrapper type would not be enough to track
        /// how long it can keep writing to the stream.
        pub fn for_stream<W: io::Write + 'static>(width: f64, height: f64, stream: W) -> Self {
            Self {
                inner: Surface::_for_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }

        /// Allows writing to a borrowed stream. The lifetime of the borrow is not tracked.
        ///
        /// # Safety
        ///
        /// The value that `stream` points to must live at least until the underlying `cairo_surface_t`
        /// (which maybe be longer then the Rust `PdfSurface` wrapper, because of reference-counting),
        /// or until the output stream is removed from the surface with [`Surface::take_output_stream`].
        ///
        /// Since the former is hard to track for sure, the latter is strongly recommended.
        /// The concrete type behind the `Box<dyn Any>` value returned by `take_output_stream`
        /// is private, so you won’t be able to downcast it.
        /// But removing it anyway ensures that later writes do no go through a dangling pointer.
        pub unsafe fn for_raw_stream<W: io::Write + 'static>(width: f64, height: f64, stream: *mut W) -> Self {
            Self {
                inner: Surface::_for_raw_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }
    };
}

impl Surface {
    pub(crate) fn _for_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: W,
    ) -> Self {
        let env_rc = Rc::new(CallbackEnvironment {
            mutable: RefCell::new(MutableCallbackEnvironment {
                stream: Some(Box::new(stream)),
                io_error: None,
                unwind_payload: None,
            }),
            saw_already_borrowed: Cell::new(false),
        });
        let env: *const CallbackEnvironment = &*env_rc;
        unsafe {
            let ptr = constructor(Some(write_callback::<W>), env as *mut c_void, width, height);
            let surface = Surface::from_raw_full(ptr);
            surface.set_user_data(&STREAM_CALLBACK_ENVIRONMENT, env_rc);
            surface
        }
    }

    pub(crate) unsafe fn _for_raw_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: *mut W,
    ) -> Self {
        Self::_for_stream(constructor, width, height, RawStream(stream))
    }

    fn with_stream_env<'surface, R, F>(&'surface self, f: F) -> Option<R>
        where F: FnOnce(RefMut<'surface, MutableCallbackEnvironment>) -> Option<R>
    {
        let env = self.get_user_data_ptr(&STREAM_CALLBACK_ENVIRONMENT)?;

        // Safety: since `STREAM_CALLBACK_ENVIRONMENT` is private and we never
        // call `set_user_data` again or `remove_user_data` with it,
        // the contract of `get_user_data_ptr` says that the user data entry
        // lives as long as the underlying `cairo_surface_t`
        // which is at least as long as this `Surface` wrapper.
        let env: &'surface _ = unsafe { &*env.as_ptr() };

        if env.saw_already_borrowed.get() {
            panic!("The output stream’s RefCell was already borrowed when cairo attempted a write")
        }

        let mut mutable = env.mutable.borrow_mut();
        if let Some(payload) = mutable.unwind_payload.take() {
            std::panic::resume_unwind(payload)
        }
        f(mutable)
    }

    /// Remove and return the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`Box::downcast`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_output_stream(&self) -> Option<Box<dyn Any>> {
        self.with_stream_env(|mut env| env.stream.take())
    }

    /// Remove and return the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_io_error(&self) -> Result<(), io::Error> {
        match self.with_stream_env(|mut env| env.io_error.take()) {
            Some(error) => Err(error),
            None => Ok(()),
        }
    }

    /// Borrow the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`<dyn Any>::downcast_mut`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_output_stream(&self) -> Option<RefMut<dyn Any>> {
        self.with_stream_env(|env| {
            env.stream.as_ref()?;
            Some(RefMut::map(env, |env| &mut **env.stream.as_mut().unwrap()))
        })
    }

    /// Borrow the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_io_error(&self) -> Option<RefMut<io::Error>> {
        self.with_stream_env(|env| {
            env.io_error.as_ref()?;
            Some(RefMut::map(env, |env| env.io_error.as_mut().unwrap()))
        })
    }
}

pub(crate) type Constructor = unsafe extern fn(
    ffi::cairo_write_func_t,
    *mut c_void,
    c_double,
    c_double,
) -> *mut ffi::cairo_surface_t;

static STREAM_CALLBACK_ENVIRONMENT: UserDataKey<CallbackEnvironment> =
    UserDataKey::new();

struct CallbackEnvironment {
    mutable: RefCell<MutableCallbackEnvironment>,
    saw_already_borrowed: Cell<bool>,
}

struct MutableCallbackEnvironment {
    stream: Option<Box<dyn Any>>,
    io_error: Option<io::Error>,
    unwind_payload: Option<Box<dyn Any + Send + 'static>>,
}

// Safety: unwinding into C is undefined behavior (https://github.com/rust-lang/rust/issues/58794)
// so code outside of the `catch_unwind` call must never panic.
extern "C" fn write_callback<W: io::Write + 'static>(
    env: *mut c_void,
    data: *mut c_uchar,
    length: c_uint,
) -> cairo_status_t {
    // This is consistent with the type of `env` in `Surface::_for_stream`.
    let env: *const CallbackEnvironment = env as _;

    // Safety: the user data entry keeps `Rc<CallbackEnvironment>` alive
    // until the surface is destroyed.
    // If this is called by cairo, the surface is still alive.
    let env: &CallbackEnvironment = unsafe { &*env };

    if let Ok(mut mutable) = env.mutable.try_borrow_mut() {
        if let MutableCallbackEnvironment {
            stream: Some(stream),
            // Don’t attempt another write if a previous one errored or panicked:
            io_error: io_error @ None,
            unwind_payload: unwind_payload @ None,
        } = &mut *mutable {
            // Safety: `write_callback<W>` was instanciated in `Surface::_for_stream`
            // with a W parameter consistent with the box that was unsized to `Box<dyn Any>`.
            let stream = unsafe {
                stream.downcast_mut_unchecked::<W>()
            };
            // Safety: this is the callback contract from cairo’s API
            let data = unsafe {
                std::slice::from_raw_parts(data, length as usize)
            };
            // Because `<W as Write>::write_all` is a generic,
            // we must conservatively assume that it can panic.
            let result = std::panic::catch_unwind(AssertUnwindSafe(|| stream.write_all(data)));
            match result {
                Ok(Ok(())) => {
                    return Status::Success.into()
                }
                Ok(Err(error)) => {
                    *io_error = Some(error);
                }
                Err(payload) => {
                    *unwind_payload = Some(payload);
                }
            }
        }
    } else {
        env.saw_already_borrowed.set(true)
    }
    Status::WriteError.into()
}

struct RawStream<W>(*mut W);

impl<W: io::Write> io::Write for RawStream<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { unsafe { (*self.0).write(buf) }}
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { unsafe { (*self.0).write_all(buf) }}
    fn flush(&mut self) -> io::Result<()> { unsafe { (*self.0).flush() } }
}

trait AnyExt {
    /// Any::downcast_mut, but YOLO
    unsafe fn downcast_mut_unchecked<T>(&mut self) -> &mut T {
        let ptr = self as *mut Self as *mut T;
        &mut *ptr
    }
}
impl AnyExt for dyn Any {}
<commit_msg>Add explicit lifetime annotations<commit_after>// Copyright 2019, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use ::ffi::{self, cairo_status_t};
use ::{Status, Surface, UserDataKey};

use libc::{c_void, c_double, c_uchar, c_uint};
use std::any::Any;
use std::cell::{Cell, RefCell, RefMut};
use std::io;
use std::panic::AssertUnwindSafe;
use std::rc::Rc;

macro_rules! for_stream_constructors {
    ($constructor_ffi: ident) => {
        /// Takes full ownership of the output stream,
        /// which is not allowed to borrow any lifetime shorter than `'static`.
        ///
        /// Because the underlying `cairo_surface_t` is reference-counted,
        /// a lifetime parameter in a Rust wrapper type would not be enough to track
        /// how long it can keep writing to the stream.
        pub fn for_stream<W: io::Write + 'static>(width: f64, height: f64, stream: W) -> Self {
            Self {
                inner: Surface::_for_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }

        /// Allows writing to a borrowed stream. The lifetime of the borrow is not tracked.
        ///
        /// # Safety
        ///
        /// The value that `stream` points to must live at least until the underlying `cairo_surface_t`
        /// (which maybe be longer then the Rust `PdfSurface` wrapper, because of reference-counting),
        /// or until the output stream is removed from the surface with [`Surface::take_output_stream`].
        ///
        /// Since the former is hard to track for sure, the latter is strongly recommended.
        /// The concrete type behind the `Box<dyn Any>` value returned by `take_output_stream`
        /// is private, so you won’t be able to downcast it.
        /// But removing it anyway ensures that later writes do no go through a dangling pointer.
        pub unsafe fn for_raw_stream<W: io::Write + 'static>(width: f64, height: f64, stream: *mut W) -> Self {
            Self {
                inner: Surface::_for_raw_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }
    };
}

impl Surface {
    pub(crate) fn _for_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: W,
    ) -> Self {
        let env_rc = Rc::new(CallbackEnvironment {
            mutable: RefCell::new(MutableCallbackEnvironment {
                stream: Some(Box::new(stream)),
                io_error: None,
                unwind_payload: None,
            }),
            saw_already_borrowed: Cell::new(false),
        });
        let env: *const CallbackEnvironment = &*env_rc;
        unsafe {
            let ptr = constructor(Some(write_callback::<W>), env as *mut c_void, width, height);
            let surface = Surface::from_raw_full(ptr);
            surface.set_user_data(&STREAM_CALLBACK_ENVIRONMENT, env_rc);
            surface
        }
    }

    pub(crate) unsafe fn _for_raw_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: *mut W,
    ) -> Self {
        Self::_for_stream(constructor, width, height, RawStream(stream))
    }

    fn with_stream_env<'surface, R, F>(&'surface self, f: F) -> Option<R>
        where F: FnOnce(RefMut<'surface, MutableCallbackEnvironment>) -> Option<R>
    {
        let env = self.get_user_data_ptr(&STREAM_CALLBACK_ENVIRONMENT)?;

        // Safety: since `STREAM_CALLBACK_ENVIRONMENT` is private and we never
        // call `set_user_data` again or `remove_user_data` with it,
        // the contract of `get_user_data_ptr` says that the user data entry
        // lives as long as the underlying `cairo_surface_t`
        // which is at least as long as this `Surface` wrapper.
        let env: &'surface _ = unsafe { &*env.as_ptr() };

        if env.saw_already_borrowed.get() {
            panic!("The output stream’s RefCell was already borrowed when cairo attempted a write")
        }

        let mut mutable = env.mutable.borrow_mut();
        if let Some(payload) = mutable.unwind_payload.take() {
            std::panic::resume_unwind(payload)
        }
        f(mutable)
    }

    /// Remove and return the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`Box::downcast`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_output_stream(&self) -> Option<Box<dyn Any>> {
        self.with_stream_env(|mut env| env.stream.take())
    }

    /// Remove and return the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_io_error(&self) -> Result<(), io::Error> {
        match self.with_stream_env(|mut env| env.io_error.take()) {
            Some(error) => Err(error),
            None => Ok(()),
        }
    }

    /// Borrow the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`<dyn Any>::downcast_mut`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_output_stream<'surface>(&'surface self) -> Option<RefMut<'surface, dyn Any>> {
        self.with_stream_env(|env| {
            env.stream.as_ref()?;
            Some(RefMut::map(env, |env| &mut **env.stream.as_mut().unwrap()))
        })
    }

    /// Borrow the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_io_error<'surface>(&'surface self) -> Option<RefMut<'surface, io::Error>> {
        self.with_stream_env(|env| {
            env.io_error.as_ref()?;
            Some(RefMut::map(env, |env| env.io_error.as_mut().unwrap()))
        })
    }
}

pub(crate) type Constructor = unsafe extern fn(
    ffi::cairo_write_func_t,
    *mut c_void,
    c_double,
    c_double,
) -> *mut ffi::cairo_surface_t;

static STREAM_CALLBACK_ENVIRONMENT: UserDataKey<CallbackEnvironment> =
    UserDataKey::new();

struct CallbackEnvironment {
    mutable: RefCell<MutableCallbackEnvironment>,
    saw_already_borrowed: Cell<bool>,
}

struct MutableCallbackEnvironment {
    stream: Option<Box<dyn Any>>,
    io_error: Option<io::Error>,
    unwind_payload: Option<Box<dyn Any + Send + 'static>>,
}

// Safety: unwinding into C is undefined behavior (https://github.com/rust-lang/rust/issues/58794)
// so code outside of the `catch_unwind` call must never panic.
extern "C" fn write_callback<W: io::Write + 'static>(
    env: *mut c_void,
    data: *mut c_uchar,
    length: c_uint,
) -> cairo_status_t {
    // This is consistent with the type of `env` in `Surface::_for_stream`.
    let env: *const CallbackEnvironment = env as _;

    // Safety: the user data entry keeps `Rc<CallbackEnvironment>` alive
    // until the surface is destroyed.
    // If this is called by cairo, the surface is still alive.
    let env: &CallbackEnvironment = unsafe { &*env };

    if let Ok(mut mutable) = env.mutable.try_borrow_mut() {
        if let MutableCallbackEnvironment {
            stream: Some(stream),
            // Don’t attempt another write if a previous one errored or panicked:
            io_error: io_error @ None,
            unwind_payload: unwind_payload @ None,
        } = &mut *mutable {
            // Safety: `write_callback<W>` was instanciated in `Surface::_for_stream`
            // with a W parameter consistent with the box that was unsized to `Box<dyn Any>`.
            let stream = unsafe {
                stream.downcast_mut_unchecked::<W>()
            };
            // Safety: this is the callback contract from cairo’s API
            let data = unsafe {
                std::slice::from_raw_parts(data, length as usize)
            };
            // Because `<W as Write>::write_all` is a generic,
            // we must conservatively assume that it can panic.
            let result = std::panic::catch_unwind(AssertUnwindSafe(|| stream.write_all(data)));
            match result {
                Ok(Ok(())) => {
                    return Status::Success.into()
                }
                Ok(Err(error)) => {
                    *io_error = Some(error);
                }
                Err(payload) => {
                    *unwind_payload = Some(payload);
                }
            }
        }
    } else {
        env.saw_already_borrowed.set(true)
    }
    Status::WriteError.into()
}

struct RawStream<W>(*mut W);

impl<W: io::Write> io::Write for RawStream<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { unsafe { (*self.0).write(buf) }}
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { unsafe { (*self.0).write_all(buf) }}
    fn flush(&mut self) -> io::Result<()> { unsafe { (*self.0).flush() } }
}

trait AnyExt {
    /// Any::downcast_mut, but YOLO
    unsafe fn downcast_mut_unchecked<T>(&mut self) -> &mut T {
        let ptr = self as *mut Self as *mut T;
        &mut *ptr
    }
}
impl AnyExt for dyn Any {}
<|endoftext|>


<commit_before>use std::sync::RwLock;
use std::path::{Path, PathBuf};
use std::fs;

use slog::Logger;
use kite_rocksdb::RocksDBIndexStore;

use index::Index;
use index::registry::IndexRegistry;

pub struct System {
    pub log: Logger,
    data_dir: PathBuf,
    pub indices: RwLock<IndexRegistry>,
}

impl System {
    pub fn new(log: Logger, data_dir: PathBuf) -> System {
        System {
            log: log,
            data_dir: data_dir,
            indices: RwLock::new(IndexRegistry::new()),
        }
    }

    pub fn get_indices_dir(&self) -> PathBuf {
        let mut dir = self.data_dir.clone();
        dir.push("indices");
        dir
    }

    fn load_index(&self, name: String, path: &Path) -> Index {
        Index::new(name, RocksDBIndexStore::open(path).unwrap())
    }

    pub fn load_indices(&self) {
        let indices_dir = self.get_indices_dir();
        match fs::read_dir(indices_dir.clone()) {
            Ok(files) => {
                for file in files {
                    let path = file.unwrap().path();
                    let index_name: String = path.file_stem().unwrap().to_str().unwrap().to_owned();

                    if let Some(ext) = path.extension() {
                        if ext.to_str() == Some("rsi") {
                            self.log.info("[sys] loaded index", b!("index" => index_name));

                            let mut indices_w = self.indices.write().unwrap();
                            let index_ref = indices_w.insert(self.load_index(index_name.clone().to_owned(), path.as_path()));
                            indices_w.aliases.insert(index_name, index_ref);
                        }
                    }
                }
            }
            Err(error) => {
                self.log.warn("[sys] cannot open indices directory", b!(
                    "dir" => indices_dir.to_str().unwrap(),
                    "error" => format!("{}", error)
                ));
            }
        }
    }
}
<commit_msg>Display index open errors<commit_after>use std::sync::RwLock;
use std::path::{Path, PathBuf};
use std::fs;

use slog::Logger;
use kite_rocksdb::RocksDBIndexStore;

use index::Index;
use index::registry::IndexRegistry;

pub struct System {
    pub log: Logger,
    data_dir: PathBuf,
    pub indices: RwLock<IndexRegistry>,
}

impl System {
    pub fn new(log: Logger, data_dir: PathBuf) -> System {
        System {
            log: log,
            data_dir: data_dir,
            indices: RwLock::new(IndexRegistry::new()),
        }
    }

    pub fn get_indices_dir(&self) -> PathBuf {
        let mut dir = self.data_dir.clone();
        dir.push("indices");
        dir
    }

    fn load_index(&self, name: String, path: &Path) -> Result<Index, String> {
        let store = try!(RocksDBIndexStore::open(path));

        Ok(Index::new(name, store))
    }

    pub fn load_indices(&self) {
        let indices_dir = self.get_indices_dir();
        match fs::read_dir(indices_dir.clone()) {
            Ok(files) => {
                for file in files {
                    let path = file.unwrap().path();
                    let index_name: String = path.file_stem().unwrap().to_str().unwrap().to_owned();

                    if let Some(ext) = path.extension() {
                        if ext.to_str() == Some("rsi") {
                            self.log.info("[sys] loaded index", b!("index" => index_name));

                            match self.load_index(index_name.clone().to_owned(), path.as_path()) {
                                Ok(index) => {
                                    let mut indices_w = self.indices.write().unwrap();
                                    let index_ref = indices_w.insert(index);
                                    indices_w.aliases.insert(index_name, index_ref);
                                }
                                Err(e) => {
                                    self.log.error("[sys] failed to open index", b!(
                                        "index" => index_name,
                                        "error" => e
                                    ));
                                }
                            }
                        }
                    }
                }
            }
            Err(error) => {
                self.log.error("[sys] cannot open indices directory", b!(
                    "dir" => indices_dir.to_str().unwrap(),
                    "error" => format!("{}", error)
                ));
            }
        }
    }
}
<|endoftext|>


<commit_before>use bytes::Bytes;
use bytes::BytesMut;

use crate::error;
use crate::result;
use crate::solicit::frame::headers::HeadersFlag;
use crate::solicit::frame::push_promise::PushPromiseFlag;
use crate::solicit::frame::unpack_header_from_slice;
use crate::solicit::frame::HeadersFrame;
use crate::solicit::frame::HttpFrame;
use crate::solicit::frame::HttpFrameType;
use crate::solicit::frame::PushPromiseFrame;
use crate::solicit::frame::RawFrame;
use crate::solicit::frame::FRAME_HEADER_LEN;
use crate::solicit::stream_id::StreamId;
use crate::ErrorCode;
use futures::task::Context;
use std::pin::Pin;
use std::task::Poll;
use tokio::io::AsyncRead;

/// Buffered read for reading HTTP/2 frames.
pub struct HttpFramedRead<R: AsyncRead + Unpin> {
    read: R,
    buf: BytesMut,
}

impl<R: AsyncRead + Unpin> HttpFramedRead<R> {
    pub fn new(read: R) -> HttpFramedRead<R> {
        HttpFramedRead {
            read,
            buf: BytesMut::new(),
        }
    }

    fn fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<result::Result<()>> {
        let mut_self = self.get_mut();
        mut_self.buf.reserve(8192);
        let n = match Pin::new(&mut mut_self.read).poll_read_buf(cx, &mut mut_self.buf)? {
            Poll::Ready(n) => n,
            Poll::Pending => return Poll::Pending,
        };
        if n == 0 {
            return Poll::Ready(Err(error::Error::EofFromStream));
        }
        Poll::Ready(Ok(()))
    }

    fn fill_buff_to_at_least(
        &mut self,
        cx: &mut Context<'_>,
        at_least: usize,
    ) -> Poll<result::Result<()>> {
        while self.buf.len() < at_least {
            if let Poll::Pending = Pin::new(&mut *self).fill_buf(cx)? {
                return Poll::Pending;
            }
        }
        Poll::Ready(Ok(()))
    }

    fn poll_raw_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<RawFrame>> {
        if let Poll::Pending = self.fill_buff_to_at_least(cx, FRAME_HEADER_LEN)? {
            return Poll::Pending;
        }

        let header = {
            let header = &self.buf[..FRAME_HEADER_LEN];
            unpack_header_from_slice(header)
        };

        if header.payload_len > max_frame_size {
            warn!(
                "closing conn because peer sent frame with size: {}, max_frame_size: {}",
                header.payload_len, max_frame_size
            );
            return Poll::Ready(Err(error::Error::CodeError(ErrorCode::FrameSizeError)));
        }

        let total_len = FRAME_HEADER_LEN + header.payload_len as usize;

        if let Poll::Pending = self.fill_buff_to_at_least(cx, total_len)? {
            return Poll::Pending;
        }

        Poll::Ready(Ok(RawFrame {
            raw_content: self.buf.split_to(total_len).freeze(),
        }))
    }

    fn poll_http_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<HttpFrame>> {
        match self.poll_raw_frame(cx, max_frame_size)? {
            Poll::Ready(frame) => Poll::Ready(Ok(HttpFrame::from_raw(&frame)?)),
            Poll::Pending => Poll::Pending,
        }
    }
}

enum Continuable {
    Headers(HeadersFrame),
    PushPromise(PushPromiseFrame),
}

impl Continuable {
    fn into_frame(self) -> HttpFrame {
        match self {
            Continuable::Headers(headers) => HttpFrame::Headers(headers),
            Continuable::PushPromise(push_promise) => HttpFrame::PushPromise(push_promise),
        }
    }

    fn extend_header_fragment(&mut self, bytes: Bytes) {
        let header_fragment = match self {
            &mut Continuable::Headers(ref mut headers) => &mut headers.header_fragment,
            &mut Continuable::PushPromise(ref mut push_promise) => {
                &mut push_promise.header_fragment
            }
        };

        // TODO: quadratic
        let mut bm = BytesMut::new();
        bm.extend_from_slice(&header_fragment[..]);
        bm.extend_from_slice(&bytes[..]);
        *header_fragment = bm.freeze();
    }

    fn set_end_headers(&mut self) {
        match self {
            &mut Continuable::Headers(ref mut headers) => {
                headers.flags.set(HeadersFlag::EndHeaders)
            }
            &mut Continuable::PushPromise(ref mut push_promise) => {
                push_promise.flags.set(PushPromiseFlag::EndHeaders)
            }
        }
    }

    fn get_stream_id(&self) -> StreamId {
        match self {
            &Continuable::Headers(ref headers) => headers.stream_id,
            &Continuable::PushPromise(ref push_promise) => push_promise.stream_id,
        }
    }
}

pub struct HttpFramedJoinContinuationRead<R: AsyncRead + Unpin> {
    framed_read: HttpFramedRead<R>,
    // TODO: check total size is not exceeded some limit
    header_opt: Option<Continuable>,
}

impl<R: AsyncRead + Unpin> HttpFramedJoinContinuationRead<R> {
    pub fn new(read: R) -> Self {
        HttpFramedJoinContinuationRead {
            framed_read: HttpFramedRead::new(read),
            header_opt: None,
        }
    }

    pub fn poll_http_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<HttpFrame>> {
        loop {
            let frame = match self.framed_read.poll_http_frame(cx, max_frame_size)? {
                Poll::Pending => return Poll::Pending,
                Poll::Ready(frame) => frame,
            };

            match frame {
                HttpFrame::Headers(h) => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            HttpFrameType::Headers,
                        )));
                    } else {
                        if h.flags.is_set(HeadersFlag::EndHeaders) {
                            return Poll::Ready(Ok(HttpFrame::Headers(h)));
                        } else {
                            self.header_opt = Some(Continuable::Headers(h));
                            continue;
                        }
                    }
                }
                HttpFrame::PushPromise(p) => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            HttpFrameType::PushPromise,
                        )));
                    } else {
                        if p.flags.is_set(PushPromiseFlag::EndHeaders) {
                            return Poll::Ready(Ok(HttpFrame::PushPromise(p)));
                        } else {
                            self.header_opt = Some(Continuable::PushPromise(p));
                            continue;
                        }
                    }
                }
                HttpFrame::Continuation(c) => {
                    if let Some(mut h) = self.header_opt.take() {
                        if h.get_stream_id() != c.stream_id {
                            return Poll::Ready(Err(
                                error::Error::ExpectingContinuationGotDifferentStreamId(
                                    h.get_stream_id(),
                                    c.stream_id,
                                ),
                            ));
                        } else {
                            let header_end = c.is_headers_end();
                            h.extend_header_fragment(c.header_fragment);
                            if header_end {
                                h.set_end_headers();
                                return Poll::Ready(Ok(h.into_frame()));
                            } else {
                                self.header_opt = Some(h);
                                continue;
                            }
                        }
                    } else {
                        return Poll::Ready(Err(error::Error::ContinuationFrameWithoutHeaders));
                    }
                }
                f => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            f.frame_type(),
                        )));
                    } else {
                        return Poll::Ready(Ok(f));
                    }
                }
            };
        }
    }
}
<commit_msg>Fix quadradic continuation frame handling<commit_after>use bytes::Bytes;
use bytes::BytesMut;

use crate::error;
use crate::result;
use crate::solicit::frame::headers::HeadersFlag;
use crate::solicit::frame::push_promise::PushPromiseFlag;
use crate::solicit::frame::unpack_header_from_slice;
use crate::solicit::frame::HeadersFrame;
use crate::solicit::frame::HttpFrame;
use crate::solicit::frame::HttpFrameType;
use crate::solicit::frame::PushPromiseFrame;
use crate::solicit::frame::RawFrame;
use crate::solicit::frame::FRAME_HEADER_LEN;
use crate::solicit::stream_id::StreamId;
use crate::ErrorCode;
use futures::task::Context;
use std::pin::Pin;
use std::task::Poll;
use tokio::io::AsyncRead;

/// Buffered read for reading HTTP/2 frames.
pub struct HttpFramedRead<R: AsyncRead + Unpin> {
    read: R,
    buf: BytesMut,
}

impl<R: AsyncRead + Unpin> HttpFramedRead<R> {
    pub fn new(read: R) -> HttpFramedRead<R> {
        HttpFramedRead {
            read,
            buf: BytesMut::new(),
        }
    }

    fn fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<result::Result<()>> {
        let mut_self = self.get_mut();
        mut_self.buf.reserve(8192);
        let n = match Pin::new(&mut mut_self.read).poll_read_buf(cx, &mut mut_self.buf)? {
            Poll::Ready(n) => n,
            Poll::Pending => return Poll::Pending,
        };
        if n == 0 {
            return Poll::Ready(Err(error::Error::EofFromStream));
        }
        Poll::Ready(Ok(()))
    }

    fn fill_buff_to_at_least(
        &mut self,
        cx: &mut Context<'_>,
        at_least: usize,
    ) -> Poll<result::Result<()>> {
        while self.buf.len() < at_least {
            if let Poll::Pending = Pin::new(&mut *self).fill_buf(cx)? {
                return Poll::Pending;
            }
        }
        Poll::Ready(Ok(()))
    }

    fn poll_raw_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<RawFrame>> {
        if let Poll::Pending = self.fill_buff_to_at_least(cx, FRAME_HEADER_LEN)? {
            return Poll::Pending;
        }

        let header = {
            let header = &self.buf[..FRAME_HEADER_LEN];
            unpack_header_from_slice(header)
        };

        if header.payload_len > max_frame_size {
            warn!(
                "closing conn because peer sent frame with size: {}, max_frame_size: {}",
                header.payload_len, max_frame_size
            );
            return Poll::Ready(Err(error::Error::CodeError(ErrorCode::FrameSizeError)));
        }

        let total_len = FRAME_HEADER_LEN + header.payload_len as usize;

        if let Poll::Pending = self.fill_buff_to_at_least(cx, total_len)? {
            return Poll::Pending;
        }

        Poll::Ready(Ok(RawFrame {
            raw_content: self.buf.split_to(total_len).freeze(),
        }))
    }

    fn poll_http_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<HttpFrame>> {
        match self.poll_raw_frame(cx, max_frame_size)? {
            Poll::Ready(frame) => Poll::Ready(Ok(HttpFrame::from_raw(&frame)?)),
            Poll::Pending => Poll::Pending,
        }
    }
}

enum ContinuableFrame {
    Headers(HeadersFrame),
    PushPromise(PushPromiseFrame),
}

struct Continuable {
    header_fragment: BytesMut,
    /// Note frame contatains a header fragment, but it is not used
    frame: ContinuableFrame,
}

impl Continuable {
    fn headers(header: HeadersFrame) -> Continuable {
        Continuable {
            header_fragment: BytesMut::from(&header.header_fragment[..]),
            frame: ContinuableFrame::Headers(header),
        }
    }

    fn push_promise(push_promise: PushPromiseFrame) -> Continuable {
        Continuable {
            header_fragment: BytesMut::from(&push_promise.header_fragment[..]),
            frame: ContinuableFrame::PushPromise(push_promise),
        }
    }

    fn into_frame(mut self) -> HttpFrame {
        let header_fragment = match &mut self.frame {
            ContinuableFrame::Headers(headers) => &mut headers.header_fragment,
            ContinuableFrame::PushPromise(push_promise) => &mut push_promise.header_fragment,
        };
        *header_fragment = self.header_fragment.freeze();
        match self.frame {
            ContinuableFrame::Headers(headers) => HttpFrame::Headers(headers),
            ContinuableFrame::PushPromise(push_promise) => HttpFrame::PushPromise(push_promise),
        }
    }

    fn extend_header_fragment(&mut self, bytes: Bytes) {
        self.header_fragment.extend_from_slice(&bytes[..]);
    }

    fn set_end_headers(&mut self) {
        match self.frame {
            ContinuableFrame::Headers(ref mut headers) => {
                headers.flags.set(HeadersFlag::EndHeaders)
            }
            ContinuableFrame::PushPromise(ref mut push_promise) => {
                push_promise.flags.set(PushPromiseFlag::EndHeaders)
            }
        }
    }

    fn get_stream_id(&self) -> StreamId {
        match self.frame {
            ContinuableFrame::Headers(ref headers) => headers.stream_id,
            ContinuableFrame::PushPromise(ref push_promise) => push_promise.stream_id,
        }
    }
}

pub struct HttpFramedJoinContinuationRead<R: AsyncRead + Unpin> {
    framed_read: HttpFramedRead<R>,
    // TODO: check total size is not exceeded some limit
    header_opt: Option<Continuable>,
}

impl<R: AsyncRead + Unpin> HttpFramedJoinContinuationRead<R> {
    pub fn new(read: R) -> Self {
        HttpFramedJoinContinuationRead {
            framed_read: HttpFramedRead::new(read),
            header_opt: None,
        }
    }

    pub fn poll_http_frame(
        &mut self,
        cx: &mut Context<'_>,
        max_frame_size: u32,
    ) -> Poll<result::Result<HttpFrame>> {
        loop {
            let frame = match self.framed_read.poll_http_frame(cx, max_frame_size)? {
                Poll::Pending => return Poll::Pending,
                Poll::Ready(frame) => frame,
            };

            match frame {
                HttpFrame::Headers(h) => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            HttpFrameType::Headers,
                        )));
                    } else {
                        if h.flags.is_set(HeadersFlag::EndHeaders) {
                            return Poll::Ready(Ok(HttpFrame::Headers(h)));
                        } else {
                            self.header_opt = Some(Continuable::headers(h));
                            continue;
                        }
                    }
                }
                HttpFrame::PushPromise(p) => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            HttpFrameType::PushPromise,
                        )));
                    } else {
                        if p.flags.is_set(PushPromiseFlag::EndHeaders) {
                            return Poll::Ready(Ok(HttpFrame::PushPromise(p)));
                        } else {
                            self.header_opt = Some(Continuable::push_promise(p));
                            continue;
                        }
                    }
                }
                HttpFrame::Continuation(c) => {
                    if let Some(mut h) = self.header_opt.take() {
                        if h.get_stream_id() != c.stream_id {
                            return Poll::Ready(Err(
                                error::Error::ExpectingContinuationGotDifferentStreamId(
                                    h.get_stream_id(),
                                    c.stream_id,
                                ),
                            ));
                        } else {
                            let header_end = c.is_headers_end();
                            h.extend_header_fragment(c.header_fragment);
                            if header_end {
                                h.set_end_headers();
                                return Poll::Ready(Ok(h.into_frame()));
                            } else {
                                self.header_opt = Some(h);
                                continue;
                            }
                        }
                    } else {
                        return Poll::Ready(Err(error::Error::ContinuationFrameWithoutHeaders));
                    }
                }
                f => {
                    if let Some(_) = self.header_opt {
                        return Poll::Ready(Err(error::Error::ExpectingContinuationGot(
                            f.frame_type(),
                        )));
                    } else {
                        return Poll::Ready(Ok(f));
                    }
                }
            };
        }
    }
}
<|endoftext|>


<commit_before>use serde::{Deserialize, Serialize};

pub trait EvolutionState {
    fn update(&mut self);
}

pub trait Cell: Clone + Serialize + Deserialize {
    type Coord: Coord;
    type State: EvolutionState;

    fn update<'a, I>(&'a self, neighbors: I, &Self::State) -> Self
        where I: Iterator<Item = Option<&'a Self>>;

    fn with_coord<C: Coord>(C) -> Self;
    fn coord(&self) -> &Self::Coord;
    fn set_coord<C: Coord>(&mut self, &C);
}

pub trait Nhood {
    type Coord: Coord;

    fn neighbors(&self, &Self::Coord) -> Vec<Self::Coord>;
    fn neighbors_count(&self) -> usize;
}

pub trait Coord: Clone + Serialize + Deserialize {
    fn from_2d(x: i32, y: i32) -> Self;

    fn x(&self) -> i32;
    fn y(&self) -> i32;
    fn z(&self) -> i32 { 0 }
}

pub trait Grid {
    type Cell: Cell;
    type Coord: Coord;

    fn update(&mut self);

    fn state(&self) -> &<<Self as Grid>::Cell as Cell>::State;
    fn cells(&self) -> &Vec<Self::Cell>;
    fn dimensions(&self) -> Self::Coord;

    fn set_cells(&mut self, Vec<Self::Cell>);
}

pub trait Consumer {
    type Cell: Cell;

    fn consume<G>(&mut self, &mut G)
        where G: Grid<Cell = Self::Cell>;
}

pub trait Engine {
    fn run_times(&mut self, i64);
}
<commit_msg>Cell is mutable.<commit_after>use serde::{Deserialize, Serialize};

pub trait EvolutionState {
    fn update(&mut self);
}

pub trait Cell: Clone + Serialize + Deserialize {
    type Coord: Coord;
    type State: EvolutionState;

    fn update<'a, I>(&'a mut self, old: &'a Self, neighbors: I, &Self::State)
        where I: Iterator<Item = Option<&'a Self>>;

    fn with_coord<C: Coord>(C) -> Self;
    fn coord(&self) -> &Self::Coord;
    fn set_coord<C: Coord>(&mut self, &C);
}

pub trait Nhood {
    type Coord: Coord;

    fn neighbors(&self, &Self::Coord) -> Vec<Self::Coord>;
    fn neighbors_count(&self) -> usize;
}

pub trait Coord: Clone + Serialize + Deserialize {
    fn from_2d(x: i32, y: i32) -> Self;

    fn x(&self) -> i32;
    fn y(&self) -> i32;
    fn z(&self) -> i32 { 0 }
}

pub trait Grid {
    type Cell: Cell;
    type Coord: Coord;

    fn update(&mut self);

    fn state(&self) -> &<<Self as Grid>::Cell as Cell>::State;
    fn cells(&self) -> &Vec<Self::Cell>;
    fn dimensions(&self) -> Self::Coord;

    fn set_cells(&mut self, Vec<Self::Cell>);
}

pub trait Consumer {
    type Cell: Cell;

    fn consume<G>(&mut self, &mut G)
        where G: Grid<Cell = Self::Cell>;
}

pub trait Engine {
    fn run_times(&mut self, i64);
}
<|endoftext|>


<commit_before>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) 2015, Nathan Zadoks <nathan@nathan7.eu>
// See the LICENSE file included in this distribution.
extern crate libc;
use core::prelude::*;
use self::libc::{c_void, c_int, size_t};
use self::libc::{mmap, mprotect, munmap};
use self::libc::MAP_FAILED;
use super::page_size;

use core::ptr;

#[cold]
#[inline(never)]
pub fn sys_page_size() -> usize {
  unsafe {
    libc::sysconf(libc::_SC_PAGESIZE) as usize
  }
}

const GUARD_PROT:  c_int = libc::PROT_NONE;
const STACK_PROT:  c_int = libc::PROT_READ
                         | libc::PROT_WRITE;
#[cfg(not(any(target_os = "freebsd", target_os = "dragonfly")))]
const STACK_FLAGS: c_int = libc::MAP_STACK
                         | libc::MAP_PRIVATE
                         | libc::MAP_ANON;
// workaround for http://lists.freebsd.org/pipermail/freebsd-bugs/2011-July/044840.html
// according to libgreen, DragonFlyBSD suffers from this too
#[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
const STACK_FLAGS: c_int = libc::MAP_PRIVATE
                         | libc::MAP_ANON;

pub unsafe fn map_stack(len: usize) -> Option<*mut u8> {
  let ptr = mmap(ptr::null_mut(), len as size_t,
                 STACK_PROT, STACK_FLAGS, -1, 0);
  if ptr != MAP_FAILED {
    Some(ptr as *mut u8)
  }
  else {
    None
  }
}

pub unsafe fn protect_stack(ptr: *mut u8) -> bool {
  mprotect(ptr as *mut c_void, page_size() as size_t, GUARD_PROT) == 0
}

pub unsafe fn unmap_stack(ptr: *mut u8, len: usize) -> bool {
  munmap(ptr as *mut c_void, len as size_t) == 0
}
<commit_msg>it's fine to inline a simple FFI call<commit_after>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) 2015, Nathan Zadoks <nathan@nathan7.eu>
// See the LICENSE file included in this distribution.
extern crate libc;
use core::prelude::*;
use self::libc::{c_void, c_int, size_t};
use self::libc::{mmap, mprotect, munmap};
use self::libc::MAP_FAILED;
use super::page_size;

use core::ptr;

#[cold]
pub fn sys_page_size() -> usize {
  unsafe {
    libc::sysconf(libc::_SC_PAGESIZE) as usize
  }
}

const GUARD_PROT:  c_int = libc::PROT_NONE;
const STACK_PROT:  c_int = libc::PROT_READ
                         | libc::PROT_WRITE;
#[cfg(not(any(target_os = "freebsd", target_os = "dragonfly")))]
const STACK_FLAGS: c_int = libc::MAP_STACK
                         | libc::MAP_PRIVATE
                         | libc::MAP_ANON;
// workaround for http://lists.freebsd.org/pipermail/freebsd-bugs/2011-July/044840.html
// according to libgreen, DragonFlyBSD suffers from this too
#[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
const STACK_FLAGS: c_int = libc::MAP_PRIVATE
                         | libc::MAP_ANON;

pub unsafe fn map_stack(len: usize) -> Option<*mut u8> {
  let ptr = mmap(ptr::null_mut(), len as size_t,
                 STACK_PROT, STACK_FLAGS, -1, 0);
  if ptr != MAP_FAILED {
    Some(ptr as *mut u8)
  }
  else {
    None
  }
}

pub unsafe fn protect_stack(ptr: *mut u8) -> bool {
  mprotect(ptr as *mut c_void, page_size() as size_t, GUARD_PROT) == 0
}

pub unsafe fn unmap_stack(ptr: *mut u8, len: usize) -> bool {
  munmap(ptr as *mut c_void, len as size_t) == 0
}
<|endoftext|>


<commit_before>//Following semver here
#[derive(Clone, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub major: u16,
    pub minor: u16,
    pub patch: u32
}

#[derive(Clone, RustcDecodable, RustcEncodable)]
pub struct Package {
    pub name: String,
    pub version: Version,
    pub crc32: u32,
    pub child_url: String, // Examples: /some_dir/something.exe
    pub dependencies: Vec<Package>
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct PackageList {
    packages: Vec<Package>,
    parent_url: String, // Examples: https://www.foo.bar
    version: u32
}

impl PartialEq for Version {
    fn eq(&self, other: &Version) -> bool {
        self.major == other.major &&
        self.minor == other.minor &&
        self.patch == other.patch
    }
}

impl PartialEq for Package {
    fn eq(&self, other: &Package) -> bool {
        self.name == other.name &&
        self.version == other.version &&
        self.crc32 == other.crc32 &&
        self.dependencies == other.dependencies
    }
}<commit_msg>Make variables of PackageList public<commit_after>//Following semver here
#[derive(Clone, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub major: u16,
    pub minor: u16,
    pub patch: u32
}

#[derive(Clone, RustcDecodable, RustcEncodable)]
pub struct Package {
    pub name: String,
    pub version: Version,
    pub crc32: u32,
    pub child_url: String, // Examples: /some_dir/something.exe
    pub dependencies: Vec<Package>
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct PackageList {
    pub packages: Vec<Package>,
    pub parent_url: String, // Examples: https://www.foo.bar
    pub version: u32
}

impl PartialEq for Version {
    fn eq(&self, other: &Version) -> bool {
        self.major == other.major &&
        self.minor == other.minor &&
        self.patch == other.patch
    }
}

impl PartialEq for Package {
    fn eq(&self, other: &Package) -> bool {
        self.name == other.name &&
        self.version == other.version &&
        self.crc32 == other.crc32 &&
        self.dependencies == other.dependencies
    }
}<|endoftext|>


<commit_before>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [u32;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as u32) | ((value as u32) << 4);
        }
      }
    }

    if use_bigtable {
      self.initialize_bigtable();
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val = self.numbits[code as usize];
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    //First attempt to do complete decode, by using the first 14 bits
    let code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    let len = try!(self.huff_len(pump));
    Ok(self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<u32,String> {
    let mut code = pump.peek_bits(8) as usize;
    let val = self.numbits[code as usize] as u32;
    let mut l = val & 15;
    if l != 0 {
      pump.consume_bits(l);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    l = 8;
    while code as i32 > self.maxcode[l as usize] {
      let temp = pump.get_bits(1) as usize;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l as usize] as usize +
        (code - (self.mincode[l as usize] as usize)) as usize
      ]);
    }
  }

  pub fn huff_diff(&self, pump: &mut BitPump, len: u32) -> i32 {
    match len {
      0 => 0,
      16 => {
        if self.dng_compatible {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // Section F.2.2.1: decode the difference and
        // Figure F.12: extend sign bit
        let mut x: i32 = pump.get_bits(len as u32) as i32;
        if (x & (1 << (len - 1))) == 0 {
          x -= (1 << len) - 1;
        }
        x
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<commit_msg>ljpeg: remove more casts<commit_after>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [u32;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as u32) | ((value as u32) << 4);
        }
      }
    }

    if use_bigtable {
      self.initialize_bigtable();
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val = self.numbits[code as usize];
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    //First attempt to do complete decode, by using the first 14 bits
    let code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    let len = try!(self.huff_len(pump));
    Ok(self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<u32,String> {
    let mut code = pump.peek_bits(8) as usize;
    let val = self.numbits[code as usize] as u32;
    let len = val & 15;
    if len != 0 {
      pump.consume_bits(len);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    let mut l: usize = 8;
    while code as i32 > self.maxcode[l] {
      let temp = pump.get_bits(1) as usize;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision || self.valptr[l] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l] as usize +
        (code - (self.mincode[l] as usize)) as usize
      ]);
    }
  }

  pub fn huff_diff(&self, pump: &mut BitPump, len: u32) -> i32 {
    match len {
      0 => 0,
      16 => {
        if self.dng_compatible {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // Section F.2.2.1: decode the difference and
        // Figure F.12: extend sign bit
        let mut x: i32 = pump.get_bits(len as u32) as i32;
        if (x & (1 << (len - 1))) == 0 {
          x -= (1 << len) - 1;
        }
        x
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<|endoftext|>


<commit_before>use std::fmt;
use crate::decoders::basics::*;

const DECODE_TABLE_BITS: u32 = 11;

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // Represent the weird shifts that are needed for some NEF files
  pub shiftval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  nbits: u32,
  hufftable: Vec<(u16,u16,u16)>,
  decodetable: [Option<(u16,i32)>; 1<< DECODE_TABLE_BITS],
  pub dng_bug: bool,
  pub initialized: bool,
}

struct MockPump {
  bits: u64,
  nbits: u32,
}

impl MockPump {
  pub fn empty() -> Self {
    MockPump {
      bits: 0,
      nbits: 0,
    }
  }

  pub fn set(&mut self, bits: u32, nbits: u32) {
    self.bits = (bits as u64) << 32;
    self.nbits = nbits + 32;
  }

  pub fn validbits(&self) -> i32 {
    self.nbits as i32 - 32
  }
}

impl BitPump for MockPump {
  fn peek_bits(&mut self, num: u32) -> u32 {
    (self.bits >> (self.nbits-num)) as u32
  }

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

impl HuffTable {
  pub fn empty() -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      shiftval: [0;256],
      nbits: 0,
      hufftable: Vec::new(),
      decodetable: [None; 1 << DECODE_TABLE_BITS],
      dng_bug: false,
      initialized: false,
    }
  }

  pub fn new(bits: [u32;17], huffval: [u32;256], dng_bug: bool) -> Result<HuffTable,String> {
    let mut tbl = HuffTable {
      bits: bits,
      huffval: huffval,
      shiftval: [0;256],
      nbits: 0,
      hufftable: Vec::new(),
      decodetable: [None; 1 << DECODE_TABLE_BITS],
      dng_bug: dng_bug,
      initialized: false,
    };
    tbl.initialize()?;
    Ok(tbl)
  }

  pub fn initialize(&mut self) -> Result<(), String> {
    // Create the decoding table for the huffman lengths
    self.nbits = 16;
    for i in 0..16 {
      if self.bits[16-i] != 0 {
        break;
      }
      self.nbits -= 1;
    }

    let tblsize = 1 << self.nbits;
    self.hufftable = vec![(0,0,0); tblsize];

    let mut h = 0;
    let mut pos = 0;
    for len in 0..self.nbits {
      for _ in 0..self.bits[len as usize + 1] {
        for _ in 0..(1 << (self.nbits-len-1)) {
          self.hufftable[h] = (len as u16 + 1, self.huffval[pos] as u16, self.shiftval[pos] as u16);
          h += 1;
        }
        pos += 1;
      }
    }

    // Now bootstrap the full decode table
    let mut pump = MockPump::empty();
    let mut i = 0;
    loop {
      pump.set(i, DECODE_TABLE_BITS);
      let decode = self.huff_decode_slow(&mut pump);
      if pump.validbits() >= 0 {
        self.decodetable[i as usize] = Some(decode);
      }
      i += 1;
      if i >= 1 << DECODE_TABLE_BITS {
        break;
      }
    }

    self.initialized = true;
    Ok(())
  }

  pub fn huff_decode(&self, pump: &mut dyn BitPump) -> Result<i32,String> {
    let code = pump.peek_bits(DECODE_TABLE_BITS) as usize;
    if let Some((bits,decode)) = self.decodetable[code] {
      pump.consume_bits(bits as u32);
      Ok(decode)
    } else {
      let decode = self.huff_decode_slow(pump);
      Ok(decode.1)
    }
  }

  pub fn huff_decode_slow(&self, pump: &mut dyn BitPump) -> (u16,i32) {
    let len = self.huff_len(pump);
    (len.0+len.1, self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut dyn BitPump) -> (u16,u16,u16) {
    let code = pump.peek_bits(self.nbits) as usize;
    let (bits, len, shift) = self.hufftable[code];
    pump.consume_bits(bits as u32);
    (bits, len, shift)
  }

  pub fn huff_diff(&self, pump: &mut dyn BitPump, input: (u16,u16,u16)) -> i32 {
    let (_, len, shift) = input;

    match len {
      0 => 0,
      16 => {
        if self.dng_bug {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // decode the difference and extend sign bit
        let fulllen: i32 = len as i32 + shift as i32;
        let shift: i32 = shift as i32;
        let bits = pump.get_bits(len as u32) as i32;
        let mut diff: i32 = ((bits << 1) + 1) << shift >> 1;
        if (diff & (1 << (fulllen - 1))) == 0 {
          diff -= (1 << fulllen) - ((shift == 0) as i32);
        }
        diff
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}
<commit_msg>Cleanup and document the huffman code<commit_after>use std::fmt;
use crate::decoders::basics::*;

const DECODE_CACHE_BITS: u32 = 11;

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // Represent the weird shifts that are needed for some NEF files
  pub shiftval: [u32;256],

  // Enable the workaround for 16 bit decodes in DNG that need to consume those
  // bits instead of the value being implied
  pub dng_bug: bool,

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private

  // The max number of bits in a huffman code and the table that converts those
  // bits into how many bits to consume and the decoded length and shift
  nbits: u32,
  hufftable: Vec<(u16,u16,u16)>,

  // A pregenerated table that goes straight to decoding a diff without first
  // finding a length, fetching bits, and sign extending them. The table is
  // sized by DECODE_CACHE_BITS and can have 97%+ hit rate with 11 bits
  decodecache: [Option<(u16,i32)>; 1<< DECODE_CACHE_BITS],

  initialized: bool,
}

struct MockPump {
  bits: u64,
  nbits: u32,
}

impl MockPump {
  pub fn empty() -> Self {
    MockPump {
      bits: 0,
      nbits: 0,
    }
  }

  pub fn set(&mut self, bits: u32, nbits: u32) {
    self.bits = (bits as u64) << 32;
    self.nbits = nbits + 32;
  }

  pub fn validbits(&self) -> i32 {
    self.nbits as i32 - 32
  }
}

impl BitPump for MockPump {
  fn peek_bits(&mut self, num: u32) -> u32 {
    (self.bits >> (self.nbits-num)) as u32
  }

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

impl HuffTable {
  pub fn empty() -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      shiftval: [0;256],
      dng_bug: false,

      nbits: 0,
      hufftable: Vec::new(),
      decodecache: [None; 1 << DECODE_CACHE_BITS],
      initialized: false,
    }
  }

  pub fn new(bits: [u32;17], huffval: [u32;256], dng_bug: bool) -> Result<HuffTable,String> {
    let mut tbl = HuffTable {
      bits: bits,
      huffval: huffval,
      shiftval: [0;256],
      dng_bug: dng_bug,

      nbits: 0,
      hufftable: Vec::new(),
      decodecache: [None; 1 << DECODE_CACHE_BITS],
      initialized: false,
    };
    tbl.initialize()?;
    Ok(tbl)
  }

  pub fn initialize(&mut self) -> Result<(), String> {
    // Find out the max code length and allocate a table with that size
    self.nbits = 16;
    for i in 0..16 {
      if self.bits[16-i] != 0 {
        break;
      }
      self.nbits -= 1;
    }
    self.hufftable = vec![(0,0,0); 1 << self.nbits];

    // Fill in the table itself
    let mut h = 0;
    let mut pos = 0;
    for len in 0..self.nbits {
      for _ in 0..self.bits[len as usize + 1] {
        for _ in 0..(1 << (self.nbits-len-1)) {
          self.hufftable[h] = (len as u16 + 1, self.huffval[pos] as u16, self.shiftval[pos] as u16);
          h += 1;
        }
        pos += 1;
      }
    }

    // Create the decode cache by running the slow code over all the possible
    // values DECODE_CACHE_BITS wide
    let mut pump = MockPump::empty();
    let mut i = 0;
    loop {
      pump.set(i, DECODE_CACHE_BITS);
      let decode = self.huff_decode_slow(&mut pump);
      if pump.validbits() >= 0 {
        self.decodecache[i as usize] = Some(decode);
      }
      i += 1;
      if i >= 1 << DECODE_CACHE_BITS {
        break;
      }
    }

    self.initialized = true;
    Ok(())
  }

  pub fn huff_decode(&self, pump: &mut dyn BitPump) -> Result<i32,String> {
    let code = pump.peek_bits(DECODE_CACHE_BITS) as usize;
    if let Some((bits,decode)) = self.decodecache[code] {
      pump.consume_bits(bits as u32);
      Ok(decode)
    } else {
      let decode = self.huff_decode_slow(pump);
      Ok(decode.1)
    }
  }

  pub fn huff_decode_slow(&self, pump: &mut dyn BitPump) -> (u16,i32) {
    let len = self.huff_len(pump);
    (len.0+len.1, self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut dyn BitPump) -> (u16,u16,u16) {
    let code = pump.peek_bits(self.nbits) as usize;
    let (bits, len, shift) = self.hufftable[code];
    pump.consume_bits(bits as u32);
    (bits, len, shift)
  }

  pub fn huff_diff(&self, pump: &mut dyn BitPump, input: (u16,u16,u16)) -> i32 {
    let (_, len, shift) = input;

    match len {
      0 => 0,
      16 => {
        if self.dng_bug {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // decode the difference and extend sign bit
        let fulllen: i32 = len as i32 + shift as i32;
        let shift: i32 = shift as i32;
        let bits = pump.get_bits(len as u32) as i32;
        let mut diff: i32 = ((bits << 1) + 1) << shift >> 1;
        if (diff & (1 << (fulllen - 1))) == 0 {
          diff -= (1 << fulllen) - ((shift == 0) as i32);
        }
        diff
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}
<|endoftext|>