<commit_before>use std::borrow::Cow;
use std::io::Write;

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    pub fn encode(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        let config = self.config(color);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType) -> Config {
        Config {
            quality: 100,
            alpha_quality: 100,
            speed: 1,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<commit_msg>Add some small description to missing docs<commit_after>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn encode(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        let config = self.config(color);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType) -> Config {
        Config {
            quality: 100,
            alpha_quality: 100,
            speed: 1,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

extern crate rustc_middle;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;

use std::convert::TryFrom;
use std::env;
use std::str::FromStr;

use hex::FromHexError;
use log::debug;

use rustc_session::CtfeBacktrace;
use rustc_driver::Compilation;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_interface::{interface, Queries};
use rustc_middle::ty::TyCtxt;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &interface::Compiler,
        queries: &'tcx Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();

        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            init_late_loggers(tcx);
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            if let Some(return_code) = miri::eval_main(tcx, entry_def_id, config) {
                std::process::exit(
                    i32::try_from(return_code).expect("Return value was too large!"),
                );
            }
        });

        compiler.session().abort_if_errors();

        Compilation::Stop
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var("RUSTC_LOG").is_ok() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers(tcx: TyCtxt<'_>) {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var("RUSTC_LOG").is_err() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var(
                    "RUSTC_LOG",
                    &format!("rustc_middle::mir::interpret={0},rustc_mir::interpret={0}", var),
                );
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUSTC_CTFE_BACKTRACE` is not, set `RUSTC_CTFE_BACKTRACE`.
    // Do this late, so we ideally only apply this to Miri's errors.
    if let Ok(val) = env::var("MIRI_BACKTRACE") {
        let ctfe_backtrace = match &*val {
            "immediate" => CtfeBacktrace::Immediate,
            "0" => CtfeBacktrace::Disabled,
            _ => CtfeBacktrace::Capture,
        };
        *tcx.sess.ctfe_backtrace.borrow_mut() = ctfe_backtrace;
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc would
    // end up somewhere in the build dir.
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
            .to_owned(),
    })
}

fn main() {
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut stacked_borrows = true;
    let mut check_alignment = true;
    let mut communicate = false;
    let mut ignore_leaks = false;
    let mut seed: Option<u64> = None;
    let mut tracked_pointer_tag: Option<miri::PtrId> = None;
    let mut tracked_alloc_id: Option<miri::AllocId> = None;
    let mut rustc_args = vec![];
    let mut miri_args = vec![];
    let mut after_dashdash = false;
    let mut excluded_env_vars = vec![];
    for arg in std::env::args() {
        if rustc_args.is_empty() {
            // Very first arg: for `rustc`.
            rustc_args.push(arg);
        } else if after_dashdash {
            // Everything that comes after are `miri` args.
            miri_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                }
                "-Zmiri-disable-stacked-borrows" => {
                    stacked_borrows = false;
                }
                "-Zmiri-disable-alignment-check" => {
                    check_alignment = false;
                }
                "-Zmiri-disable-isolation" => {
                    communicate = true;
                }
                "-Zmiri-ignore-leaks" => {
                    ignore_leaks = true;
                }
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.trim_start_matches("-Zmiri-seed="))
                        .unwrap_or_else(|err| match err {
                            FromHexError::InvalidHexCharacter { .. } => panic!(
                                "-Zmiri-seed should only contain valid hex digits [0-9a-fA-F]"
                            ),
                            FromHexError::OddLength =>
                                panic!("-Zmiri-seed should have an even number of digits"),
                            err => panic!("Unknown error decoding -Zmiri-seed as hex: {:?}", err),
                        });
                    if seed_raw.len() > 8 {
                        panic!(format!(
                            "-Zmiri-seed must be at most 8 bytes, was {}",
                            seed_raw.len()
                        ));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));
                }
                arg if arg.starts_with("-Zmiri-env-exclude=") => {
                    excluded_env_vars
                        .push(arg.trim_start_matches("-Zmiri-env-exclude=").to_owned());
                }
                arg if arg.starts_with("-Zmiri-track-pointer-tag=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-pointer-tag=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-pointer-tag requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    if let Some(id) = miri::PtrId::new(id) {
                        tracked_pointer_tag = Some(id);
                    } else {
                        panic!("-Zmiri-track-pointer-tag must be a nonzero id");
                    }
                }
                arg if arg.starts_with("-Zmiri-track-alloc-id=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-alloc-id=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-alloc-id requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    tracked_alloc_id = Some(miri::AllocId(id));
                }
                _ => {
                    rustc_args.push(arg);
                }
            }
        }
    }

    // Determine sysroot if needed.  Make sure we always call `compile_time_sysroot`
    // as that also does some sanity-checks of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error, but
    // CTFE does not seem powerful enough for that yet.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot";
        if !rustc_args.iter().any(|e| e == sysroot_flag) {
            // We need to overwrite the default that librustc would compute.
            rustc_args.push(sysroot_flag.to_owned());
            rustc_args.push(sysroot);
        }
    }

    // Finally, add the default flags all the way in the beginning, but after the binary name.
    rustc_args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("miri arguments: {:?}", miri_args);
    let miri_config = miri::MiriConfig {
        validate,
        stacked_borrows,
        check_alignment,
        communicate,
        ignore_leaks,
        excluded_env_vars,
        seed,
        args: miri_args,
        tracked_pointer_tag,
        tracked_alloc_id,
    };
    rustc_driver::install_ice_hook();
    let result = rustc_driver::catch_fatal_errors(move || {
        rustc_driver::run_compiler(&rustc_args, &mut MiriCompilerCalls { miri_config }, None, None)
    })
    .and_then(|result| result);
    let exit_code = match result {
        Ok(_) => rustc_driver::EXIT_SUCCESS,
        Err(_) => rustc_driver::EXIT_FAILURE,
    };
    std::process::exit(exit_code);
}
<commit_msg>make sure we do not discard information on the Ok exit side<commit_after>#![feature(rustc_private)]

extern crate rustc_middle;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;

use std::convert::TryFrom;
use std::env;
use std::str::FromStr;

use hex::FromHexError;
use log::debug;

use rustc_session::CtfeBacktrace;
use rustc_driver::Compilation;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_interface::{interface, Queries};
use rustc_middle::ty::TyCtxt;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &interface::Compiler,
        queries: &'tcx Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();

        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            init_late_loggers(tcx);
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            if let Some(return_code) = miri::eval_main(tcx, entry_def_id, config) {
                std::process::exit(
                    i32::try_from(return_code).expect("Return value was too large!"),
                );
            }
        });

        compiler.session().abort_if_errors();

        Compilation::Stop
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var("RUSTC_LOG").is_ok() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers(tcx: TyCtxt<'_>) {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var("RUSTC_LOG").is_err() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var(
                    "RUSTC_LOG",
                    &format!("rustc_middle::mir::interpret={0},rustc_mir::interpret={0}", var),
                );
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUSTC_CTFE_BACKTRACE` is not, set `RUSTC_CTFE_BACKTRACE`.
    // Do this late, so we ideally only apply this to Miri's errors.
    if let Ok(val) = env::var("MIRI_BACKTRACE") {
        let ctfe_backtrace = match &*val {
            "immediate" => CtfeBacktrace::Immediate,
            "0" => CtfeBacktrace::Disabled,
            _ => CtfeBacktrace::Capture,
        };
        *tcx.sess.ctfe_backtrace.borrow_mut() = ctfe_backtrace;
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc would
    // end up somewhere in the build dir.
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
            .to_owned(),
    })
}

