<commit_before>use std::io::{self, Read, Seek, SeekFrom, BufReader};
use std::fs::{File, Permissions};
use std::fs::{create_dir_all, set_permissions, hard_link};
use std::os::unix::fs::{PermissionsExt, symlink};
use std::path::{Path, PathBuf};

use tar::Archive;
use flate2::read::GzDecoder;
use xz2::read::XzDecoder;
use bzip2::read::BzDecoder;
use libmount::BindMount;
use quire::validate as V;

use capsule::download::{maybe_download_and_check_hashsum};
use container::mount::{unmount};
use builder::context::Context;
use builder::commands::generic::run_command_at;
use builder::dns::revert_name_files;
use file_util::{Dir, read_visible_entries, copy_stream, set_owner_group};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

#[derive(Deserialize, Debug)]
pub struct Tar {
    pub url: String,
    pub sha256: Option<String>,
    pub path: PathBuf,
    pub subdir: PathBuf,
}

impl Tar {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("path", V::Directory::new().absolute(true).default("/"))
        .member("subdir", V::Directory::new().default("").absolute(false))
    }
}

#[derive(Deserialize, Debug)]
pub struct TarInstall {
    pub url: String,
    pub sha256: Option<String>,
    pub subdir: Option<PathBuf>,
    pub script: String,
}

impl TarInstall {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("subdir", V::Directory::new().optional().absolute(false))
        .member("script", V::Scalar::new()
                .default("./configure --prefix=/usr\n\
                          make\n\
                          make install\n"))
    }
}