fn main() {
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut stacked_borrows = true;
    let mut check_alignment = true;
    let mut communicate = false;
    let mut ignore_leaks = false;
    let mut seed: Option<u64> = None;
    let mut tracked_pointer_tag: Option<miri::PtrId> = None;
    let mut tracked_alloc_id: Option<miri::AllocId> = None;
    let mut rustc_args = vec![];
    let mut miri_args = vec![];
    let mut after_dashdash = false;
    let mut excluded_env_vars = vec![];
    for arg in std::env::args() {
        if rustc_args.is_empty() {
            // Very first arg: for `rustc`.
            rustc_args.push(arg);
        } else if after_dashdash {
            // Everything that comes after are `miri` args.
            miri_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                }
                "-Zmiri-disable-stacked-borrows" => {
                    stacked_borrows = false;
                }
                "-Zmiri-disable-alignment-check" => {
                    check_alignment = false;
                }
                "-Zmiri-disable-isolation" => {
                    communicate = true;
                }
                "-Zmiri-ignore-leaks" => {
                    ignore_leaks = true;
                }
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.trim_start_matches("-Zmiri-seed="))
                        .unwrap_or_else(|err| match err {
                            FromHexError::InvalidHexCharacter { .. } => panic!(
                                "-Zmiri-seed should only contain valid hex digits [0-9a-fA-F]"
                            ),
                            FromHexError::OddLength =>
                                panic!("-Zmiri-seed should have an even number of digits"),
                            err => panic!("Unknown error decoding -Zmiri-seed as hex: {:?}", err),
                        });
                    if seed_raw.len() > 8 {
                        panic!(format!(
                            "-Zmiri-seed must be at most 8 bytes, was {}",
                            seed_raw.len()
                        ));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));
                }
                arg if arg.starts_with("-Zmiri-env-exclude=") => {
                    excluded_env_vars
                        .push(arg.trim_start_matches("-Zmiri-env-exclude=").to_owned());
                }
                arg if arg.starts_with("-Zmiri-track-pointer-tag=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-pointer-tag=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-pointer-tag requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    if let Some(id) = miri::PtrId::new(id) {
                        tracked_pointer_tag = Some(id);
                    } else {
                        panic!("-Zmiri-track-pointer-tag must be a nonzero id");
                    }
                }
                arg if arg.starts_with("-Zmiri-track-alloc-id=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-alloc-id=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-alloc-id requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    tracked_alloc_id = Some(miri::AllocId(id));
                }
                _ => {
                    rustc_args.push(arg);
                }
            }
        }
    }

    // Determine sysroot if needed.  Make sure we always call `compile_time_sysroot`
    // as that also does some sanity-checks of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error, but
    // CTFE does not seem powerful enough for that yet.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot";
        if !rustc_args.iter().any(|e| e == sysroot_flag) {
            // We need to overwrite the default that librustc would compute.
            rustc_args.push(sysroot_flag.to_owned());
            rustc_args.push(sysroot);
        }
    }

    // Finally, add the default flags all the way in the beginning, but after the binary name.
    rustc_args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("miri arguments: {:?}", miri_args);
    let miri_config = miri::MiriConfig {
        validate,
        stacked_borrows,
        check_alignment,
        communicate,
        ignore_leaks,
        excluded_env_vars,
        seed,
        args: miri_args,
        tracked_pointer_tag,
        tracked_alloc_id,
    };
    rustc_driver::install_ice_hook();
    let result = rustc_driver::catch_fatal_errors(move || {
        rustc_driver::run_compiler(&rustc_args, &mut MiriCompilerCalls { miri_config }, None, None)
    })
    .and_then(|result| result);
    let exit_code = match result {
        Ok(()) => rustc_driver::EXIT_SUCCESS,
        Err(_) => rustc_driver::EXIT_FAILURE,
    };
    std::process::exit(exit_code);
}
<|endoftext|>


<commit_before>extern crate ilog2;
extern crate num;

use std::fmt::Binary;
use self::ilog2::Bitops;

pub trait Bitalloc {
    fn clear_all(&mut self);
    fn fill_all(&mut self);
    fn alloc(&mut self, size: usize) -> isize;
    fn free(&mut self, from: usize, size: usize);
}

impl<T: ilog2::Bitops + Binary> Bitalloc for [T] {
    fn clear_all(&mut self) {
        for item in self.iter_mut() {
            *item = T::zero();
        }
    }
    fn fill_all(&mut self) {
        for item in self.iter_mut() {
            *item = ilog2::bit_mask::<T>();
        }
    }
    fn alloc(&mut self, size: usize) -> isize {
        let block_size = ilog2::bit_length::<T>() as usize;
        let block_count = size / block_size;
        let rem_count = size % block_size;

        let mut blocks_allocated = 0usize;
        let mut rem_allocated = 0usize;
        let mut blocks_needed = block_count;
        let mut rem_needed = rem_count;

        enum State {
            FindFirst,
            FindMid,
            FindLast,
        };

        let mut found_index = 0usize;
        let mut found_offset = 0usize;
        let mut found = false;
        let mut state = State::FindFirst;
        let mut index = 0usize;
        let limit = self.len();
        loop {
            if index == limit {
                break;
            }
            let val = self[index];

            match state {
                State::FindFirst => {
                    //reset
                    blocks_allocated = 0;
                    rem_allocated = 0;
                    blocks_needed = block_count;
                    rem_needed = rem_count;
                    found_index = 0;
                    found_offset = 0;

                    if size <= block_size { // allocated region may inside an integer
                        let mut possible = val;
                        let mut shifted = 0usize;
                        loop {
                            let current_available = possible.leading_zeros() as usize;
                            //println!("debug2 {} {} {} {:064b}", current_available, size, shifted, possible);
                            if current_available >= size {
                                found = true;
                                found_index = index;
                                found_offset = shifted;
                                break;
                            }
                            possible = possible << current_available; //discard current;
                            shifted += current_available;
                            //println!("debug3 {} {} {} {:064b}", current_available, size, shifted, possible);
                            let current_skip = possible.leading_ones() as usize;
                            shifted += current_skip;
                            if (shifted + size) > block_size {
                                break;
                            }
                            possible = possible << current_skip;
                            possible = possible | ilog2::bit_mask::<T>() >> (block_size - shifted);
                            //println!("debug4 {} {} {} {:064b}", current_available, size, shifted, possible);
                        }
                        if found {
                            break;
                        }
                    }
                    //size exceeded block size or failed to allocated inside a block
                    //find zeros from lsb position
                    let current_available = val.trailing_zeros() as usize;
                    if current_available > 0 {
                        //possible start location
                        found_index = index;
                        found_offset = block_size - current_available;
                        if current_available > rem_needed {
                            assert!(blocks_needed > 0);
                            rem_needed = rem_needed + block_size - current_available;
                            blocks_needed -= 1;
                        }
                        rem_allocated += current_available;
                        state = State::FindMid;
                    }
                },
                State::FindMid => {
                    if blocks_allocated == blocks_needed {
                        state = State::FindLast;
                        continue;
                    } else if val == T::zero() {
                        blocks_allocated += 1;
                    } else {
                        state = State::FindFirst;
                        continue;
                    }
                },
                State::FindLast => {
                    let current_available = val.leading_zeros() as usize;
                    rem_allocated += current_available;
                    if rem_allocated >= rem_needed {
                        found = true;
                        break;
                    }
                    state = State::FindFirst;
                    continue;
                },
            }
            index += 1;
        }
        if found {
            let mut allocated = 0usize;
            let mut first_mask = ilog2::bit_mask::<T>();
            first_mask = first_mask >> found_offset;
            allocated += block_size - found_offset;
            if allocated > size {
                let diff = allocated - size;
                first_mask = first_mask >> diff;
                first_mask = first_mask << diff;
                allocated = size;
            }
            let mut update_index = found_index;
            assert!(self[update_index] & first_mask == T::zero());
            //println!("Allocate first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
            self[update_index] = self[update_index] | first_mask;
            //println!("Allocate first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
            update_index += 1;

            let rem = size - allocated;
            let fill = rem / block_size;
            let last_fill = rem % block_size;

            for i in 0 .. fill {
                //println!("Allocate mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
                let mid_mask = ilog2::bit_mask::<T>();
                assert!(self[update_index + i] & mid_mask == T::zero());
                self[update_index + i] = self[update_index + i] | mid_mask;
                allocated += block_size;
                //println!("Allocate mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
            }
            if last_fill > 0 {
                //println!("Allocate last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
                let mut last_mask = ilog2::bit_mask::<T>();
                let diff = block_size - last_fill;
                last_mask = last_mask >> diff;
                last_mask = last_mask << diff;
                assert!(self[update_index + fill] & last_mask == T::zero());
                self[update_index + fill] = self[update_index + fill] | last_mask;
                allocated += last_fill;
                //println!("Allocate last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
            }
            assert!(allocated == size);

            ((found_index * block_size) + found_offset) as isize
        } else {
            -1isize
        }
    }
    fn free(&mut self, from: usize, size: usize) {
        let block_size = ilog2::bit_length::<T>() as usize;
        let mut freed = 0usize;
        let mut first_mask = ilog2::bit_mask::<T>();
        let found_offset = from % block_size;
        let found_index = from / block_size;
        first_mask = first_mask >> found_offset;
        freed += block_size - found_offset;
        if freed > size {
            let diff = freed - size;
            first_mask = first_mask >> diff;
            first_mask = first_mask << diff;
            freed = size;
        }
        let mut update_index = found_index;
        assert!(self[update_index] & first_mask == first_mask);
        //println!("Free first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
        self[update_index] = self[update_index] ^ first_mask;
        //println!("Free first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
        update_index += 1;

        let rem = size - freed;
        let fill = rem / block_size;
        let last_fill = rem % block_size;

        for i in 0 .. fill {
            //println!("Free mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
            let mid_mask = ilog2::bit_mask::<T>();
            assert!(self[update_index + i] & mid_mask == mid_mask);
            self[update_index + i] = self[update_index + i] ^ mid_mask;
            freed += block_size;
            //println!("Free mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
        }
        if last_fill > 0 {
            //println!("Free last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
            let mut last_mask = ilog2::bit_mask::<T>();
            let diff = block_size - last_fill;
            last_mask = last_mask >> diff;
            last_mask = last_mask << diff;
            assert!(self[update_index + fill] & last_mask == last_mask);
            self[update_index + fill] = self[update_index + fill] ^ last_mask;
            freed += last_fill;
            //println!("Free last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
        }
        assert!(freed == size);
    }
}<commit_msg>Remove warnings.<commit_after>extern crate ilog2;
extern crate num;

use std::fmt::Binary;

pub trait Bitalloc {
    fn clear_all(&mut self);
    fn fill_all(&mut self);
    fn alloc(&mut self, size: usize) -> isize;
    fn free(&mut self, from: usize, size: usize);
}

impl<T: ilog2::Bitops + Binary> Bitalloc for [T] {
    fn clear_all(&mut self) {
        for item in self.iter_mut() {
            *item = T::zero();
        }
    }
    fn fill_all(&mut self) {
        for item in self.iter_mut() {
            *item = ilog2::bit_mask::<T>();
        }
    }
    fn alloc(&mut self, size: usize) -> isize {
        let block_size = ilog2::bit_length::<T>() as usize;
        let block_count = size / block_size;
        let rem_count = size % block_size;

        let mut blocks_allocated = 0usize;
        let mut rem_allocated = 0usize;
        let mut blocks_needed = block_count;
        let mut rem_needed = rem_count;

        enum State {
            FindFirst,
            FindMid,
            FindLast,
        };

        let mut found_index = 0usize;
        let mut found_offset = 0usize;
        let mut found = false;
        let mut state = State::FindFirst;
        let mut index = 0usize;
        let limit = self.len();
        loop {
            if index == limit {
                break;
            }
            let val = self[index];

            match state {
                State::FindFirst => {
                    //reset
                    blocks_allocated = 0;
                    rem_allocated = 0;
                    blocks_needed = block_count;
                    rem_needed = rem_count;
                    found_index = 0;
                    found_offset = 0;

                    if size <= block_size { // allocated region may inside an integer
                        let mut possible = val;
                        let mut shifted = 0usize;
                        loop {
                            let current_available = possible.leading_zeros() as usize;
                            //println!("debug2 {} {} {} {:064b}", current_available, size, shifted, possible);
                            if current_available >= size {
                                found = true;
                                found_index = index;
                                found_offset = shifted;
                                break;
                            }
                            possible = possible << current_available; //discard current;
                            shifted += current_available;
                            //println!("debug3 {} {} {} {:064b}", current_available, size, shifted, possible);
                            let current_skip = possible.leading_ones() as usize;
                            shifted += current_skip;
                            if (shifted + size) > block_size {
                                break;
                            }
                            possible = possible << current_skip;
                            possible = possible | ilog2::bit_mask::<T>() >> (block_size - shifted);
                            //println!("debug4 {} {} {} {:064b}", current_available, size, shifted, possible);
                        }
                        if found {
                            break;
                        }
                    }
                    //size exceeded block size or failed to allocated inside a block
                    //find zeros from lsb position
                    let current_available = val.trailing_zeros() as usize;
                    if current_available > 0 {
                        //possible start location
                        found_index = index;
                        found_offset = block_size - current_available;
                        if current_available > rem_needed {
                            assert!(blocks_needed > 0);
                            rem_needed = rem_needed + block_size - current_available;
                            blocks_needed -= 1;
                        }
                        rem_allocated += current_available;
                        state = State::FindMid;
                    }
                },
                State::FindMid => {
                    if blocks_allocated == blocks_needed {
                        state = State::FindLast;
                        continue;
                    } else if val == T::zero() {
                        blocks_allocated += 1;
                    } else {
                        state = State::FindFirst;
                        continue;
                    }
                },
                State::FindLast => {
                    let current_available = val.leading_zeros() as usize;
                    rem_allocated += current_available;
                    if rem_allocated >= rem_needed {
                        found = true;
                        break;
                    }
                    state = State::FindFirst;
                    continue;
                },
            }
            index += 1;
        }
        if found {
            let mut allocated = 0usize;
            let mut first_mask = ilog2::bit_mask::<T>();
            first_mask = first_mask >> found_offset;
            allocated += block_size - found_offset;
            if allocated > size {
                let diff = allocated - size;
                first_mask = first_mask >> diff;
                first_mask = first_mask << diff;
                allocated = size;
            }
            let mut update_index = found_index;
            assert!(self[update_index] & first_mask == T::zero());
            //println!("Allocate first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
            self[update_index] = self[update_index] | first_mask;
            //println!("Allocate first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
            update_index += 1;

            let rem = size - allocated;
            let fill = rem / block_size;
            let last_fill = rem % block_size;

            for i in 0 .. fill {
                //println!("Allocate mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
                let mid_mask = ilog2::bit_mask::<T>();
                assert!(self[update_index + i] & mid_mask == T::zero());
                self[update_index + i] = self[update_index + i] | mid_mask;
                allocated += block_size;
                //println!("Allocate mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
            }
            if last_fill > 0 {
                //println!("Allocate last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
                let mut last_mask = ilog2::bit_mask::<T>();
                let diff = block_size - last_fill;
                last_mask = last_mask >> diff;
                last_mask = last_mask << diff;
                assert!(self[update_index + fill] & last_mask == T::zero());
                self[update_index + fill] = self[update_index + fill] | last_mask;
                allocated += last_fill;
                //println!("Allocate last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
            }
            assert!(allocated == size);

            ((found_index * block_size) + found_offset) as isize
        } else {
            -1isize
        }
    }
    fn free(&mut self, from: usize, size: usize) {
        let block_size = ilog2::bit_length::<T>() as usize;
        let mut freed = 0usize;
        let mut first_mask = ilog2::bit_mask::<T>();
        let found_offset = from % block_size;
        let found_index = from / block_size;
        first_mask = first_mask >> found_offset;
        freed += block_size - found_offset;
        if freed > size {
            let diff = freed - size;
            first_mask = first_mask >> diff;
            first_mask = first_mask << diff;
            freed = size;
        }
        let mut update_index = found_index;
        assert!(self[update_index] & first_mask == first_mask);
        //println!("Free first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
        self[update_index] = self[update_index] ^ first_mask;
        //println!("Free first at block {} [{},{}] : {:064b}", update_index, found_index, found_offset, self[update_index]);
        update_index += 1;

        let rem = size - freed;
        let fill = rem / block_size;
        let last_fill = rem % block_size;

        for i in 0 .. fill {
            //println!("Free mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
            let mid_mask = ilog2::bit_mask::<T>();
            assert!(self[update_index + i] & mid_mask == mid_mask);
            self[update_index + i] = self[update_index + i] ^ mid_mask;
            freed += block_size;
            //println!("Free mid at block {} [{},{}] : {:064b}", update_index + i, found_index, found_offset, self[update_index + i]);
        }
        if last_fill > 0 {
            //println!("Free last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
            let mut last_mask = ilog2::bit_mask::<T>();
            let diff = block_size - last_fill;
            last_mask = last_mask >> diff;
            last_mask = last_mask << diff;
            assert!(self[update_index + fill] & last_mask == last_mask);
            self[update_index + fill] = self[update_index + fill] ^ last_mask;
            freed += last_fill;
            //println!("Free last at block {} [{},{}] : {:064b}", update_index + fill, found_index, found_offset, self[update_index + fill]);
        }
        assert!(freed == size);
    }
}<|endoftext|>