pub fn unpack_file(_ctx: &mut Context, src: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{

    info!("Unpacking {:?} -> {:?}", src, tgt);
    let read_err = |e| format!("Error reading {:?}: {}", src, e);

    let mut file = BufReader::new(File::open(&src).map_err(&read_err)?);

    let mut buf = [0u8; 8];
    let nbytes = file.read(&mut buf).map_err(&read_err)?;
    file.seek(SeekFrom::Start(0)).map_err(&read_err)?;
    let magic = &buf[..nbytes];
    if magic.len() >= 2 && magic[..2] == [0x1f, 0x8b] {
        return unpack_stream(GzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 6 && magic[..6] ==
        [ 0xFD, b'7', b'z', b'X', b'Z', 0x00]
    {
        return unpack_stream(XzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 3 && magic[..3] == [ b'B', b'Z', b'h'] {
        return unpack_stream(BzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else {
        return Err(format!("unpacking {:?}: unexpected compression", src));
    }

}

fn unpack_stream<F: Read>(file: F, srcpath: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{
    let read_err = |e| format!("Error reading {:?}: {}", srcpath, e);
    let mut arc = Archive::new(file);
    let mut hardlinks = Vec::new();

    for item in arc.entries().map_err(&read_err)? {
        let mut src = item.map_err(&read_err)?;
        let path_ref = src.path().map_err(&read_err)?
            .to_path_buf();
        let mut orig_path: &Path = &path_ref;
        if orig_path.is_absolute() {
            orig_path = orig_path.strip_prefix("/").unwrap();
        }
        if includes.len() > 0 {
            if !includes.iter().any(|x| orig_path.starts_with(x)) {
                continue;
            }
        }
        if excludes.iter().any(|x| orig_path.starts_with(x)) {
            continue;
        }
        let path = tgt.join(orig_path);
        let write_err = |e| format!("Error writing {:?}: {}", path, e);
        let entry = src.header().entry_type();

        // Some archives don't have uids
        // TODO(tailhook) should this be handled in tar-rs?
        let uid = src.header().uid().unwrap_or(0);
        let gid = src.header().gid().unwrap_or(0);

        if entry.is_dir() {
            let mode = src.header().mode().map_err(&read_err)?;
            let mut dir_builder = Dir::new(&path);
            dir_builder.recursive(true).mode(mode);
            if preserve_owner {
                dir_builder.uid(uid).gid(gid);
            }
            dir_builder.create().map_err(&write_err)?;
        } else if entry.is_symlink() {
            let src = src.header().link_name().map_err(&read_err)?
                .ok_or(format!("Error unpacking {:?}, broken symlink", path))?;
            match symlink(&src, &path) {
                Ok(_) => {},
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            Dir::new(parent).recursive(true).create()
                                .map_err(&write_err)?;
                            symlink(&src, &path).map_err(&write_err)?
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
        } else if entry.is_hard_link() {
            let link = src.link_name().map_err(&read_err)?
                .ok_or(format!("Error unpacking {:?}, broken symlink", path))?;
            let link = if link.is_absolute() {
                link.strip_prefix("/").unwrap()
            } else {
                &*link
            };
            hardlinks.push((tgt.join(link).to_path_buf(), path.to_path_buf()));
        } else if entry.is_pax_global_extensions() ||
                  entry.is_pax_local_extensions() ||
                  entry.is_gnu_longname() ||
                  entry.is_gnu_longlink()
        {
            // nothing to do
        } else {
            let mut dest = match File::create(&path) {
                Ok(x) => x,
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            Dir::new(parent).recursive(true).create()
                                .map_err(&write_err)?;
                            File::create(&path).map_err(&write_err)?
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
            copy_stream(&mut src, &mut dest).map_err(|e|
                format!("Error unpacking {:?} -> {:?}: {}",
                        srcpath, path, e))?;
            let mode = src.header().mode().map_err(&read_err)?;
            set_permissions(&path, Permissions::from_mode(mode))
                .map_err(&write_err)?;
            if preserve_owner {
                set_owner_group(&path, uid, gid).map_err(&write_err)?;
            }
        }
    }
    for (src, dst) in hardlinks.into_iter() {
        let write_err = |e| {
            format!("Error hardlinking {:?} - {:?}: {}", &src, &dst, e)
        };
        match hard_link(&src, &dst) {
            Ok(_) => {},
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    if let Some(parent) = dst.parent() {
                        Dir::new(parent).recursive(true).create()
                            .map_err(&write_err)?;
                        hard_link(&src, &dst).map_err(&write_err)?
                    } else {
                        return Err(write_err(e));
                    }
                } else {
                    return Err(write_err(e));
                }
            }
        };
    }
    Ok(())
}

pub fn unpack_subdir(ctx: &mut Context, filename: &Path, dest: &Path,
    subdir: &Path)
    -> Result<(), String>
{
    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    let tmpsub = tmppath.join(subdir);
    try_msg!(Dir::new(&tmpsub).recursive(true).create(),
        "Error making dir: {err}");
    if !dest.exists() {
        try_msg!(Dir::new(&dest).recursive(true).create(),
            "Error making dir: {err}");
    }
    try_msg!(BindMount::new(&dest, &tmpsub).mount(),
        "temporary tar mount: {err}");
    let res = if subdir == Path::new("") {
        unpack_file(ctx, &filename, &tmppath, &[], &[], false)
    } else {
        unpack_file(ctx, &filename, &tmppath, &[subdir], &[], false)
    };
    unmount(&tmpsub)?;
    res
}

pub fn tar_command(ctx: &mut Context, tar: &Tar) -> Result<(), String>
{
    let fpath = PathBuf::from("/vagga/root")
        .join(tar.path.strip_prefix("/").unwrap());
    let (filename, _) = maybe_download_and_check_hashsum(
        &mut ctx.capsule, &tar.url, tar.sha256.clone(), false)?;

    if &Path::new(&tar.subdir) == &Path::new(".") {
        unpack_file(ctx, &filename, &fpath, &[], &[], false)?;
    } else {
        unpack_subdir(ctx, &filename, &fpath, &tar.subdir)?;
    }
    if tar.path == Path::new("/") {
        revert_name_files()?;
    }
    Ok(())
}

pub fn tar_install(ctx: &mut Context, tar: &TarInstall)
    -> Result<(), String>
{
    let (filename, _) = maybe_download_and_check_hashsum(
        &mut ctx.capsule, &tar.url, tar.sha256.clone(), false)?;

    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    create_dir_all(&tmppath)
         .map_err(|e| format!("Error making dir: {}", e))?;
    set_permissions(&tmppath, Permissions::from_mode(0o755))
         .map_err(|e| format!("Error setting permissions: {}", e))?;
    unpack_file(ctx, &filename, &tmppath, &[], &[], false)?;
    let workdir = if let Some(ref subpath) = tar.subdir {
        tmppath.join(subpath)
    } else {
        let items = read_visible_entries(&tmppath)
            .map_err(|e| format!("Error reading dir: {}", e))?;
        if items.len() != 1 {
            if items.len() == 0 {
                return Err("Tar archive was empty".to_string());
            } else {
                return Err("Multiple directories was unpacked. \
                    If thats expected use `subdir: \".\"` or any \
                    other directory".to_string());
            }
        }
        items.into_iter().next().unwrap()
    };
    let workdir = PathBuf::from("/").join(
        workdir.strip_prefix("/vagga/root").unwrap());
    return run_command_at(ctx, &[
        "/bin/sh".to_string(),
        "-exc".to_string(),
        tar.script.to_string()],
        &workdir);
}

impl BuildStep for Tar {
    fn name(&self) -> &'static str { "Tar" }
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.field("path", &self.path);
        hash.field("subdir", &self.subdir);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            tar_command(&mut guard.ctx, self)?;
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}

impl BuildStep for TarInstall {
    fn name(&self) -> &'static str { "TarInstall" }
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.opt_field("subdir", &self.subdir);
        hash.field("script", &self.script);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            tar_install(&mut guard.ctx, self)?;
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}
<commit_msg>prevents truncation of long symbolic links when unpacking<commit_after>use std::io::{self, Read, Seek, SeekFrom, BufReader};
use std::fs::{File, Permissions};
use std::fs::{create_dir_all, set_permissions, hard_link};
use std::os::unix::fs::{PermissionsExt, symlink};
use std::path::{Path, PathBuf};

use tar::Archive;
use flate2::read::GzDecoder;
use xz2::read::XzDecoder;
use bzip2::read::BzDecoder;
use libmount::BindMount;
use quire::validate as V;

use capsule::download::{maybe_download_and_check_hashsum};
use container::mount::{unmount};
use builder::context::Context;
use builder::commands::generic::run_command_at;
use builder::dns::revert_name_files;
use file_util::{Dir, read_visible_entries, copy_stream, set_owner_group};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

#[derive(Deserialize, Debug)]
pub struct Tar {
    pub url: String,
    pub sha256: Option<String>,
    pub path: PathBuf,
    pub subdir: PathBuf,
}

impl Tar {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("path", V::Directory::new().absolute(true).default("/"))
        .member("subdir", V::Directory::new().default("").absolute(false))
    }
}

#[derive(Deserialize, Debug)]
pub struct TarInstall {
    pub url: String,
    pub sha256: Option<String>,
    pub subdir: Option<PathBuf>,
    pub script: String,
}

impl TarInstall {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("subdir", V::Directory::new().optional().absolute(false))
        .member("script", V::Scalar::new()
                .default("./configure --prefix=/usr\n\
                          make\n\
                          make install\n"))
    }
}

pub fn unpack_file(_ctx: &mut Context, src: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{

    info!("Unpacking {:?} -> {:?}", src, tgt);
    let read_err = |e| format!("Error reading {:?}: {}", src, e);

    let mut file = BufReader::new(File::open(&src).map_err(&read_err)?);

    let mut buf = [0u8; 8];
    let nbytes = file.read(&mut buf).map_err(&read_err)?;
    file.seek(SeekFrom::Start(0)).map_err(&read_err)?;
    let magic = &buf[..nbytes];
    if magic.len() >= 2 && magic[..2] == [0x1f, 0x8b] {
        return unpack_stream(GzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 6 && magic[..6] ==
        [ 0xFD, b'7', b'z', b'X', b'Z', 0x00]
    {
        return unpack_stream(XzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 3 && magic[..3] == [ b'B', b'Z', b'h'] {
        return unpack_stream(BzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else {
        return Err(format!("unpacking {:?}: unexpected compression", src));
    }

}

fn unpack_stream<F: Read>(file: F, srcpath: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{
    let read_err = |e| format!("Error reading {:?}: {}", srcpath, e);
    let mut arc = Archive::new(file);
    let mut hardlinks = Vec::new();

    for item in arc.entries().map_err(&read_err)? {
        let mut src = item.map_err(&read_err)?;
        let path_ref = src.path().map_err(&read_err)?
            .to_path_buf();
        let mut orig_path: &Path = &path_ref;
        if orig_path.is_absolute() {
            orig_path = orig_path.strip_prefix("/").unwrap();
        }
        if includes.len() > 0 {
            if !includes.iter().any(|x| orig_path.starts_with(x)) {
                continue;
            }
        }
        if excludes.iter().any(|x| orig_path.starts_with(x)) {
            continue;
        }
        let path = tgt.join(orig_path);
        let write_err = |e| format!("Error writing {:?}: {}", path, e);
        let entry = src.header().entry_type();

        // Some archives don't have uids
        // TODO(tailhook) should this be handled in tar-rs?
        let uid = src.header().uid().unwrap_or(0);
        let gid = src.header().gid().unwrap_or(0);

        if entry.is_dir() {
            let mode = src.header().mode().map_err(&read_err)?;
            let mut dir_builder = Dir::new(&path);
            dir_builder.recursive(true).mode(mode);
            if preserve_owner {
                dir_builder.uid(uid).gid(gid);
            }
            dir_builder.create().map_err(&write_err)?;
        } else if entry.is_symlink() {
            let src = src.link_name().map_err(&read_err)?
                .ok_or(format!("Error unpacking {:?}, broken symlink", path))?;
            match symlink(&src, &path) {
                Ok(_) => {},
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            Dir::new(parent).recursive(true).create()
                                .map_err(&write_err)?;
                            symlink(&src, &path).map_err(&write_err)?
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
        } else if entry.is_hard_link() {
            let link = src.link_name().map_err(&read_err)?
                .ok_or(format!("Error unpacking {:?}, broken symlink", path))?;
            let link = if link.is_absolute() {
                link.strip_prefix("/").unwrap()
            } else {
                &*link
            };
            hardlinks.push((tgt.join(link).to_path_buf(), path.to_path_buf()));
        } else if entry.is_pax_global_extensions() ||
                  entry.is_pax_local_extensions() ||
                  entry.is_gnu_longname() ||
                  entry.is_gnu_longlink()
        {
            // nothing to do
        } else {
            let mut dest = match File::create(&path) {
                Ok(x) => x,
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            Dir::new(parent).recursive(true).create()
                                .map_err(&write_err)?;
                            File::create(&path).map_err(&write_err)?
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
            copy_stream(&mut src, &mut dest).map_err(|e|
                format!("Error unpacking {:?} -> {:?}: {}",
                        srcpath, path, e))?;
            let mode = src.header().mode().map_err(&read_err)?;
            set_permissions(&path, Permissions::from_mode(mode))
                .map_err(&write_err)?;
            if preserve_owner {
                set_owner_group(&path, uid, gid).map_err(&write_err)?;
            }
        }
    }
    for (src, dst) in hardlinks.into_iter() {
        let write_err = |e| {
            format!("Error hardlinking {:?} - {:?}: {}", &src, &dst, e)
        };
        match hard_link(&src, &dst) {
            Ok(_) => {},
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    if let Some(parent) = dst.parent() {
                        Dir::new(parent).recursive(true).create()
                            .map_err(&write_err)?;
                        hard_link(&src, &dst).map_err(&write_err)?
                    } else {
                        return Err(write_err(e));
                    }
                } else {
                    return Err(write_err(e));
                }
            }
        };
    }
    Ok(())
}

pub fn unpack_subdir(ctx: &mut Context, filename: &Path, dest: &Path,
    subdir: &Path)
    -> Result<(), String>
{
    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    let tmpsub = tmppath.join(subdir);
    try_msg!(Dir::new(&tmpsub).recursive(true).create(),
        "Error making dir: {err}");
    if !dest.exists() {
        try_msg!(Dir::new(&dest).recursive(true).create(),
            "Error making dir: {err}");
    }
    try_msg!(BindMount::new(&dest, &tmpsub).mount(),
        "temporary tar mount: {err}");
    let res = if subdir == Path::new("") {
        unpack_file(ctx, &filename, &tmppath, &[], &[], false)
    } else {
        unpack_file(ctx, &filename, &tmppath, &[subdir], &[], false)
    };
    unmount(&tmpsub)?;
    res
}

pub fn tar_command(ctx: &mut Context, tar: &Tar) -> Result<(), String>
{
    let fpath = PathBuf::from("/vagga/root")
        .join(tar.path.strip_prefix("/").unwrap());
    let (filename, _) = maybe_download_and_check_hashsum(
        &mut ctx.capsule, &tar.url, tar.sha256.clone(), false)?;

    if &Path::new(&tar.subdir) == &Path::new(".") {
        unpack_file(ctx, &filename, &fpath, &[], &[], false)?;
    } else {
        unpack_subdir(ctx, &filename, &fpath, &tar.subdir)?;
    }
    if tar.path == Path::new("/") {
        revert_name_files()?;
    }
    Ok(())
}

pub fn tar_install(ctx: &mut Context, tar: &TarInstall)
    -> Result<(), String>
{
    let (filename, _) = maybe_download_and_check_hashsum(
        &mut ctx.capsule, &tar.url, tar.sha256.clone(), false)?;

    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    create_dir_all(&tmppath)
         .map_err(|e| format!("Error making dir: {}", e))?;
    set_permissions(&tmppath, Permissions::from_mode(0o755))
         .map_err(|e| format!("Error setting permissions: {}", e))?;
    unpack_file(ctx, &filename, &tmppath, &[], &[], false)?;
    let workdir = if let Some(ref subpath) = tar.subdir {
        tmppath.join(subpath)
    } else {
        let items = read_visible_entries(&tmppath)
            .map_err(|e| format!("Error reading dir: {}", e))?;
        if items.len() != 1 {
            if items.len() == 0 {
                return Err("Tar archive was empty".to_string());
            } else {
                return Err("Multiple directories was unpacked. \
                    If thats expected use `subdir: \".\"` or any \
                    other directory".to_string());
            }
        }
        items.into_iter().next().unwrap()
    };
    let workdir = PathBuf::from("/").join(
        workdir.strip_prefix("/vagga/root").unwrap());
    return run_command_at(ctx, &[
        "/bin/sh".to_string(),
        "-exc".to_string(),
        tar.script.to_string()],
        &workdir);
}

impl BuildStep for Tar {
    fn name(&self) -> &'static str { "Tar" }
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.field("path", &self.path);
        hash.field("subdir", &self.subdir);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            tar_command(&mut guard.ctx, self)?;
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}

impl BuildStep for TarInstall {
    fn name(&self) -> &'static str { "TarInstall" }
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.opt_field("subdir", &self.subdir);
        hash.field("script", &self.script);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            tar_install(&mut guard.ctx, self)?;
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bit_vec;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bit_vec::BitVec;
//!
//! let mut bits = BitVec::from_elem(64, false);
//! bits.set(5, true);
//! bits.set(32, true);
//! let rs = RankSelect::new(&bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use bit_vec::BitVec;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: Vec<u8>,
    superblocks_1: Vec<u32>,
    superblocks_0: Vec<u32>,
    s: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: &BitVec, k: usize) -> RankSelect {
        let n = bits.len();
        let bits = bits.to_bytes();
        let s = k * 32;

        RankSelect {
            n,
            s,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &[u8] {
        &self.bits
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: usize) -> Option<u32> {
        if i >= self.n {
            None
        } else {
            let s = i / self.s; // the superblock
            let b = i / 8; // the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s];
            // add the rank within the block
            rank += (self.bits[b] >> (7 - i % 8)).count_ones();
            // add the popcounts of blocks in between
            rank += self.bits[s * 32 / 8..b]
                .iter()
                .map(|&a| a.count_ones())
                .sum::<u32>();

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: usize) -> Option<u32> {
        self.rank_1(i).map(|r| (i as u32 + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: usize) -> Option<u32> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u32) -> Option<usize> {
        let mut superblock = match self.superblocks_1.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_1[superblock];

        let first_block = superblock * self.s / 8;
        for (block, &b) in self.bits[first_block..].iter().enumerate() {
            let p = b.count_ones();
            if rank + p >= j {
                let mut bit = 0b1000_0000;
                for i in 0..8usize {
                    rank += (b & bit > 0) as u32;
                    if rank == j {
                        return Some((first_block + block) * 8 + i);
                    }
                    bit >>= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u32) -> Option<usize> {
        let mut superblock = match self.superblocks_0.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_0[superblock];

        let first_block = superblock * self.s / 8;
        for (block, &b) in self.bits[first_block..].iter().enumerate() {
            let p = b.count_zeros();
            if rank + p >= j {
                let mut bit = 0b1000_0000;
                for i in 0..8usize {
                    rank += (b & bit == 0) as u32;
                    if rank == j {
                        return Some((first_block + block) * 8 + i);
                    }
                    bit >>= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u32) -> Option<usize> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, raw_bits: &[u8]) -> Vec<u32> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u32 = 0;
    let mut i = 0;
    for &b in raw_bits.iter() {
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank +=
            if t { b.count_ones() } else { b.count_zeros() };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bit_vec::BitVec;

    #[test]
    fn test_rank_select() {
        let mut bits = BitVec::from_elem(64, false);
        bits.set(5, true);
        bits.set(32, true);
        let rs = RankSelect::new(&bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
    }
}
<commit_msg>Allow to get the i-th bit.<commit_after>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bit_vec;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bit_vec::BitVec;
//!
//! let mut bits = BitVec::from_elem(64, false);
//! bits.set(5, true);
//! bits.set(32, true);
//! let rs = RankSelect::new(&bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use bit_vec::BitVec;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: Vec<u8>,
    superblocks_1: Vec<u32>,
    superblocks_0: Vec<u32>,
    s: usize,
    k: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: &BitVec, k: usize) -> RankSelect {
        let n = bits.len();
        let bits = bits.to_bytes();
        let s = k * 32;

        RankSelect {
            n,
            s,
            k,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Return the used k (see `RankSelect::new()`).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Get i-th bit. Returns `None` if i is bigger than the number of bits.
    pub fn get(&self, i: usize) -> Option<bool> {
        if i >= self.bits.len() * 8 {
            None
        } else {
            let one = 1 << (7 - (i % 8));
            Some((self.bits.get(i / 8).unwrap() & one) != 0)
        }
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &[u8] {
        &self.bits
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: usize) -> Option<u32> {
        if i >= self.n {
            None
        } else {
            let s = i / self.s; // the superblock
            let b = i / 8; // the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s];
            // add the rank within the block
            rank += (self.bits[b] >> (7 - i % 8)).count_ones();
            // add the popcounts of blocks in between
            rank += self.bits[s * 32 / 8..b]
                .iter()
                .map(|&a| a.count_ones())
                .sum::<u32>();

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: usize) -> Option<u32> {
        self.rank_1(i).map(|r| (i as u32 + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: usize) -> Option<u32> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u32) -> Option<usize> {
        let mut superblock = match self.superblocks_1.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_1[superblock];

        let first_block = superblock * self.s / 8;
        for (block, &b) in self.bits[first_block..].iter().enumerate() {
            let p = b.count_ones();
            if rank + p >= j {
                let mut bit = 0b1000_0000;
                for i in 0..8usize {
                    rank += (b & bit > 0) as u32;
                    if rank == j {
                        return Some((first_block + block) * 8 + i);
                    }
                    bit >>= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u32) -> Option<usize> {
        let mut superblock = match self.superblocks_0.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_0[superblock];

        let first_block = superblock * self.s / 8;
        for (block, &b) in self.bits[first_block..].iter().enumerate() {
            let p = b.count_zeros();
            if rank + p >= j {
                let mut bit = 0b1000_0000;
                for i in 0..8usize {
                    rank += (b & bit == 0) as u32;
                    if rank == j {
                        return Some((first_block + block) * 8 + i);
                    }
                    bit >>= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u32) -> Option<usize> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, raw_bits: &[u8]) -> Vec<u32> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u32 = 0;
    let mut i = 0;
    for &b in raw_bits.iter() {
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank +=
            if t { b.count_ones() } else { b.count_zeros() };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bit_vec::BitVec;

    #[test]
    fn test_rank_select() {
        let mut bits = BitVec::from_elem(64, false);
        bits.set(5, true);
        bits.set(32, true);
        let rs = RankSelect::new(&bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
        assert_eq!(rs.get(5).unwrap(), true);
        assert_eq!(rs.get(1).unwrap(), false);
        assert_eq!(rs.get(32).unwrap(), true);
    }
}
<|endoftext|>


<commit_before>use std::slice;
use std::fmt;
use std::ops::Index;

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVRange {
    Limited,
    Full
}

impl fmt::Display for YUVRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            YUVRange::Limited => write!(f, "Limited range"),
            YUVRange::Full => write!(f, "Full range")
        }
    }
}

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVSystem {
    YCbCr(YUVRange),
    YCoCg,
    ICtCp
}

impl fmt::Display for YUVSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::YUVSystem::*;
        match *self {
            YCbCr(range) => write!(f, "YCbCr ({})", range),
            YCoCg => write!(f, "YCbCg"),
            ICtCp => write!(f, "ICtCp"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum TrichromaticEncodingSystem {
    RGB,
    YUV(YUVSystem),
    XYZ
}

impl fmt::Display for TrichromaticEncodingSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::TrichromaticEncodingSystem::*;
        match *self {
            YUV(system) => write!(f, "{}", system),
            RGB => write!(f, "RGB"),
            XYZ => write!(f, "XYZ"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum ColorModel {
    Trichromatic(TrichromaticEncodingSystem),
    CMYK,
    HSV,
    LAB,
}

impl fmt::Display for ColorModel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::ColorModel::*;
        match *self {
            Trichromatic(system) => write!(f, "{}", system),
            CMYK => write!(f, "CMYK"),
            HSV => write!(f, "HSV"),
            LAB => write!(f, "LAB"),
        }
    }
}

impl ColorModel {
    pub fn get_default_components(&self) -> usize {
        match *self {
            ColorModel::CMYK => 4,
            _ => 3,
        }
    }
}

#[derive(Clone,Copy,Debug,PartialEq)]
pub struct Chromaton {
    h_ss: u8,
    v_ss: u8,
    packed: bool,
    depth: u8,
    shift: u8,
    comp_offs: u8,
    next_elem: u8,
}

fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

impl Chromaton {
    pub fn get_subsampling(&self) -> (u8, u8) {
        (self.h_ss, self.v_ss)
    }
    pub fn is_packed(&self) -> bool {
        self.packed
    }
    pub fn get_depth(&self) -> u8 {
        self.depth
    }
    pub fn get_shift(&self) -> u8 {
        self.shift
    }
    pub fn get_offset(&self) -> u8 {
        self.comp_offs
    }
    pub fn get_step(&self) -> u8 {
        self.next_elem
    }

    pub fn get_width(&self, width: usize) -> usize {
        (width + ((1 << self.h_ss) - 1)) >> self.h_ss
    }
    pub fn get_height(&self, height: usize) -> usize {
        (height + ((1 << self.v_ss) - 1)) >> self.v_ss
    }
    pub fn get_linesize(&self, width: usize, alignment: usize) -> usize {
        let d = self.depth as usize;
        align((self.get_width(width) * d + d - 1) >> 3, alignment)
    }
    pub fn get_data_size(&self, width: usize, height: usize, align: usize) -> usize {
        let nh = (height + ((1 << self.v_ss) - 1)) >> self.v_ss;
        self.get_linesize(width, align) * nh
    }
}

impl fmt::Display for Chromaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pfmt = if self.packed {
            let mask = ((1 << self.depth) - 1) << self.shift;
            format!("packed(+{},{:X}, step {})",
                    self.comp_offs,
                    mask,
                    self.next_elem)
        } else {
            format!("planar({},{})", self.comp_offs, self.next_elem)
        };
        write!(f, "({}x{}, {})", self.h_ss, self.v_ss, pfmt)
    }
}

#[derive(Clone,Copy,PartialEq,Debug)]
pub struct Formaton {
    model: ColorModel,
    components: u8,
    comp_info: [Option<Chromaton>; 5],
    elem_size: u8,
    be: bool,
    alpha: bool,
    palette: bool,
}

bitflags! {
    pub flags Flags: u8 {
        const BE       = 0x01,
        const ALPHA    = 0x02,
        const PALETTE  = 0x04,
    }
}

impl Formaton {
    pub fn new(model: ColorModel, components: &[Chromaton], flags: Flags, elem_size: u8) -> Self {
        let be = flags.contains(BE);
        let alpha = flags.contains(ALPHA);
        let palette = flags.contains(PALETTE);
        let mut c: [Option<Chromaton>; 5] = [None; 5];

        if components.len() > 5 {
            panic!("too many components");
        }

        for (i, v) in components.iter().enumerate() {
            c[i] = Some(*v);
        }

        Formaton {
            model: model,
            components: components.len() as u8,
            comp_info: c,
            elem_size: elem_size,
            be: be,
            alpha: alpha,
            palette: palette,
        }
    }

    pub fn get_model(&self) -> ColorModel {
        self.model
    }
    pub fn get_num_comp(&self) -> usize {
        self.components as usize
    }
    pub fn get_chromaton(&self, idx: usize) -> Option<Chromaton> {
        if idx < self.comp_info.len() {
            return self.comp_info[idx];
        }
        None
    }
    pub fn is_be(&self) -> bool {
        self.be
    }
    pub fn has_alpha(&self) -> bool {
        self.alpha
    }
    pub fn is_paletted(&self) -> bool {
        self.palette
    }
    pub fn get_elem_size(&self) -> u8 {
        self.elem_size
    }
}

impl<'a> Index<usize> for &'a Formaton {
    type Output = Option<Chromaton>;

    fn index(&self, index: usize) -> &Self::Output {
        self.comp_info.index(index)
    }
}

impl<'a> IntoIterator for &'a Formaton {
    type Item = &'a Option<Chromaton>;
    type IntoIter = slice::Iter<'a, Option<Chromaton>>;

    fn into_iter(self) -> Self::IntoIter {
        self.comp_info.iter()
    }
}

impl fmt::Display for Formaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let end = if self.be {
            "BE"
        } else {
            "LE"
        };
        let palstr = if self.palette {
            "palette "
        } else {
            ""
        };
        let astr = if self.alpha {
            "alpha "
        } else {
            ""
        };
        let mut str = format!("Formaton for {} ({}{}elem {} size {}): ",
                              self.model,
                              palstr,
                              astr,
                              end,
                              self.elem_size);
        for &i in self.into_iter() {
            if let Some(chr) = i {
                str = format!("{} {}", str, chr);
            }
        }
        write!(f, "[{}]", str)
    }
}

macro_rules! chromaton {
    ($hs: expr, $vs: expr, $pck: expr, $d: expr, $sh: expr, $co: expr, $ne: expr) => ({
        Some(Chromaton {
                h_ss: $hs,
                v_ss: $vs,
                packed: $pck,
                depth: $d,
                shift: $sh,
                comp_offs: $co,
                next_elem: $ne })
    });
    (yuv8; $hs: expr, $vs: expr, $co: expr) => ({
        chromaton!($hs, $vs, false, 8, 0, $co, 1)
    });
    (packrgb; $d: expr, $s: expr, $co: expr, $ne: expr) => ({
        chromaton!(0, 0, true, $d, $s, $co, $ne)
    });
    (pal8; $co: expr) => ({
        chromaton!(0, 0, true, 8, 0, $co, 3)
    });
}

pub mod formats {
    use pixel::*;
    use self::ColorModel::*;
    use self::TrichromaticEncodingSystem::*;
    use self::YUVSystem::*;
    use self::YUVRange::*;

    pub const YUV420: Formaton = Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 1, 1, 1),
                    chromaton!(yuv8; 1, 1, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const YUV410: Formaton = Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 2, 2, 1),
                    chromaton!(yuv8; 2, 2, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const PAL8: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(pal8; 0), chromaton!(pal8; 1), chromaton!(pal8; 2), None, None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: true,
    };

    pub const RGB565: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 5, 11, 0, 2),
                    chromaton!(packrgb; 6,  5, 0, 2),
                    chromaton!(packrgb; 5,  0, 0, 2),
                    None,
                    None],
        elem_size: 2,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const RGB24: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 8, 0, 2, 3),
                    chromaton!(packrgb; 8, 0, 1, 3),
                    chromaton!(packrgb; 8, 0, 0, 3),
                    None,
                    None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: false,
    };
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn fmt() {
        println!("formaton yuv- {}", formats::YUV420);
        println!("formaton pal- {}", formats::PAL8);
        println!("formaton rgb565- {}", formats::RGB565);
    }
}
<commit_msg>pixel: Use const ref for the well-known formats<commit_after>use std::slice;
use std::fmt;
use std::ops::Index;

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVRange {
    Limited,
    Full
}

impl fmt::Display for YUVRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            YUVRange::Limited => write!(f, "Limited range"),
            YUVRange::Full => write!(f, "Full range")
        }
    }
}

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVSystem {
    YCbCr(YUVRange),
    YCoCg,
    ICtCp
}

impl fmt::Display for YUVSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::YUVSystem::*;
        match *self {
            YCbCr(range) => write!(f, "YCbCr ({})", range),
            YCoCg => write!(f, "YCbCg"),
            ICtCp => write!(f, "ICtCp"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum TrichromaticEncodingSystem {
    RGB,
    YUV(YUVSystem),
    XYZ
}

impl fmt::Display for TrichromaticEncodingSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::TrichromaticEncodingSystem::*;
        match *self {
            YUV(system) => write!(f, "{}", system),
            RGB => write!(f, "RGB"),
            XYZ => write!(f, "XYZ"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum ColorModel {
    Trichromatic(TrichromaticEncodingSystem),
    CMYK,
    HSV,
    LAB,
}

impl fmt::Display for ColorModel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::ColorModel::*;
        match *self {
            Trichromatic(system) => write!(f, "{}", system),
            CMYK => write!(f, "CMYK"),
            HSV => write!(f, "HSV"),
            LAB => write!(f, "LAB"),
        }
    }
}

impl ColorModel {
    pub fn get_default_components(&self) -> usize {
        match *self {
            ColorModel::CMYK => 4,
            _ => 3,
        }
    }
}

#[derive(Clone,Copy,Debug,PartialEq)]
pub struct Chromaton {
    h_ss: u8,
    v_ss: u8,
    packed: bool,
    depth: u8,
    shift: u8,
    comp_offs: u8,
    next_elem: u8,
}

fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

impl Chromaton {
    pub fn get_subsampling(&self) -> (u8, u8) {
        (self.h_ss, self.v_ss)
    }
    pub fn is_packed(&self) -> bool {
        self.packed
    }
    pub fn get_depth(&self) -> u8 {
        self.depth
    }
    pub fn get_shift(&self) -> u8 {
        self.shift
    }
    pub fn get_offset(&self) -> u8 {
        self.comp_offs
    }
    pub fn get_step(&self) -> u8 {
        self.next_elem
    }

    pub fn get_width(&self, width: usize) -> usize {
        (width + ((1 << self.h_ss) - 1)) >> self.h_ss
    }
    pub fn get_height(&self, height: usize) -> usize {
        (height + ((1 << self.v_ss) - 1)) >> self.v_ss
    }
    pub fn get_linesize(&self, width: usize, alignment: usize) -> usize {
        let d = self.depth as usize;
        align((self.get_width(width) * d + d - 1) >> 3, alignment)
    }
    pub fn get_data_size(&self, width: usize, height: usize, align: usize) -> usize {
        let nh = (height + ((1 << self.v_ss) - 1)) >> self.v_ss;
        self.get_linesize(width, align) * nh
    }
}

impl fmt::Display for Chromaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pfmt = if self.packed {
            let mask = ((1 << self.depth) - 1) << self.shift;
            format!("packed(+{},{:X}, step {})",
                    self.comp_offs,
                    mask,
                    self.next_elem)
        } else {
            format!("planar({},{})", self.comp_offs, self.next_elem)
        };
        write!(f, "({}x{}, {})", self.h_ss, self.v_ss, pfmt)
    }
}

#[derive(Clone,Copy,PartialEq,Debug)]
pub struct Formaton {
    model: ColorModel,
    components: u8,
    comp_info: [Option<Chromaton>; 5],
    elem_size: u8,
    be: bool,
    alpha: bool,
    palette: bool,
}

bitflags! {
    pub flags Flags: u8 {
        const BE       = 0x01,
        const ALPHA    = 0x02,
        const PALETTE  = 0x04,
    }
}

impl Formaton {
    pub fn new(model: ColorModel, components: &[Chromaton], flags: Flags, elem_size: u8) -> Self {
        let be = flags.contains(BE);
        let alpha = flags.contains(ALPHA);
        let palette = flags.contains(PALETTE);
        let mut c: [Option<Chromaton>; 5] = [None; 5];

        if components.len() > 5 {
            panic!("too many components");
        }

        for (i, v) in components.iter().enumerate() {
            c[i] = Some(*v);
        }

        Formaton {
            model: model,
            components: components.len() as u8,
            comp_info: c,
            elem_size: elem_size,
            be: be,
            alpha: alpha,
            palette: palette,
        }
    }

    pub fn get_model(&self) -> ColorModel {
        self.model
    }
    pub fn get_num_comp(&self) -> usize {
        self.components as usize
    }
    pub fn get_chromaton(&self, idx: usize) -> Option<Chromaton> {
        if idx < self.comp_info.len() {
            return self.comp_info[idx];
        }
        None
    }
    pub fn is_be(&self) -> bool {
        self.be
    }
    pub fn has_alpha(&self) -> bool {
        self.alpha
    }
    pub fn is_paletted(&self) -> bool {
        self.palette
    }
    pub fn get_elem_size(&self) -> u8 {
        self.elem_size
    }
}

impl<'a> Index<usize> for &'a Formaton {
    type Output = Option<Chromaton>;

    fn index(&self, index: usize) -> &Self::Output {
        self.comp_info.index(index)
    }
}

impl<'a> IntoIterator for &'a Formaton {
    type Item = &'a Option<Chromaton>;
    type IntoIter = slice::Iter<'a, Option<Chromaton>>;

    fn into_iter(self) -> Self::IntoIter {
        self.comp_info.iter()
    }
}

impl fmt::Display for Formaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let end = if self.be {
            "BE"
        } else {
            "LE"
        };
        let palstr = if self.palette {
            "palette "
        } else {
            ""
        };
        let astr = if self.alpha {
            "alpha "
        } else {
            ""
        };
        let mut str = format!("Formaton for {} ({}{}elem {} size {}): ",
                              self.model,
                              palstr,
                              astr,
                              end,
                              self.elem_size);
        for &i in self.into_iter() {
            if let Some(chr) = i {
                str = format!("{} {}", str, chr);
            }
        }
        write!(f, "[{}]", str)
    }
}

macro_rules! chromaton {
    ($hs: expr, $vs: expr, $pck: expr, $d: expr, $sh: expr, $co: expr, $ne: expr) => ({
        Some(Chromaton {
                h_ss: $hs,
                v_ss: $vs,
                packed: $pck,
                depth: $d,
                shift: $sh,
                comp_offs: $co,
                next_elem: $ne })
    });
    (yuv8; $hs: expr, $vs: expr, $co: expr) => ({
        chromaton!($hs, $vs, false, 8, 0, $co, 1)
    });
    (packrgb; $d: expr, $s: expr, $co: expr, $ne: expr) => ({
        chromaton!(0, 0, true, $d, $s, $co, $ne)
    });
    (pal8; $co: expr) => ({
        chromaton!(0, 0, true, 8, 0, $co, 3)
    });
}

pub mod formats {
    use pixel::*;
    use self::ColorModel::*;
    use self::TrichromaticEncodingSystem::*;
    use self::YUVSystem::*;
    use self::YUVRange::*;

    pub const YUV420: &Formaton = &Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 1, 1, 1),
                    chromaton!(yuv8; 1, 1, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const YUV410: &Formaton = &Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 2, 2, 1),
                    chromaton!(yuv8; 2, 2, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const PAL8: &Formaton = &Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(pal8; 0), chromaton!(pal8; 1), chromaton!(pal8; 2), None, None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: true,
    };

    pub const RGB565: &Formaton = &Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 5, 11, 0, 2),
                    chromaton!(packrgb; 6,  5, 0, 2),
                    chromaton!(packrgb; 5,  0, 0, 2),
                    None,
                    None],
        elem_size: 2,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const RGB24: &Formaton = &Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 8, 0, 2, 3),
                    chromaton!(packrgb; 8, 0, 1, 3),
                    chromaton!(packrgb; 8, 0, 0, 3),
                    None,
                    None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: false,
    };
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn fmt() {
        println!("formaton yuv- {}", formats::YUV420);
        println!("formaton pal- {}", formats::PAL8);
        println!("formaton rgb565- {}", formats::RGB565);
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering::*;
use std::ops::Not;

#[derive(Debug, Clone)]
pub struct BST<T> {
    nodes: Vec<Option<Node<T>>>,
    root: Option<Ptr>,
    deleted_indices: Vec<Ptr>
}

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

#[derive(Debug, Clone, Copy)]
enum Color {Red, Black}

#[derive(Debug, Clone)]
struct Node<T> {
    elem: T,
    color: Color,
    left: Option<Ptr>,
    right: Option<Ptr>,
}

impl<T: Ord> Node<T> {
    fn new(elem: T, color: Color) -> Self {
        Node { elem: elem, color: color, left: None, right: None }
    }
}

impl Not for Color {
    type Output = Color;
    fn not(self) -> Self {
        match self {
            Color::Red => Color::Black,
            Color::Black => Color::Red,
        }
    }
}

impl<T: Ord> BST<T> {
    fn deref(&self, i: &Ptr) -> &Node<T> {
        self.nodes[i.0].as_ref().expect("deref encounters a reference to a deleted node")
    }

    fn deref_mut(&mut self, i: &Ptr) -> &mut Node<T> {
        self.nodes[i.0].as_mut().expect("deref_mut encounters a reference to a deleted node")
    }

    pub fn new() -> Self {
        BST{ nodes: Vec::new(), root: None, deleted_indices: Vec::new() }
    }

    pub fn singleton(elem: T) -> Self {
        BST{ nodes: vec![Some(Node::new(elem, Color::Black))], root: Some(Ptr(0)), deleted_indices: Vec::new() }
    }

    fn member_impl(&self, ptr: &Option<Ptr>, elem: &T) -> bool {
        match *ptr {
            None => false,
            Some(ref ptr) => {
                let node = self.deref(ptr);
                match node.elem.cmp(elem) {
                    Less => self.member_impl(&node.right, elem),
                    Greater => self.member_impl(&node.left, elem),
                    Equal => true,
                }
            }
        }
    }

    pub fn member(&self, elem: &T) -> bool {
        self.member_impl(&self.root, elem)
    }

    fn is_red(&self, ptr: &Option<Ptr>) -> bool {
        ptr.as_ref().map_or(false, |p| if let Color::Red = self.deref(p).color { true } else { false })
    }

    fn rotate_left(&mut self, h: Ptr) -> Ptr {
        let x : Ptr = self.deref(&h).right.expect("rotate left on node whose left child is nil");
        self.deref_mut(&h).right = self.deref(&x).left;
        self.deref_mut(&x).left = Some(h);
        self.deref_mut(&x).color = self.deref(&h).color;
        self.deref_mut(&h).color = Color::Red;
        x
    }

    fn rotate_right(&mut self, h: Ptr) -> Ptr {
        let x : Ptr = self.deref(&h).left.expect("rotate right on node whose left child is nil");
        self.deref_mut(&h).left = self.deref(&x).right;
        self.deref_mut(&x).right = Some(h);
        self.deref_mut(&x).color = self.deref(&h).color;
        self.deref_mut(&h).color = Color::Red;
        x
    }

    fn move_red_up_or_down(&mut self, h: Ptr) {
        self.deref_mut(&h).color = !self.deref(&h).color;
        let left : Ptr = self.deref(&h).left.expect("move red up/down on node whose left child is nil");
        self.deref_mut(&left).color = !self.deref(&left).color;
        let right: Ptr = self.deref(&h).right.expect("move red up/down on node whose right child is nil");
        self.deref_mut(&right).color = !self.deref(&right).color;
    }

    fn fixup(&mut self, mut node: Ptr) -> Ptr {
        if self.is_red(&self.deref(&node).right) && !self.is_red(&self.deref(&node).left) {
            node = self.rotate_left(node);
        }
        if self.is_red(&self.deref(&node).left) && self.is_red(&self.deref(&self.deref(&node).left.unwrap()).left) {
            node = self.rotate_right(node);
        }
        if self.is_red(&self.deref(&node).left) && self.is_red(&self.deref(&node).right) {
            self.move_red_up_or_down(node);
        }
        node
    }

    fn insert_impl(&mut self, node: Option<Ptr>, elem: T) -> Ptr {
        match node {
            None => {
                let new = Some(Node::new(elem, Color::Red));
                if let Some(index) = self.deleted_indices.pop() {
                    self.nodes[index.0] = new;
                    index
                } else {
                    self.nodes.push(new);
                    Ptr(self.nodes.len() - 1)
                }
            },
            Some(node) => {
                match self.deref(&node).elem.cmp(&elem) {
                    Less => {
                        let right : Option<Ptr> = self.deref(&node).right;
                        let new_right : Ptr = self.insert_impl(right, elem);
                        self.deref_mut(&node).right = Some(new_right);
                    },
                    Greater => {
                        let left : Option<Ptr> = self.deref(&node).left;
                        let new_left : Ptr = self.insert_impl(left, elem);
                        self.deref_mut(&node).left = Some(new_left);
                    },
                    Equal => self.deref_mut(&node).elem = elem,
                }
                self.fixup(node)
            }
        }
    }

    pub fn insert(&mut self, elem: T) {
        let old_root : Option<Ptr> = self.root;
        let new_root : Ptr = self.insert_impl(old_root, elem);
        self.root = Some(new_root);
        self.deref_mut(&new_root).color = Color::Black;
    }

    pub fn clear(&mut self) {
        self.root = None;
        self.nodes.clear();
        self.deleted_indices.clear();
    }

    fn move_red_left(&mut self, mut h: Ptr) -> Ptr {
        self.move_red_up_or_down(h);
        if self.is_red(&self.deref(&self.deref(&h).right.unwrap()).left) {
            self.deref_mut(&h).right = self.deref(&h).right.map(|right| self.rotate_right(right));
            h = self.rotate_left(h);
            self.move_red_up_or_down(h);
        }
        h
    }

    fn take_min_impl(&mut self, mut node: Ptr) -> (T, Option<Ptr>) {
        match self.deref(&node).left {
            None => {
                // The current node is the minimum in the tree.
                self.deleted_indices.push(node);
                (self.nodes[node.0].take().expect("take_min_impl: leftmost node is already deleted").elem, None)
            },
            Some(left) => {
                // We need to make sure the next node is not a 2-node.
                // Making the next node not a 2-node means either it or
                // its left child is red (or both, in the case of a 4-node).
                // This checks if this is violated.
                if !self.is_red(&Some(left)) && !self.is_red(&self.deref(&left).left) {
                    node = self.move_red_left(node);
                }
                let left = self.deref_mut(&node).left.unwrap();
                let (min, new_left) = self.take_min_impl(left);
                self.deref_mut(&node).left = new_left;
                (min, Some(self.fixup(node)))
            }
        }
    }

    pub fn take_min(&mut self) -> Option<T> {
        self.root.and_then(
            |root|
            if self.deref(&root).left.is_none() {
                // The tree has only one element.
                let rv = self.nodes.swap_remove(root.0).map(|node| node.elem);
                self.root = None;
                self.deleted_indices.clear();
                self.nodes.clear();
                rv
            } else {
                // The tree has more than one element.
                let (min, new_root) = self.take_min_impl(root);
                self.root = new_root;
                self.deref_mut(&new_root.unwrap()).color = Color::Black;
                Some(min)
            })
    }

    fn print_structure_inner(&self, node: Option<Ptr>) {
        match node {
            None => print!("[missing]"),
            Some(node_id) => {
                print!("{{ node ");
                let node = self.deref(&node_id);
                if let Color::Red = node.color {
                    print!("[draw=red]");
                }
                print!("{{{:?}}} ", node_id.0); // Prints order of insertion
                if let Color::Red = node.color {
                    print!("edge from parent[red]");
                }
                print!(" child ");
                self.print_structure_inner(node.left);
                print!(" child ");
                self.print_structure_inner(node.right);
                print!(" }}");
            }
        }
    }

    pub fn print_structure(&self) {
        match self.root {
            None => (),
            Some(ref node_id) => {
                println!("%% Put these in your preamble\n\
                          \\usepackage{{tikz}}\n\
                          \\usetikzlibrary{{graphdrawing}}\n\
                          \\usegdlibrary{{trees}}\n\
                          \\definecolor{{red}}{{RGB}}{{171,50,37}}\n\n\
                          %% Put these in the document body\n\
                          \\tikz [binary tree layout, nodes={{draw,circle}}, font=\\sffamily, semithick] \
                          \\node");
                let node = self.deref(&node_id);
                print!("{{{:?}}} child ", node_id.0); // Prints order of insertion
                self.print_structure_inner(node.left);
                print!(" child ");
                self.print_structure_inner(node.right);
                println!(";");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::BST;

    #[test]
    fn basics() {
        let e: BST<i32> = BST::new();
        let s = BST::singleton(2);

        assert_eq!(false, e.member(&1));
        assert_eq!(false, e.member(&2));
        assert_eq!(false, e.member(&3));

        assert_eq!(false, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));
    }

    #[test]
    fn insertion() {
        let mut s = BST::singleton(2);
        s.insert(1);
        assert_eq!(true, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));

        s.insert(4);
        assert_eq!(true, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));
        assert_eq!(true, s.member(&4));

        {
            let mut thousand : BST<i32> = BST::new();
            for i in 0..1000 {
                thousand.insert(i);
            }
            for i in 0..1000 {
                assert_eq!(true, thousand.member(&i));
            }
            assert_eq!(false, thousand.member(&1000));
        }

        {
            let mut thousand : BST<i32> = BST::new();
            for i in 0..1000 {
                if i % 2 == 0 {
                    thousand.insert(i);
                }
            }
            for i in 0..1000 {
                assert_eq!(i % 2 == 0, thousand.member(&i));
            }
            for i in 0..1000 {
                if i % 2 == 1 {
                    thousand.insert(i);
                }
            }
            for i in 0..1000 {
                assert_eq!(true, thousand.member(&i));
            }
        }

        {
            let mut ex : BST<i32> = BST::new();
            for c in 0..64 {
                ex.insert(c);
                ex.print_structure();
                println!("");
            }
        }

        {
            let mut ex : BST<i32> = BST::new();
            let v: [i32; 20] = [14, 9, 12, 6, 2, 10, 1, 18, 16, 5, 8, 17, 13, 3, 11, 15, 7, 19, 4, 20];
            for c in &v {
                ex.insert(*c);
                ex.print_structure();
                println!("");
            }
        }
    }

    #[test]
    fn taking_minimum() {
        {
            let mut tree: BST<i32> = BST::new();
            tree.insert(2);
            tree.insert(3);
            tree.insert(5);
            assert_eq!(tree.take_min(), Some(2)); // [3,5]
            tree.insert(1);                       // [1,3,5]
            assert_eq!(tree.take_min(), Some(1)); // [3,5]
            assert_eq!(tree.take_min(), Some(3)); // [5]
            tree.insert(2);                       // [2,5]
            tree.insert(1);                       // [1,2,5]
            assert_eq!(tree.take_min(), Some(1));
            assert_eq!(tree.take_min(), Some(2));
            assert_eq!(tree.take_min(), Some(5));
            assert_eq!(tree.take_min(), None);
            println!("Passed: take_min for a manual test case with interleaved insertions and deletions")
        }

        {
            let mut tree : BST<i32> = BST::new();
            for size in 0..10 {
                for i in 0..size {
                    tree.insert(i);
                }
                for i in 0..size {
                    assert_eq!(tree.take_min(), Some(i));
                }
                assert_eq!(tree.take_min(), None);
                println!("Passed: take_min for tree with {} element(s)", size);
            }
        }
    }
}
<commit_msg>Add .len() and .is_empty()<commit_after>use std::cmp::Ordering::*;
use std::ops::Not;

#[derive(Debug, Clone)]
pub struct BST<T> {
    nodes: Vec<Option<Node<T>>>,
    root: Option<Ptr>,
    deleted_indices: Vec<Ptr>
}

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

#[derive(Debug, Clone, Copy)]
enum Color {Red, Black}

#[derive(Debug, Clone)]
struct Node<T> {
    elem: T,
    color: Color,
    left: Option<Ptr>,
    right: Option<Ptr>,
}

impl<T: Ord> Node<T> {
    fn new(elem: T, color: Color) -> Self {
        Node { elem: elem, color: color, left: None, right: None }
    }
}

impl Not for Color {
    type Output = Color;
    fn not(self) -> Self {
        match self {
            Color::Red => Color::Black,
            Color::Black => Color::Red,
        }
    }
}

impl<T: Ord> BST<T> {
    fn deref(&self, i: &Ptr) -> &Node<T> {
        self.nodes[i.0].as_ref().expect("deref encounters a reference to a deleted node")
    }

    fn deref_mut(&mut self, i: &Ptr) -> &mut Node<T> {
        self.nodes[i.0].as_mut().expect("deref_mut encounters a reference to a deleted node")
    }

    pub fn new() -> Self {
        BST{ nodes: Vec::new(), root: None, deleted_indices: Vec::new() }
    }

    pub fn singleton(elem: T) -> Self {
        BST{ nodes: vec![Some(Node::new(elem, Color::Black))], root: Some(Ptr(0)), deleted_indices: Vec::new() }
    }

    pub fn len(&self) -> usize {
        self.nodes.len() - self.deleted_indices.len()
    }

    pub fn is_empty(&self) -> bool {
        self.root.is_none()
    }

    fn member_impl(&self, ptr: &Option<Ptr>, elem: &T) -> bool {
        match *ptr {
            None => false,
            Some(ref ptr) => {
                let node = self.deref(ptr);
                match node.elem.cmp(elem) {
                    Less => self.member_impl(&node.right, elem),
                    Greater => self.member_impl(&node.left, elem),
                    Equal => true,
                }
            }
        }
    }

    pub fn member(&self, elem: &T) -> bool {
        self.member_impl(&self.root, elem)
    }

    fn is_red(&self, ptr: &Option<Ptr>) -> bool {
        ptr.as_ref().map_or(false, |p| if let Color::Red = self.deref(p).color { true } else { false })
    }

    fn rotate_left(&mut self, h: Ptr) -> Ptr {
        let x : Ptr = self.deref(&h).right.expect("rotate left on node whose left child is nil");
        self.deref_mut(&h).right = self.deref(&x).left;
        self.deref_mut(&x).left = Some(h);
        self.deref_mut(&x).color = self.deref(&h).color;
        self.deref_mut(&h).color = Color::Red;
        x
    }

    fn rotate_right(&mut self, h: Ptr) -> Ptr {
        let x : Ptr = self.deref(&h).left.expect("rotate right on node whose left child is nil");
        self.deref_mut(&h).left = self.deref(&x).right;
        self.deref_mut(&x).right = Some(h);
        self.deref_mut(&x).color = self.deref(&h).color;
        self.deref_mut(&h).color = Color::Red;
        x
    }

    fn move_red_up_or_down(&mut self, h: Ptr) {
        self.deref_mut(&h).color = !self.deref(&h).color;
        let left : Ptr = self.deref(&h).left.expect("move red up/down on node whose left child is nil");
        self.deref_mut(&left).color = !self.deref(&left).color;
        let right: Ptr = self.deref(&h).right.expect("move red up/down on node whose right child is nil");
        self.deref_mut(&right).color = !self.deref(&right).color;
    }

    fn fixup(&mut self, mut node: Ptr) -> Ptr {
        if self.is_red(&self.deref(&node).right) && !self.is_red(&self.deref(&node).left) {
            node = self.rotate_left(node);
        }
        if self.is_red(&self.deref(&node).left) && self.is_red(&self.deref(&self.deref(&node).left.unwrap()).left) {
            node = self.rotate_right(node);
        }
        if self.is_red(&self.deref(&node).left) && self.is_red(&self.deref(&node).right) {
            self.move_red_up_or_down(node);
        }
        node
    }

    fn insert_impl(&mut self, node: Option<Ptr>, elem: T) -> Ptr {
        match node {
            None => {
                let new = Some(Node::new(elem, Color::Red));
                if let Some(index) = self.deleted_indices.pop() {
                    self.nodes[index.0] = new;
                    index
                } else {
                    self.nodes.push(new);
                    Ptr(self.nodes.len() - 1)
                }
            },
            Some(node) => {
                match self.deref(&node).elem.cmp(&elem) {
                    Less => {
                        let right : Option<Ptr> = self.deref(&node).right;
                        let new_right : Ptr = self.insert_impl(right, elem);
                        self.deref_mut(&node).right = Some(new_right);
                    },
                    Greater => {
                        let left : Option<Ptr> = self.deref(&node).left;
                        let new_left : Ptr = self.insert_impl(left, elem);
                        self.deref_mut(&node).left = Some(new_left);
                    },
                    Equal => self.deref_mut(&node).elem = elem,
                }
                self.fixup(node)
            }
        }
    }

    pub fn insert(&mut self, elem: T) {
        let old_root : Option<Ptr> = self.root;
        let new_root : Ptr = self.insert_impl(old_root, elem);
        self.root = Some(new_root);
        self.deref_mut(&new_root).color = Color::Black;
    }

    pub fn clear(&mut self) {
        self.root = None;
        self.nodes.clear();
        self.deleted_indices.clear();
    }

    fn move_red_left(&mut self, mut h: Ptr) -> Ptr {
        self.move_red_up_or_down(h);
        if self.is_red(&self.deref(&self.deref(&h).right.unwrap()).left) {
            self.deref_mut(&h).right = self.deref(&h).right.map(|right| self.rotate_right(right));
            h = self.rotate_left(h);
            self.move_red_up_or_down(h);
        }
        h
    }

    fn take_min_impl(&mut self, mut node: Ptr) -> (T, Option<Ptr>) {
        match self.deref(&node).left {
            None => {
                // The current node is the minimum in the tree.
                self.deleted_indices.push(node);
                (self.nodes[node.0].take().expect("take_min_impl: leftmost node is already deleted").elem, None)
            },
            Some(left) => {
                // We need to make sure the next node is not a 2-node.
                // Making the next node not a 2-node means either it or
                // its left child is red (or both, in the case of a 4-node).
                // This checks if this is violated.
                if !self.is_red(&Some(left)) && !self.is_red(&self.deref(&left).left) {
                    node = self.move_red_left(node);
                }
                let left = self.deref_mut(&node).left.unwrap();
                let (min, new_left) = self.take_min_impl(left);
                self.deref_mut(&node).left = new_left;
                (min, Some(self.fixup(node)))
            }
        }
    }

    pub fn take_min(&mut self) -> Option<T> {
        self.root.and_then(
            |root|
            if self.deref(&root).left.is_none() {
                // The tree has only one element.
                let rv = self.nodes.swap_remove(root.0).map(|node| node.elem);
                self.root = None;
                self.deleted_indices.clear();
                self.nodes.clear();
                rv
            } else {
                // The tree has more than one element.
                let (min, new_root) = self.take_min_impl(root);
                self.root = new_root;
                self.deref_mut(&new_root.unwrap()).color = Color::Black;
                Some(min)
            })
    }

    fn print_structure_inner(&self, node: Option<Ptr>) {
        match node {
            None => print!("[missing]"),
            Some(node_id) => {
                print!("{{ node ");
                let node = self.deref(&node_id);
                if let Color::Red = node.color {
                    print!("[draw=red]");
                }
                print!("{{{:?}}} ", node_id.0); // Prints order of insertion
                if let Color::Red = node.color {
                    print!("edge from parent[red]");
                }
                print!(" child ");
                self.print_structure_inner(node.left);
                print!(" child ");
                self.print_structure_inner(node.right);
                print!(" }}");
            }
        }
    }

    pub fn print_structure(&self) {
        match self.root {
            None => (),
            Some(ref node_id) => {
                println!("%% Put these in your preamble\n\
                          \\usepackage{{tikz}}\n\
                          \\usetikzlibrary{{graphdrawing}}\n\
                          \\usegdlibrary{{trees}}\n\
                          \\definecolor{{red}}{{RGB}}{{171,50,37}}\n\n\
                          %% Put these in the document body\n\
                          \\tikz [binary tree layout, nodes={{draw,circle}}, font=\\sffamily, semithick] \
                          \\node");
                let node = self.deref(&node_id);
                print!("{{{:?}}} child ", node_id.0); // Prints order of insertion
                self.print_structure_inner(node.left);
                print!(" child ");
                self.print_structure_inner(node.right);
                println!(";");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::BST;

    #[test]
    fn basics() {
        let e: BST<i32> = BST::new();
        let s = BST::singleton(2);

        assert_eq!(false, e.member(&1));
        assert_eq!(false, e.member(&2));
        assert_eq!(false, e.member(&3));

        assert_eq!(false, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));
    }

    #[test]
    fn insertion() {
        let mut s = BST::singleton(2);
        s.insert(1);
        assert_eq!(true, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));

        s.insert(4);
        assert_eq!(true, s.member(&1));
        assert_eq!(true, s.member(&2));
        assert_eq!(false, s.member(&3));
        assert_eq!(true, s.member(&4));

        {
            let mut thousand : BST<i32> = BST::new();
            for i in 0..1000 {
                thousand.insert(i);
            }
            for i in 0..1000 {
                assert_eq!(true, thousand.member(&i));
            }
            assert_eq!(false, thousand.member(&1000));
        }

        {
            let mut thousand : BST<i32> = BST::new();
            for i in 0..1000 {
                if i % 2 == 0 {
                    thousand.insert(i);
                }
            }
            for i in 0..1000 {
                assert_eq!(i % 2 == 0, thousand.member(&i));
            }
            for i in 0..1000 {
                if i % 2 == 1 {
                    thousand.insert(i);
                }
            }
            for i in 0..1000 {
                assert_eq!(true, thousand.member(&i));
            }
        }

        {
            let mut ex : BST<i32> = BST::new();
            for c in 0..64 {
                ex.insert(c);
                ex.print_structure();
                println!("");
            }
        }

        {
            let mut ex : BST<i32> = BST::new();
            let v: [i32; 20] = [14, 9, 12, 6, 2, 10, 1, 18, 16, 5, 8, 17, 13, 3, 11, 15, 7, 19, 4, 20];
            for c in &v {
                ex.insert(*c);
                ex.print_structure();
                println!("");
            }
        }
    }

    #[test]
    fn taking_minimum() {
        {
            let mut tree: BST<i32> = BST::new();
            tree.insert(2);
            tree.insert(3);
            tree.insert(5);
            assert_eq!(tree.take_min(), Some(2)); // [3,5]
            tree.insert(1);                       // [1,3,5]
            assert_eq!(tree.take_min(), Some(1)); // [3,5]
            assert_eq!(tree.take_min(), Some(3)); // [5]
            tree.insert(2);                       // [2,5]
            tree.insert(1);                       // [1,2,5]
            assert_eq!(tree.take_min(), Some(1));
            assert_eq!(tree.take_min(), Some(2));
            assert_eq!(tree.take_min(), Some(5));
            assert_eq!(tree.take_min(), None);
            println!("Passed: take_min for a manual test case with interleaved insertions and deletions")
        }

        {
            let mut tree : BST<i32> = BST::new();
            for size in 0..10 {
                for i in 0..size {
                    tree.insert(i);
                }
                for i in 0..size {
                    assert_eq!(tree.take_min(), Some(i));
                }
                assert_eq!(tree.take_min(), None);
                println!("Passed: take_min for tree with {} element(s)", size);
            }
        }
    }
}
<|endoftext|>


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

#![forbid(unsafe_code)]

use crate::{aws::Aws, cluster::Cluster};
use anyhow::{bail, format_err, Result};
use rusoto_core::RusotoError;
use rusoto_ecr::{
    BatchGetImageRequest, DescribeImagesRequest, DescribeImagesResponse, Ecr, Image,
    ImageIdentifier, PutImageError, PutImageRequest,
};
use rusoto_ecs::{Ecs, UpdateServiceRequest};
use slog_scope::{info, warn};
use std::{env, thread, time::Duration};
use util::retry;

#[derive(Clone)]
pub struct DeploymentManager {
    aws: Aws,
    cluster: Cluster,
    running_tag: String,
}

const VALIDATOR_IMAGE_REPO: &str = "libra_validator";
const NIGHTLY_PREFIX: &str = "nightly";
const UPSTREAM_PREFIX: &str = "upstream_";
const MASTER_PREFIX: &str = "master_";

impl DeploymentManager {
    pub fn new(aws: Aws, cluster: Cluster) -> Self {
        let running_tag = env::var("TAG").unwrap_or_else(|_| aws.workspace().to_string());
        if running_tag.starts_with(NIGHTLY_PREFIX)
            || running_tag.starts_with(UPSTREAM_PREFIX)
            || running_tag.starts_with(MASTER_PREFIX)
        {
            panic!(
                "Cluster test can not deploy if workspace is configured to use {} tag.\
                 Use custom tag for your workspace to use --deploy",
                running_tag
            );
        }
        info!("Will use {} tag for deployment", running_tag);
        Self {
            aws,
            cluster,
            running_tag,
        }
    }

    pub fn redeploy(&mut self, hash: String) -> Result<()> {
        info!("Will deploy with digest {}", hash);
        self.tag_image(
            VALIDATOR_IMAGE_REPO,
            &ImageIdentifier {
                image_digest: Some(hash),
                image_tag: None,
            },
            &self.running_tag,
        )?;
        self.update_all_services()?;
        Ok(())
    }

    pub fn update_all_services(&self) -> Result<()> {
        for instance in self.cluster.all_instances() {
            let mut request = UpdateServiceRequest::default();
            request.cluster = Some(self.aws.workspace().clone());
            request.force_new_deployment = Some(true);
            request.service = instance.peer_name().to_string();

            self.aws
                .ecs()
                .update_service(request)
                .sync()
                .map_err(|e| format_err!("Failed to update {}: {:?}", instance, e))?;
            thread::sleep(Duration::from_millis(200));
        }
        Ok(())
    }

    fn image_digest_by_tag(&self, tag: &str) -> Result<String> {
        let result = self.describe_images(tag);
        let images = result
            .image_details
            .ok_or_else(|| format_err!("No image_details in ECR response"))?;
        if images.len() != 1 {
            bail!("Ecr returned {} images for libra_e2e:nightly", images.len());
        }
        let image = images.into_iter().next().unwrap();
        image
            .image_digest
            .ok_or_else(|| format_err!("No image_digest"))
    }

    fn describe_images(&self, tag: &str) -> DescribeImagesResponse {
        let mut retry = 0usize;
        loop {
            let mut request = DescribeImagesRequest::default();
            request.repository_name = VALIDATOR_IMAGE_REPO.into();
            let image_id = ImageIdentifier {
                image_digest: None,
                image_tag: Some(tag.to_string()),
            };
            request.image_ids = Some(vec![image_id]);
            match self.aws.ecr().describe_images(request).sync() {
                Ok(r) => return r,
                Err(e) => {
                    if retry > 10 {
                        panic!("Failed describe_images after 10 attempts");
                    } else {
                        warn!("Transient failure in describe_images: {}", e);
                        thread::sleep(Duration::from_secs(10));
                        retry += 1;
                    }
                }
            }
        }
    }

    fn get_images(&self, repository: &str, image_id: &ImageIdentifier) -> Result<Vec<Image>> {
        let mut get_request = BatchGetImageRequest::default();
        get_request.repository_name = repository.to_string();
        get_request.image_ids = vec![image_id.clone()];
        // Retry upto 10 times, waiting 10 sec between retries
        let response = retry::retry(retry::fixed_retry_strategy(10_000, 10), || {
            self.aws
                .ecr()
                .batch_get_image(get_request.clone())
                .sync()
                .map_err(|e| {
                    warn!(
                        "Failed to get image from repository: {:?}. Retrying...",
                        get_request.repository_name
                    );
                    e
                })
        })
        .map_err(|e| {
            format_err!(
                "Failed to get image from repository: {:?} after 10 tries: {:?}",
                get_request.repository_name,
                e
            )
        })?;
        response
            .images
            .ok_or_else(|| format_err!("No images in batch_get_image response"))
    }

    fn tag_image(&self, repository: &str, image_id: &ImageIdentifier, new_tag: &str) -> Result<()> {
        let images = self.get_images(repository, &image_id)?;
        let image = images
            .into_iter()
            .next()
            .ok_or_else(|| format_err!("get_images returned 0 images"))?;
        let manifest = image
            .image_manifest
            .ok_or_else(|| format_err!("no manifest in batch_get_image response"))?;
        let mut put_request = PutImageRequest::default();
        put_request.image_manifest = manifest;
        put_request.repository_name = repository.to_string();
        put_request.image_tag = Some(new_tag.to_string());
        let result = self.aws.ecr().put_image(put_request).sync();
        if let Err(e) = result {
            if let RusotoError::Service(PutImageError::ImageAlreadyExists(_)) = e {
                info!("Tag for Image already exists {}:{}", repository, new_tag);
                Ok(())
            } else {
                Err(format_err!("Failed to tag image: {:?}", e))
            }
        } else {
            Ok(())
        }
    }

    pub fn resolve(&self, hash_or_tag: &str) -> Result<String> {
        if hash_or_tag.starts_with("sha256:") {
            Ok(hash_or_tag.to_string())
        } else {
            self.image_digest_by_tag(hash_or_tag)
        }
    }
}
<commit_msg>[cluster-test] Retry flaky update_service AWS call<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::instance::Instance;
use crate::{aws::Aws, cluster::Cluster};
use anyhow::{bail, format_err, Result};
use rusoto_core::RusotoError;
use rusoto_ecr::{
    BatchGetImageRequest, DescribeImagesRequest, DescribeImagesResponse, Ecr, Image,
    ImageIdentifier, PutImageError, PutImageRequest,
};
use rusoto_ecs::{Ecs, UpdateServiceRequest};
use slog_scope::{info, warn};
use std::{env, thread, time::Duration};
use util::retry;

#[derive(Clone)]
pub struct DeploymentManager {
    aws: Aws,
    cluster: Cluster,
    running_tag: String,
}

const VALIDATOR_IMAGE_REPO: &str = "libra_validator";
const NIGHTLY_PREFIX: &str = "nightly";
const UPSTREAM_PREFIX: &str = "upstream_";
const MASTER_PREFIX: &str = "master_";

impl DeploymentManager {
    pub fn new(aws: Aws, cluster: Cluster) -> Self {
        let running_tag = env::var("TAG").unwrap_or_else(|_| aws.workspace().to_string());
        if running_tag.starts_with(NIGHTLY_PREFIX)
            || running_tag.starts_with(UPSTREAM_PREFIX)
            || running_tag.starts_with(MASTER_PREFIX)
        {
            panic!(
                "Cluster test can not deploy if workspace is configured to use {} tag.\
                 Use custom tag for your workspace to use --deploy",
                running_tag
            );
        }
        info!("Will use {} tag for deployment", running_tag);
        Self {
            aws,
            cluster,
            running_tag,
        }
    }

    pub fn redeploy(&mut self, hash: String) -> Result<()> {
        info!("Will deploy with digest {}", hash);
        self.tag_image(
            VALIDATOR_IMAGE_REPO,
            &ImageIdentifier {
                image_digest: Some(hash),
                image_tag: None,
            },
            &self.running_tag,
        )?;
        self.update_all_services()?;
        Ok(())
    }

    pub fn update_all_services(&self) -> Result<()> {
        for instance in self.cluster.all_instances() {
            self.update_service(instance)?;
            thread::sleep(Duration::from_millis(200));
        }
        Ok(())
    }

    fn update_service(&self, instance: &Instance) -> Result<()> {
        for _ in 0..10 {
            let mut request = UpdateServiceRequest::default();
            request.cluster = Some(self.aws.workspace().clone());
            request.force_new_deployment = Some(true);
            request.service = instance.peer_name().to_string();

            let res = self.aws.ecs().update_service(request).sync();
            match res {
                Ok(..) => return Ok(()),
                Err(e) => println!("Failed to update {}: {}, retrying", instance, e),
            }
            thread::sleep(Duration::from_millis(1000));
        }
        Err(format_err!(
            "Failed to update {}: no more retries",
            instance
        ))
    }

    fn image_digest_by_tag(&self, tag: &str) -> Result<String> {
        let result = self.describe_images(tag);
        let images = result
            .image_details
            .ok_or_else(|| format_err!("No image_details in ECR response"))?;
        if images.len() != 1 {
            bail!("Ecr returned {} images for libra_e2e:nightly", images.len());
        }
        let image = images.into_iter().next().unwrap();
        image
            .image_digest
            .ok_or_else(|| format_err!("No image_digest"))
    }

    fn describe_images(&self, tag: &str) -> DescribeImagesResponse {
        let mut retry = 0usize;
        loop {
            let mut request = DescribeImagesRequest::default();
            request.repository_name = VALIDATOR_IMAGE_REPO.into();
            let image_id = ImageIdentifier {
                image_digest: None,
                image_tag: Some(tag.to_string()),
            };
            request.image_ids = Some(vec![image_id]);
            match self.aws.ecr().describe_images(request).sync() {
                Ok(r) => return r,
                Err(e) => {
                    if retry > 10 {
                        panic!("Failed describe_images after 10 attempts");
                    } else {
                        warn!("Transient failure in describe_images: {}", e);
                        thread::sleep(Duration::from_secs(10));
                        retry += 1;
                    }
                }
            }
        }
    }

    fn get_images(&self, repository: &str, image_id: &ImageIdentifier) -> Result<Vec<Image>> {
        let mut get_request = BatchGetImageRequest::default();
        get_request.repository_name = repository.to_string();
        get_request.image_ids = vec![image_id.clone()];
        // Retry upto 10 times, waiting 10 sec between retries
        let response = retry::retry(retry::fixed_retry_strategy(10_000, 10), || {
            self.aws
                .ecr()
                .batch_get_image(get_request.clone())
                .sync()
                .map_err(|e| {
                    warn!(
                        "Failed to get image from repository: {:?}. Retrying...",
                        get_request.repository_name
                    );
                    e
                })
        })
        .map_err(|e| {
            format_err!(
                "Failed to get image from repository: {:?} after 10 tries: {:?}",
                get_request.repository_name,
                e
            )
        })?;
        response
            .images
            .ok_or_else(|| format_err!("No images in batch_get_image response"))
    }

    fn tag_image(&self, repository: &str, image_id: &ImageIdentifier, new_tag: &str) -> Result<()> {
        let images = self.get_images(repository, &image_id)?;
        let image = images
            .into_iter()
            .next()
            .ok_or_else(|| format_err!("get_images returned 0 images"))?;
        let manifest = image
            .image_manifest
            .ok_or_else(|| format_err!("no manifest in batch_get_image response"))?;
        let mut put_request = PutImageRequest::default();
        put_request.image_manifest = manifest;
        put_request.repository_name = repository.to_string();
        put_request.image_tag = Some(new_tag.to_string());
        let result = self.aws.ecr().put_image(put_request).sync();
        if let Err(e) = result {
            if let RusotoError::Service(PutImageError::ImageAlreadyExists(_)) = e {
                info!("Tag for Image already exists {}:{}", repository, new_tag);
                Ok(())
            } else {
                Err(format_err!("Failed to tag image: {:?}", e))
            }
        } else {
            Ok(())
        }
    }

    pub fn resolve(&self, hash_or_tag: &str) -> Result<String> {
        if hash_or_tag.starts_with("sha256:") {
            Ok(hash_or_tag.to_string())
        } else {
            self.image_digest_by_tag(hash_or_tag)
        }
    }
}
<|endoftext|>


<commit_before>use tag::TagType;
use core::{mem, str, slice};

#[repr(packed)]
pub struct BootCommandLineTag {
	tag_type: u32,
	size: u32,
	string: u8, // the first byte of the string
}

impl BootCommandLineTag {
	/// Get the boot loader name.
	fn string(&self) -> &str {
		let length = self.size as usize - mem::size_of::<BootCommandLineTag>();
		unsafe {
			let byte_slice = slice::from_raw_parts((&self.string) as *const u8, length);

			str::from_utf8_unchecked(byte_slice)
		}

	}

	/// Validate the input `BootLoaderNameTag`.
	fn is_valid(&self) -> bool {
		self.tag_type == TagType::BootLoaderName as u32
	}
}
<commit_msg>Fix validation code for BootCommandLine.<commit_after>use tag::TagType;
use core::{mem, str, slice};

#[repr(packed)]
pub struct BootCommandLineTag {
	tag_type: u32,
	size: u32,
	string: u8, // the first byte of the string
}

impl BootCommandLineTag {
	/// Get the boot loader name.
	fn string(&self) -> &str {
		let length = self.size as usize - mem::size_of::<BootCommandLineTag>();
		unsafe {
			let byte_slice = slice::from_raw_parts((&self.string) as *const u8, length);

			str::from_utf8_unchecked(byte_slice)
		}

	}

	/// Validate the input `BootLoaderNameTag`.
	fn is_valid(&self) -> bool {
		self.tag_type == TagType::BootCommandLine as u32
	}
}
<|endoftext|>