<commit_before>//! Use this module to interact with Bitstamp exchange.
//! Please see examples for more informations.

use hyper_native_tls::NativeTlsClient;
use hyper::Client;
use hyper::net::HttpsConnector;

use serde_json;
use serde_json::Value;
use serde_json::value::Map;

use std::collections::HashMap;
use std::io::Read;
use std::thread;
use std::time::Duration;
use std::path::PathBuf;
use std::fs::File;

use helpers;

header! {
    #[doc(hidden)]
    (KeyHeader, "Key") => [String]
}

header! {
    #[doc(hidden)]
    (SignHeader, "Sign") => [String]
}

header! {
    #[doc(hidden)]
    (ContentHeader, "Content-Type") => [String]
}

pub struct BitstampApi {
    last_request: i64, // unix timestamp in ms, to avoid ban
    api_key: String,
    api_secret: String,
    http_client: Client,
}

impl BitstampApi {
    /// Create a new BitstampApi by providing an API key & API secret
    pub fn new(api_key: &str, api_secret: &str) -> BitstampApi {
        let ssl = NativeTlsClient::new().unwrap();
        let connector = HttpsConnector::new(ssl);

        BitstampApi {
            last_request: 0,
            api_key: api_key.to_string(),
            api_secret: api_secret.to_string(),
            http_client: Client::with_connector(connector),
        }
    }

    /// Create a new BitstampApi from a json configuration file. This file must follow this
    /// structure:
    ///
    /// ```ignore
    /// {
    ///     "account_kraken": {
    ///         "exchange"  : "kraken",
    ///         "api_key"   : "123456789ABCDEF",
    ///         "api_secret": "ABC&EF?abcdef"
    ///     },
    ///     "account_bitstamp": {
    ///         "exchange"  : "bitstamp",
    ///         "api_key"   : "XYXY-XYXY-XYXY-XY",
    ///         "api_secret": "A0A0B1B1C2C2"
    ///     }
    /// }
    /// ```
    /// For this example, you could use load your Bitstamp account with
    /// `new_from_file("account_bitstamp", Path::new("/keys.json"))`
    pub fn new_from_file(config_name: &str, path: PathBuf) -> BitstampApi {
        let mut f = File::open(&path).unwrap();
        let mut buffer = String::new();
        f.read_to_string(&mut buffer).unwrap();

        let data: Value = serde_json::from_str(&buffer).unwrap();
        let json_obj = data.as_object().unwrap().get(config_name).unwrap();
        let api_key = json_obj.get("api_key").unwrap().as_str().unwrap();
        let api_secret = json_obj.get("api_secret").unwrap().as_str().unwrap();

        BitstampApi::new(api_key, api_secret)
    }

    fn block_or_continue(&self) {
        let threshold = 1000; // 600 requests per 10 mins = 1 request per second
        let delay = helpers::get_unix_timestamp_ms() - self.last_request;
        if delay < threshold {
            let duration_ms = Duration::from_millis(delay as u64);
            thread::sleep(duration_ms);
        }
    }

    fn deserialize_json(&mut self, json_string: String) -> Option<Map<String, Value>> {
        let data: Value = serde_json::from_str(&json_string).unwrap();

        match data.as_object() {
            Some(value) => Some(value.clone()),
            None => None,
        }
    }

    pub fn build_url(method: &str, pair: &str) -> String {
        "https://www.bitstamp.net/api/v2/".to_string() + method + "/" + &pair + "/"
    }

    fn public_query(&mut self,
                    params: &HashMap<&str, &str>)
                    -> Option<Map<String, Value>> {
        let mut params = params.clone();
        helpers::strip_empties(&mut params);

        // TODO pass the pair as parameter, dont hardcode it #FIXME
        let method: &str = params.get("method").unwrap();
        let pair: &str = params.get("pair").unwrap();
        let url: String = BitstampApi::build_url(method, pair);

        self.block_or_continue();
        let mut response = self.http_client.get(&url).send().unwrap();
        self.last_request = helpers::get_unix_timestamp_ms();
        let mut buffer = String::new();
        response.read_to_string(&mut buffer).unwrap();
        return self.deserialize_json(buffer);
    }

    /// Sample output :
    ///
    /// ```ignore
    /// {
    /// "BTC_LTC":{
    /// "last":"0.0251","lowestAsk":"0.02589999","highestBid":"0.0251",
    /// "percentChange":"0.02390438","baseVolume":"6.16485315","quoteVolume":"245.82513926"},
    /// "BTC_NXT":{
    /// "last":"0.00005730","lowestAsk":"0.00005710","highestBid":"0.00004903",
    /// "percentChange":"0.16701570","baseVolume":"0.45347489","quoteVolume":"9094"},
    /// ... }
    /// ```
    pub fn return_ticker(&mut self) -> Option<Map<String, Value>> {
        let mut params = HashMap::new();
        params.insert("method", "ticker");
        params.insert("pair", "btcusd");
        self.public_query(&params)
    }

    /// Sample output :
    ///
    /// ```ignore
    /// {"asks":[[0.00007600,1164],[0.00007620,1300], ... ], "bids":[[0.00006901,200],
    /// [0.00006900,408], ... ], "isFrozen": 0, "seq": 18849}
    /// ```
    pub fn return_order_book(&mut self,
                             pair: &str)
                             -> Option<Map<String, Value>> {
        let mut params = HashMap::new();
        params.insert("method", "order_book");
        params.insert("pair", pair);
        self.public_query(&params)
    }
}
<commit_msg>Fix doc<commit_after>//! Use this module to interact with Bitstamp exchange.
//! Please see examples for more informations.

use hyper_native_tls::NativeTlsClient;
use hyper::Client;
use hyper::net::HttpsConnector;

use serde_json;
use serde_json::Value;
use serde_json::value::Map;

use std::collections::HashMap;
use std::io::Read;
use std::thread;
use std::time::Duration;
use std::path::PathBuf;
use std::fs::File;

use helpers;

header! {
    #[doc(hidden)]
    (KeyHeader, "Key") => [String]
}

header! {
    #[doc(hidden)]
    (SignHeader, "Sign") => [String]
}

header! {
    #[doc(hidden)]
    (ContentHeader, "Content-Type") => [String]
}

pub struct BitstampApi {
    last_request: i64, // unix timestamp in ms, to avoid ban
    api_key: String,
    api_secret: String,
    http_client: Client,
}

impl BitstampApi {
    /// Create a new BitstampApi by providing an API key & API secret
    pub fn new(api_key: &str, api_secret: &str) -> BitstampApi {
        let ssl = NativeTlsClient::new().unwrap();
        let connector = HttpsConnector::new(ssl);

        BitstampApi {
            last_request: 0,
            api_key: api_key.to_string(),
            api_secret: api_secret.to_string(),
            http_client: Client::with_connector(connector),
        }
    }

    /// Create a new BitstampApi from a json configuration file. This file must follow this
    /// structure:
    ///
    /// ```ignore
    /// {
    ///     "account_kraken": {
    ///         "exchange"  : "kraken",
    ///         "api_key"   : "123456789ABCDEF",
    ///         "api_secret": "ABC&EF?abcdef"
    ///     },
    ///     "account_bitstamp": {
    ///         "exchange"  : "bitstamp",
    ///         "api_key"   : "XYXY-XYXY-XYXY-XY",
    ///         "api_secret": "A0A0B1B1C2C2"
    ///     }
    /// }
    /// ```
    /// For this example, you could use load your Bitstamp account with
    /// `new_from_file("account_bitstamp", Path::new("/keys.json"))`
    pub fn new_from_file(config_name: &str, path: PathBuf) -> BitstampApi {
        let mut f = File::open(&path).unwrap();
        let mut buffer = String::new();
        f.read_to_string(&mut buffer).unwrap();

        let data: Value = serde_json::from_str(&buffer).unwrap();
        let json_obj = data.as_object().unwrap().get(config_name).unwrap();
        let api_key = json_obj.get("api_key").unwrap().as_str().unwrap();
        let api_secret = json_obj.get("api_secret").unwrap().as_str().unwrap();

        BitstampApi::new(api_key, api_secret)
    }

    fn block_or_continue(&self) {
        let threshold = 1000; // 600 requests per 10 mins = 1 request per second
        let delay = helpers::get_unix_timestamp_ms() - self.last_request;
        if delay < threshold {
            let duration_ms = Duration::from_millis(delay as u64);
            thread::sleep(duration_ms);
        }
    }

    fn deserialize_json(&mut self, json_string: String) -> Option<Map<String, Value>> {
        let data: Value = serde_json::from_str(&json_string).unwrap();

        match data.as_object() {
            Some(value) => Some(value.clone()),
            None => None,
        }
    }

    pub fn build_url(method: &str, pair: &str) -> String {
        "https://www.bitstamp.net/api/v2/".to_string() + method + "/" + &pair + "/"
    }

    fn public_query(&mut self,
                    params: &HashMap<&str, &str>)
                    -> Option<Map<String, Value>> {
        let mut params = params.clone();
        helpers::strip_empties(&mut params);

        // TODO pass the pair as parameter, dont hardcode it #FIXME
        let method: &str = params.get("method").unwrap();
        let pair: &str = params.get("pair").unwrap();
        let url: String = BitstampApi::build_url(method, pair);

        self.block_or_continue();
        let mut response = self.http_client.get(&url).send().unwrap();
        self.last_request = helpers::get_unix_timestamp_ms();
        let mut buffer = String::new();
        response.read_to_string(&mut buffer).unwrap();
        return self.deserialize_json(buffer);
    }

    /// Sample output :
    ///
    /// ```ignore
    /// {
    /// "BTC_LTC":{
    /// "last":"0.0251","lowestAsk":"0.02589999","highestBid":"0.0251",
    /// "percentChange":"0.02390438","baseVolume":"6.16485315","quoteVolume":"245.82513926"},
    /// "BTC_NXT":{
    /// "last":"0.00005730","lowestAsk":"0.00005710","highestBid":"0.00004903",
    /// "percentChange":"0.16701570","baseVolume":"0.45347489","quoteVolume":"9094"},
    /// ... }
    /// ```
    pub fn return_ticker(&mut self) -> Option<Map<String, Value>> {
        let mut params = HashMap::new();
        params.insert("method", "ticker");
        params.insert("pair", "btcusd");
        self.public_query(&params)
    }

    /// Sample output :
    ///
    /// ```ignore
    /// {"asks":[[0.00007600,1164],[0.00007620,1300], ... ], "bids":[[0.00006901,200],
    /// [0.00006900,408], ... ], "timestamp": 1234567890"}
    /// ```
    pub fn return_order_book(&mut self,
                             pair: &str)
                             -> Option<Map<String, Value>> {
        let mut params = HashMap::new();
        params.insert("method", "order_book");
        params.insert("pair", pair);
        self.public_query(&params)
    }
}
<|endoftext|>


<commit_before>use {alloc, Bytes, ByteBuf, ROByteBuf, Rope};
use traits::{Buf, MutBuf, MutBufExt, ByteStr, ToBytes};
use std::{cmp, ops};

/*
 *
 * ===== SeqByteStr =====
 *
 */

pub struct SeqByteStr {
    mem: alloc::MemRef,
    pos: u32,
    len: u32,
}

impl SeqByteStr {
    /// Create a new `SeqByteStr` from a byte slice.
    ///
    /// The contents of the byte slice will be copied.
    pub fn from_slice(bytes: &[u8]) -> SeqByteStr {
        let mut buf = ByteBuf::mut_with_capacity(bytes.len());

        if let Err(e) = buf.write(bytes) {
            panic!("failed to copy bytes from slice; err={:?}", e);
        }

        buf.flip().to_seq_byte_str()
    }

    /// Creates a new `SeqByteStr` from a `MemRef`, an offset, and a length.
    ///
    /// This function is unsafe as there are no guarantees that the given
    /// arguments are valid.
    pub unsafe fn from_mem_ref(mem: alloc::MemRef, pos: u32, len: u32) -> SeqByteStr {
        SeqByteStr {
            mem: mem,
            pos: pos,
            len: len,
        }
    }
}

impl ByteStr for SeqByteStr {
    type Buf = ROByteBuf;

    fn buf(&self) -> ROByteBuf {
        unsafe {
            let pos = self.pos;
            let lim = pos + self.len;

            ROByteBuf::from_mem_ref(self.mem.clone(), lim, pos, lim)
        }
    }

    fn concat<B: ByteStr+'static>(&self, other: &B) -> Bytes {
        Rope::of(self.clone()).concat(other)
    }

    fn len(&self) -> usize {
        self.len as usize
    }

    fn slice(&self, begin: usize, end: usize) -> Bytes {
        if begin >= end || begin >= self.len() {
            return Bytes::empty()
        }

        let bytes = unsafe {
            SeqByteStr::from_mem_ref(
                self.mem.clone(),
                self.pos + begin as u32,
                (end - begin) as u32)
        };

        Bytes::of(bytes)
    }
}

impl ToBytes for SeqByteStr {
    fn to_bytes(self) -> Bytes {
        Bytes::of(self)
    }
}

impl ops::Index<usize> for SeqByteStr {
    type Output = u8;

    fn index(&self, index: &usize) -> &u8 {
        assert!(*index < self.len());

        unsafe {
            &*self.mem.ptr()
                .offset(*index as isize + self.pos as isize)
        }
    }
}

impl Clone for SeqByteStr {
    fn clone(&self) -> SeqByteStr {
        SeqByteStr {
            mem: self.mem.clone(),
            pos: self.pos,
            len: self.len,
        }
    }
}

/*
 *
 * ===== SmallByteStr =====
 *
 */

#[cfg(target_pointer_width = "64")]
const MAX_LEN: usize = 7;

#[cfg(target_pointer_width = "32")]
const MAX_LEN: usize = 3;

#[derive(Clone, Copy)]
pub struct SmallByteStr {
    len: u8,
    bytes: [u8; MAX_LEN],
}

impl SmallByteStr {
    pub fn zero() -> SmallByteStr {
        use std::mem;

        SmallByteStr {
            len: 0,
            bytes: unsafe { mem::zeroed() }
        }
    }

    pub fn from_slice(bytes: &[u8]) -> Option<SmallByteStr> {
        use std::mem;
        use std::slice::bytes;

        if bytes.len() > MAX_LEN {
            return None;
        }

        let mut ret = SmallByteStr {
            len: bytes.len() as u8,
            bytes: unsafe { mem::zeroed() },
        };

        // Copy the memory
        bytes::copy_memory(&mut ret.bytes, bytes);

        Some(ret)
    }
}

impl ByteStr for SmallByteStr {
    type Buf = SmallByteStrBuf;

    fn buf(&self) -> SmallByteStrBuf {
        SmallByteStrBuf { small: self.clone() }
    }

    fn concat<B: ByteStr+'static>(&self, other: &B) -> Bytes {
        Rope::of(self.clone()).concat(other)
    }

    fn len(&self) -> usize {
        self.len as usize
    }

    fn slice(&self, _begin: usize, _end: usize) -> Bytes {
        unimplemented!();
    }

    fn split_at(&self, _mid: usize) -> (Bytes, Bytes) {
        unimplemented!();
    }
}

impl ToBytes for SmallByteStr {
    fn to_bytes(self) -> Bytes {
        Bytes::of(self)
    }
}

impl ops::Index<usize> for SmallByteStr {
    type Output = u8;

    fn index(&self, index: &usize) -> &u8 {
        assert!(*index < self.len());
        &self.bytes[*index]
    }
}

#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub struct SmallByteStrBuf {
    small: SmallByteStr,
}

impl SmallByteStrBuf {
    fn len(&self) -> usize {
        (self.small.len & 0x0F) as usize
    }

    fn pos(&self) -> usize {
        (self.small.len >> 4) as usize
    }
}

impl Buf for SmallByteStrBuf {
    fn remaining(&self) -> usize {
        self.len() - self.pos()
    }

    fn bytes(&self) -> &[u8] {
        &self.small.bytes[self.pos()..self.len()]
    }

    fn advance(&mut self, mut cnt: usize) {
        cnt = cmp::min(cnt, self.remaining());
        self.small.len += (cnt as u8) << 4;
    }
}

#[test]
pub fn test_size_of() {
    use std::mem;
    assert_eq!(mem::size_of::<SmallByteStr>(), mem::size_of::<usize>());
}
<commit_msg>Implement slice for small ByteStr<commit_after>use {alloc, Bytes, ByteBuf, ROByteBuf, Rope};
use traits::{Buf, MutBuf, MutBufExt, ByteStr, ToBytes};
use std::{cmp, ops};

/*
 *
 * ===== SeqByteStr =====
 *
 */

pub struct SeqByteStr {
    mem: alloc::MemRef,
    pos: u32,
    len: u32,
}

impl SeqByteStr {
    /// Create a new `SeqByteStr` from a byte slice.
    ///
    /// The contents of the byte slice will be copied.
    pub fn from_slice(bytes: &[u8]) -> SeqByteStr {
        let mut buf = ByteBuf::mut_with_capacity(bytes.len());

        if let Err(e) = buf.write(bytes) {
            panic!("failed to copy bytes from slice; err={:?}", e);
        }

        buf.flip().to_seq_byte_str()
    }

    /// Creates a new `SeqByteStr` from a `MemRef`, an offset, and a length.
    ///
    /// This function is unsafe as there are no guarantees that the given
    /// arguments are valid.
    pub unsafe fn from_mem_ref(mem: alloc::MemRef, pos: u32, len: u32) -> SeqByteStr {
        SeqByteStr {
            mem: mem,
            pos: pos,
            len: len,
        }
    }
}

impl ByteStr for SeqByteStr {
    type Buf = ROByteBuf;

    fn buf(&self) -> ROByteBuf {
        unsafe {
            let pos = self.pos;
            let lim = pos + self.len;

            ROByteBuf::from_mem_ref(self.mem.clone(), lim, pos, lim)
        }
    }

    fn concat<B: ByteStr+'static>(&self, other: &B) -> Bytes {
        Rope::of(self.clone()).concat(other)
    }

    fn len(&self) -> usize {
        self.len as usize
    }

    fn slice(&self, begin: usize, end: usize) -> Bytes {
        if begin >= end || begin >= self.len() {
            return Bytes::empty()
        }

        let bytes = unsafe {
            SeqByteStr::from_mem_ref(
                self.mem.clone(),
                self.pos + begin as u32,
                (end - begin) as u32)
        };

        Bytes::of(bytes)
    }
}

impl ToBytes for SeqByteStr {
    fn to_bytes(self) -> Bytes {
        Bytes::of(self)
    }
}

impl ops::Index<usize> for SeqByteStr {
    type Output = u8;

    fn index(&self, index: &usize) -> &u8 {
        assert!(*index < self.len());

        unsafe {
            &*self.mem.ptr()
                .offset(*index as isize + self.pos as isize)
        }
    }
}

impl Clone for SeqByteStr {
    fn clone(&self) -> SeqByteStr {
        SeqByteStr {
            mem: self.mem.clone(),
            pos: self.pos,
            len: self.len,
        }
    }
}

/*
 *
 * ===== SmallByteStr =====
 *
 */

#[cfg(target_pointer_width = "64")]
const MAX_LEN: usize = 7;

#[cfg(target_pointer_width = "32")]
const MAX_LEN: usize = 3;

#[derive(Clone, Copy)]
pub struct SmallByteStr {
    len: u8,
    bytes: [u8; MAX_LEN],
}

impl SmallByteStr {
    pub fn zero() -> SmallByteStr {
        use std::mem;

        SmallByteStr {
            len: 0,
            bytes: unsafe { mem::zeroed() }
        }
    }

    pub fn from_slice(bytes: &[u8]) -> Option<SmallByteStr> {
        use std::mem;
        use std::slice::bytes;

        if bytes.len() > MAX_LEN {
            return None;
        }

        let mut ret = SmallByteStr {
            len: bytes.len() as u8,
            bytes: unsafe { mem::zeroed() },
        };

        // Copy the memory
        bytes::copy_memory(&mut ret.bytes, bytes);

        Some(ret)
    }

    pub fn as_slice(&self) -> &[u8] {
        &self.bytes[..self.len as usize]
    }
}

impl ByteStr for SmallByteStr {
    type Buf = SmallByteStrBuf;

    fn buf(&self) -> SmallByteStrBuf {
        SmallByteStrBuf { small: self.clone() }
    }

    fn concat<B: ByteStr+'static>(&self, other: &B) -> Bytes {
        Rope::of(self.clone()).concat(other)
    }

    fn len(&self) -> usize {
        self.len as usize
    }

    fn slice(&self, begin: usize, end: usize) -> Bytes {
        Bytes::from_slice(&self.as_slice()[begin..end])
    }
}

impl ToBytes for SmallByteStr {
    fn to_bytes(self) -> Bytes {
        Bytes::of(self)
    }
}

impl ops::Index<usize> for SmallByteStr {
    type Output = u8;

    fn index(&self, index: &usize) -> &u8 {
        assert!(*index < self.len());
        &self.bytes[*index]
    }
}

#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub struct SmallByteStrBuf {
    small: SmallByteStr,
}

impl SmallByteStrBuf {
    fn len(&self) -> usize {
        (self.small.len & 0x0F) as usize
    }

    fn pos(&self) -> usize {
        (self.small.len >> 4) as usize
    }
}

impl Buf for SmallByteStrBuf {
    fn remaining(&self) -> usize {
        self.len() - self.pos()
    }

    fn bytes(&self) -> &[u8] {
        &self.small.bytes[self.pos()..self.len()]
    }

    fn advance(&mut self, mut cnt: usize) {
        cnt = cmp::min(cnt, self.remaining());
        self.small.len += (cnt as u8) << 4;
    }
}

#[test]
pub fn test_size_of() {
    use std::mem;
    assert_eq!(mem::size_of::<SmallByteStr>(), mem::size_of::<usize>());
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::io;
use std::str;
use std::mem;
use time;
use euclid::{Point2D, Rect};
use calx_resource::Resource;
use calx_alg::split_line;
use backend;

struct Message {
    expire_time_s: f64,
    text: String,
}

impl Message {
    fn new(text: String, time_start_s: f64) -> Message {
        const TIME_TO_READ_CHAR_S: f64 = 0.2;
        let expire_time_s = time_start_s + text.len() as f64 * TIME_TO_READ_CHAR_S;
        Message {
            expire_time_s: expire_time_s,
            text: text,
        }
    }
}

/// Output
pub struct Console {
    lines: Vec<Message>,
    input_buffer: String,
    output_buffer: String,
    done_reading_s: f64,
}

impl Console {
    pub fn new() -> Console {
        Console {
            lines: Vec::new(),
            input_buffer: String::new(),
            output_buffer: String::new(),
            done_reading_s: 0.0,
        }
    }

    /// Draw the console as a regular message display.
    pub fn draw_small(&mut self, context: &mut backend::Context, screen_area: &Rect<f32>) {
        // TODO: Store default font in context object
        let font: Resource<backend::Font> = Resource::new("default".to_string()).unwrap();
        // TODO: Ditto for color
        let color = [1.0, 1.0, 1.0, 1.0];

        let t = time::precise_time_s();
        let mut y = screen_area.max_y();
        // The log can be very long, and we're always most interested in the latest ones, so
        // do a backwards iteration with an early exist once we hit a sufficiently old item.
        for msg in self.lines.iter().rev().take_while(|m| m.expire_time_s > t) {
            // The split_line iterator can't be reversed, need to do a bit of caching here.
            let fragments = split_line(&msg.text,
                                       |c| font.0.char_width(c).unwrap_or(0.0),
                                       screen_area.size.width)
                                .map(|x| x.to_string())
                                .collect::<Vec<String>>();
            for line in fragments.iter().rev() {
                context.ui.draw_text(&font.0, Point2D::new(0.0, y), color, line);
                y -= font.0.height;
            }
        }
    }

    /// Draw the console as a big drop-down with a command prompt.
    pub fn draw_large(&mut self, context: &mut backend::Context, screen_area: &Rect<f32>) {
        // TODO: Store default font in context object
        let font: Resource<backend::Font> = Resource::new("default".to_string()).unwrap();
        // TODO: Ditto for color
        let color = [0.6, 0.6, 0.6, 1.0];
        let background = [0.0, 0.0, 0.6, 1.0];

        context.ui.fill_rect(*screen_area, background);

        let mut lines_left = (screen_area.size.height / font.0.height).ceil() as i32;

        let mut y = screen_area.max_y();

        // TODO: Handle enter with text input.
        // TODO: Command history.
        context.ui.text_input(&font.0, Point2D::new(0.0, y), color, &mut self.input_buffer);
        y -= font.0.height;
        lines_left -= 1;

        for msg in self.lines.iter().rev() {
            // XXX: Duplicated from draw_small.
            let fragments = split_line(&msg.text,
                                       |c| font.0.char_width(c).unwrap_or(0.0),
                                       screen_area.size.width)
                                .map(|x| x.to_string())
                                .collect::<Vec<String>>();
            for line in fragments.iter().rev() {
                context.ui.draw_text(&font.0, Point2D::new(0.0, y), color, line);
                y -= font.0.height;
                lines_left -= 1;
            }

            if lines_left <= 0 {
                break;
            }
        }
    }

    fn end_message(&mut self) {
        let mut message_text = String::new();
        mem::swap(&mut message_text, &mut self.output_buffer);

        let now = time::precise_time_s();
        if now > self.done_reading_s {
            self.done_reading_s = now;
        }

        let message = Message::new(message_text, self.done_reading_s);
        assert!(message.expire_time_s >= self.done_reading_s);
        self.done_reading_s = message.expire_time_s;
        self.lines.push(message);
    }

    pub fn get_input(&mut self) -> String {
        let mut ret = String::new();
        mem::swap(&mut ret, &mut self.input_buffer);
        ret
    }
}

impl Write for Console {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let s = str::from_utf8(buf).expect("Wrote non-UTF-8 to Console");
        let mut lines = s.split('\n');
        lines.next().map(|text| self.output_buffer.push_str(text));

        for line in lines {
            self.end_message();
            self.output_buffer.push_str(line);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}
<commit_msg>Translucent console<commit_after>use std::io::prelude::*;
use std::io;
use std::str;
use std::mem;
use time;
use euclid::{Point2D, Rect};
use calx_resource::Resource;
use calx_alg::split_line;
use backend;

struct Message {
    expire_time_s: f64,
    text: String,
}

impl Message {
    fn new(text: String, time_start_s: f64) -> Message {
        const TIME_TO_READ_CHAR_S: f64 = 0.2;
        let expire_time_s = time_start_s + text.len() as f64 * TIME_TO_READ_CHAR_S;
        Message {
            expire_time_s: expire_time_s,
            text: text,
        }
    }
}

/// Output
pub struct Console {
    lines: Vec<Message>,
    input_buffer: String,
    output_buffer: String,
    done_reading_s: f64,
}

impl Console {
    pub fn new() -> Console {
        Console {
            lines: Vec::new(),
            input_buffer: String::new(),
            output_buffer: String::new(),
            done_reading_s: 0.0,
        }
    }

    /// Draw the console as a regular message display.
    pub fn draw_small(&mut self, context: &mut backend::Context, screen_area: &Rect<f32>) {
        // TODO: Store default font in context object
        let font: Resource<backend::Font> = Resource::new("default".to_string()).unwrap();
        // TODO: Ditto for color
        let color = [1.0, 1.0, 1.0, 1.0];

        let t = time::precise_time_s();
        let mut y = screen_area.max_y();
        // The log can be very long, and we're always most interested in the latest ones, so
        // do a backwards iteration with an early exist once we hit a sufficiently old item.
        for msg in self.lines.iter().rev().take_while(|m| m.expire_time_s > t) {
            // The split_line iterator can't be reversed, need to do a bit of caching here.
            let fragments = split_line(&msg.text,
                                       |c| font.0.char_width(c).unwrap_or(0.0),
                                       screen_area.size.width)
                                .map(|x| x.to_string())
                                .collect::<Vec<String>>();
            for line in fragments.iter().rev() {
                context.ui.draw_text(&font.0, Point2D::new(0.0, y), color, line);
                y -= font.0.height;
            }
        }
    }

    /// Draw the console as a big drop-down with a command prompt.
    pub fn draw_large(&mut self, context: &mut backend::Context, screen_area: &Rect<f32>) {
        // TODO: Store default font in context object
        let font: Resource<backend::Font> = Resource::new("default".to_string()).unwrap();
        // TODO: Ditto for color
        let color = [0.6, 0.6, 0.6, 1.0];
        let background = [0.0, 0.0, 0.6, 0.8];

        context.ui.fill_rect(*screen_area, background);

        let mut lines_left = (screen_area.size.height / font.0.height).ceil() as i32;

        let mut y = screen_area.max_y();

        // TODO: Handle enter with text input.
        // TODO: Command history.
        context.ui.text_input(&font.0, Point2D::new(0.0, y), color, &mut self.input_buffer);
        y -= font.0.height;
        lines_left -= 1;

        for msg in self.lines.iter().rev() {
            // XXX: Duplicated from draw_small.
            let fragments = split_line(&msg.text,
                                       |c| font.0.char_width(c).unwrap_or(0.0),
                                       screen_area.size.width)
                                .map(|x| x.to_string())
                                .collect::<Vec<String>>();
            for line in fragments.iter().rev() {
                context.ui.draw_text(&font.0, Point2D::new(0.0, y), color, line);
                y -= font.0.height;
                lines_left -= 1;
            }

            if lines_left <= 0 {
                break;
            }
        }
    }

    fn end_message(&mut self) {
        let mut message_text = String::new();
        mem::swap(&mut message_text, &mut self.output_buffer);

        let now = time::precise_time_s();
        if now > self.done_reading_s {
            self.done_reading_s = now;
        }

        let message = Message::new(message_text, self.done_reading_s);
        assert!(message.expire_time_s >= self.done_reading_s);
        self.done_reading_s = message.expire_time_s;
        self.lines.push(message);
    }

    pub fn get_input(&mut self) -> String {
        let mut ret = String::new();
        mem::swap(&mut ret, &mut self.input_buffer);
        ret
    }
}

impl Write for Console {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let s = str::from_utf8(buf).expect("Wrote non-UTF-8 to Console");
        let mut lines = s.split('\n');
        lines.next().map(|text| self.output_buffer.push_str(text));

        for line in lines {
            self.end_message();
            self.output_buffer.push_str(line);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}
<|endoftext|>


<commit_before>use std::io::BufRead;
// rustc incorrectly suggests `std::os::ext::process::CommandExt`.
// <https://github.com/rust-lang/rust/issues/39175>
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::str;

use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use diesel::types::*;
use libsqlite3_sys;
use regex::Regex;

use models;
use types::Error;

const MAGIC: &'static str = r#"__duiker_import() {
    local _histignore=$HISTIGNORE
    local _histtimeformat=$HISTTIMEFORMAT
    HISTIGNORE='history*'
    HISTTIMEFORMAT='%s '
    history 1 | ./target/debug/duiker import --quiet -
    HISTIGNORE=$_histignore
    HISTTIMEFORMAT=$_histtimeformat
}"#;

pub fn head(connection: &SqliteConnection, n: i64) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    let commands = history.order(timestamp.asc()).limit(n);
    Ok(commands.load::<models::History>(connection)?)
}

fn parse_history_line(line: &str) -> Result<models::NewCommand, Error> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"^\s*?(\d+\s+)(?P<timestamp>\d+)?\s+(?P<command>.*)").unwrap();
    }
    match RE.is_match(line) {
        true => {
            let caps = RE.captures(line).unwrap();
            let timestamp = caps.name("timestamp").unwrap().as_str().parse::<i32>().unwrap();
            let command = caps.name("command").unwrap().as_str().trim();
            Ok(models::NewCommand{timestamp: timestamp, command: &command})
        }
        _ => Err(Error::InvalidHistoryLine)
    }
}

pub fn import<'a>(connection: &SqliteConnection, reader: Box<BufRead + 'a>) -> Result<usize, Error> {
    use schema::history;
    let mut n: usize = 0;
    for line in reader.lines() {
        let buf = line?;
        let command = parse_history_line(&buf)?;
        n += diesel::insert(&command).into(history::table).execute(connection)?;
    }
    Ok(n)
}

pub fn log(connection: &SqliteConnection) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    Ok(history.load::<models::History>(connection)?)
}

pub fn magic() {
    println!("{}", MAGIC);
}

pub fn search(connection: &SqliteConnection, expression: &str) -> Result<Vec<models::History>, Error> {
    use diesel::expression::sql;
    let query = sql::<(Integer, Integer, Text)>("SELECT history.*
                                                   FROM fts_history
                                                   JOIN history
                                                     ON fts_history.history_id = history.id
                                                  WHERE fts_history MATCH ?");
    Ok(query.bind::<Text, _>(expression).load::<models::History>(connection)?)
}

pub fn sqlite3(sqlite3_options: Vec<&str>) {
    Command::new("sqlite3")
        .args(&sqlite3_options)
        .exec();
}

pub fn tail(connection: &SqliteConnection, n: i64) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    let commands = history.select(id)
        .order(timestamp.desc())
        .limit(n);
    let sorted = history.filter(id.eq_any(commands))
        .order(timestamp.asc());
    Ok(sorted.load::<models::History>(connection)?)
}

pub fn top(connection: &SqliteConnection, n: i64) -> Result<Vec<models::Frequency>, Error> {
    use diesel::expression::sql;
    let query = sql::<(Integer, Text)>("SELECT COUNT(*) AS frequency, command
                                          FROM history
                                         GROUP BY command
                                         ORDER BY frequency DESC
                                         LIMIT ?");
    Ok(query.bind::<BigInt, _>(n).load::<models::Frequency>(connection)?)
}

pub fn version(verbose: bool) {
    println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
    if verbose {
        println!("SQLite3 {}",
                 str::from_utf8(libsqlite3_sys::SQLITE_VERSION).unwrap());
    }
}
<commit_msg>Remove debug path from magic output<commit_after>use std::io::BufRead;
// rustc incorrectly suggests `std::os::ext::process::CommandExt`.
// <https://github.com/rust-lang/rust/issues/39175>
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::str;

use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use diesel::types::*;
use libsqlite3_sys;
use regex::Regex;

use models;
use types::Error;

const MAGIC: &'static str = r#"__duiker_import() {
    local _histignore=$HISTIGNORE
    local _histtimeformat=$HISTTIMEFORMAT
    HISTIGNORE='history*'
    HISTTIMEFORMAT='%s '
    history 1 | duiker import --quiet -
    HISTIGNORE=$_histignore
    HISTTIMEFORMAT=$_histtimeformat
}"#;

pub fn head(connection: &SqliteConnection, n: i64) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    let commands = history.order(timestamp.asc()).limit(n);
    Ok(commands.load::<models::History>(connection)?)
}

fn parse_history_line(line: &str) -> Result<models::NewCommand, Error> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"^\s*?(\d+\s+)(?P<timestamp>\d+)?\s+(?P<command>.*)").unwrap();
    }
    match RE.is_match(line) {
        true => {
            let caps = RE.captures(line).unwrap();
            let timestamp = caps.name("timestamp").unwrap().as_str().parse::<i32>().unwrap();
            let command = caps.name("command").unwrap().as_str().trim();
            Ok(models::NewCommand{timestamp: timestamp, command: &command})
        }
        _ => Err(Error::InvalidHistoryLine)
    }
}

pub fn import<'a>(connection: &SqliteConnection, reader: Box<BufRead + 'a>) -> Result<usize, Error> {
    use schema::history;
    let mut n: usize = 0;
    for line in reader.lines() {
        let buf = line?;
        let command = parse_history_line(&buf)?;
        n += diesel::insert(&command).into(history::table).execute(connection)?;
    }
    Ok(n)
}

pub fn log(connection: &SqliteConnection) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    Ok(history.load::<models::History>(connection)?)
}

pub fn magic() {
    println!("{}", MAGIC);
}

pub fn search(connection: &SqliteConnection, expression: &str) -> Result<Vec<models::History>, Error> {
    use diesel::expression::sql;
    let query = sql::<(Integer, Integer, Text)>("SELECT history.*
                                                   FROM fts_history
                                                   JOIN history
                                                     ON fts_history.history_id = history.id
                                                  WHERE fts_history MATCH ?");
    Ok(query.bind::<Text, _>(expression).load::<models::History>(connection)?)
}

pub fn sqlite3(sqlite3_options: Vec<&str>) {
    Command::new("sqlite3")
        .args(&sqlite3_options)
        .exec();
}

pub fn tail(connection: &SqliteConnection, n: i64) -> Result<Vec<models::History>, Error> {
    use schema::history::dsl::*;
    let commands = history.select(id)
        .order(timestamp.desc())
        .limit(n);
    let sorted = history.filter(id.eq_any(commands))
        .order(timestamp.asc());
    Ok(sorted.load::<models::History>(connection)?)
}

pub fn top(connection: &SqliteConnection, n: i64) -> Result<Vec<models::Frequency>, Error> {
    use diesel::expression::sql;
    let query = sql::<(Integer, Text)>("SELECT COUNT(*) AS frequency, command
                                          FROM history
                                         GROUP BY command
                                         ORDER BY frequency DESC
                                         LIMIT ?");
    Ok(query.bind::<BigInt, _>(n).load::<models::Frequency>(connection)?)
}

pub fn version(verbose: bool) {
    println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
    if verbose {
        println!("SQLite3 {}",
                 str::from_utf8(libsqlite3_sys::SQLITE_VERSION).unwrap());
    }
}
<|endoftext|>