<commit_before>/* diosix hypervisor main entry code
 *
 * (c) Chris Williams, 2019-2021.
 *
 * See LICENSE for usage and copying.
 */

/* let the compiler know we're on our own here in bare-metal world */
#![no_std]
#![no_main]
#![feature(asm)]

#![allow(dead_code)]
#![allow(unused_unsafe)]
#![allow(improper_ctypes)]
#![feature(type_ascription)]

/* provide a framework for unit testing */
#![feature(custom_test_frameworks)]
#![test_runner(crate::run_tests)]
#![reexport_test_harness_main = "hvtests"] /* entry point for tests */

/* plug our custom heap allocator into the Rust language: Box, etc */
#![feature(alloc_error_handler)]
#![feature(box_syntax)]
#[allow(unused_imports)]
#[macro_use]
extern crate alloc;

/* needed to convert raw dtb pointer into a slice */
use core::slice;

/* needed for fast lookup tables of stuff */
extern crate hashbrown;

/* needed for elf parsing */
extern crate xmas_elf;

/* needed for device tree parsing and manipulation */
extern crate devicetree;

/* needed for parsing diosix manifest file-system (DMFS) images bundled with the hypervisor */
extern crate dmfs;

/* needed for lazyily-allocated static variables */
#[macro_use]
extern crate lazy_static;

/* this will bring in all the hardware-specific code */
extern crate platform;

/* and now for all our non-hw specific code */
#[macro_use]
mod debug;      /* get us some kind of debug output, typically to a serial port */
#[macro_use]
mod capsule;    /* manage capsules */
#[macro_use]
mod heap;       /* per-CPU private heap management */
#[macro_use]
mod physmem;    /* manage host physical memory */
mod hardware;   /* parse device trees into hardware objects */
mod panic;      /* implement panic() handlers */
mod irq;        /* handle hw interrupts and sw exceptions, collectively known as IRQs */
mod virtmem;    /* manage capsule virtual memory */
mod pcore;      /* manage CPU cores */
mod vcore;      /* virtual CPU core management... */
mod scheduler;  /* ...and scheduling */
mod loader;     /* parse and load supervisor binaries */
mod message;    /* send messages between physical cores */
mod service;    /* allow capsules to register services */
mod manifest;   /* manage capsules loaded with the hypervisor */

/* needed for exclusive locks */
mod lock;
use lock::Mutex;

/* list of error codes */
mod error;
use error::Cause;

use pcore::{PhysicalCoreID, BOOT_PCORE_ID};

/* tell Rust to use our HVallocator to allocate and free heap memory.
although we'll keep track of physical memory, we'll let Rust perform essential
tasks, such as dropping objects when it's no longer needed, borrow checking, etc */
#[global_allocator]
static HV_HEAP: heap::HVallocator = heap::HVallocator;

lazy_static!
{
    /* set to true to allow physical CPU cores to start running supervisor code */
    static ref INIT_DONE: Mutex<bool> = Mutex::new("system bring-up", false);

    /* a physical CPU core obtaining this lock when it is false must walk the DMFS, create
    capsules required to run at boot time, and set the flag to true. any other core
    obtaining it as true must release the lock and move on */
    static ref MANIFEST_UNPACKED: Mutex<bool> = Mutex::new("dmfs unpacked", false);

    /* set to true if individual cores can sound off their presence and capabilities */
    static ref ROLL_CALL: Mutex<bool> = Mutex::new("CPU roll call", false);
}

/* pointer sizes: stick to usize as much as possible: don't always assume it's a 64-bit machine */

/* hventry
   This is the official entry point of the Rust-level hypervisor.
   Call hvmain, which is where all the real work happens, and catch any errors.
   => cpu_nr = this boot-assigned CPU ID number
      dtb_ptr = pointer to start of device tree blob structure
      dtb_len = 32-bit big-endian length of the device tree blob
   <= return to infinite loop, awaiting interrupts */
#[no_mangle]
pub extern "C" fn hventry(cpu_nr: PhysicalCoreID, dtb_ptr: *const u8, dtb_len: u32)
{
    /* carry out tests if that's what we're here for */
    #[cfg(test)]
    hvtests();

    /* if not performing tests, start the system as normal */
    match hvmain(cpu_nr, dtb_ptr, dtb_len)
    {
        Err(e) =>
        {
            hvalert!("Hypervisor failed to start. Reason: {:?}", e);
            debughousekeeper!(); /* attempt to flush queued debug to output */
        },
        _ => () /* continue waiting for an IRQ to come in */
    }
}

/* hvmain
   This code runs at the hypervisor level, with full physical memory access.
   Its job is to initialize physical CPU cores and other resources so that capsules can be
   created in which supervisors run that manage their own user spaces, in which
   applications run. The hypervisor ensures capsules are kept apart using
   hardware protections.

   Assumes all physical CPU cores enter this function during startup.
   The boot CPU is chosen to initialize the system in pre-SMP mode.
   If we're on a single CPU core then everything should still run OK.
   Assumes hardware and exception interrupts are enabled and handlers
   installed.

   Also assumes all CPU cores are compatible ISA-wise. There is provision
   for marking some cores as more powerful than others for systems with
   a mix of performance and efficiency CPU cores.

   => cpu_nr = arbitrary CPU core ID number assigned by boot code,
               separate from hardware ID number.
               BOOT_PCORE_ID = boot CPU core.
      dtb_ptr = pointer to device tree in memory from bootlaoder
      dtb_len = 32-bit big endian size of the device tree
   <= return to infinite loop, waiting for interrupts
*/
fn hvmain(cpu_nr: PhysicalCoreID, dtb_ptr: *const u8, dtb_len: u32) -> Result<(), Cause>
{
    /* set up each physical processor core with its own private heap pool and any other resources.
    each private pool uses physical memory assigned by the pre-hvmain boot code. init() should be called
    first thing to set up each processor core, including the boot CPU, which then sets up the global
    resources. all non-boot CPUs should wait until global resources are ready. */
    pcore::PhysicalCore::init(cpu_nr);

    /* note that pre-physmem::init(), CPU cores rely on their pre-hventry()-assigned
    heap space. after physmem::init(), CPU cores can extend their heaps using physical memory.
    the hypervisor will become stuck pre-physmem::init() if it goes beyond its assigned heap space. */

    match cpu_nr
    {
        /* delegate to boot CPU the welcome banner and set up global resources.
        note: the platform code should ensure whichever CPU core is assigned
        BOOT_PCORE_ID as its cpu_nr can initialize the hypervisor */
        BOOT_PCORE_ID =>
        {
            /* convert the dtb pointer into a rust byte slice. assumes dtb_len is valid */
            let dtb = unsafe { slice::from_raw_parts(dtb_ptr, u32::from_be(dtb_len) as usize) };

            /* process device tree to create data structures representing system hardware,
            allowing these peripherals to be accessed by subsequent routines. this should
            also initialize any found hardware */
            hardware::parse_and_init(dtb)?;

            /* register all the available physical RAM */
            physmem::init()?;
            describe_system();

            /* allow other cores to continue */
            *(INIT_DONE.lock()) = true;
        },

        /* non-boot cores must wait here for early initialization to complete */
        _ => while *(INIT_DONE.lock()) != true {}
    }

    /* we're now ready to start creating capsules to run from the bundled DMFS image.
    the hypervisor can't make any assumptions about the underlying hardware.
    the device tree for these early capsules is derived from the host's device tree,
    modified to virtualize the peripherals. the virtual CPU cores that will run the
    capsule are based on the physical CPU core that creates it. this is more
    straightforward than the hypervisor trying to specify a hypothetical CPU core */
    {
        let mut flag = MANIFEST_UNPACKED.lock();

        if *flag == false
        {
            /* process the manifest and mark it as handled */
            manifest::unpack_at_boot()?;
            *flag = true;

            /* allow all working cores to join the roll call */
            *(ROLL_CALL.lock()) = true;
        }
    }

    /* once ROLL_CALL is set to true, acknowledge we're alive and well, and report CPU core features */
    while *(ROLL_CALL.lock()) != true {}
    hvdebug!("Physical CPU core {:?} ready to roll", pcore::PhysicalCore::describe());

    /* enable timer on this physical CPU core to start scheduling and running virtual cores */
    scheduler::start()?;

    /* initialization complete. fall through to infinite loop waiting for a timer interrupt
    to come in. when it does fire, this stack will be flattened, a virtual CPU loaded up to run,
    and this boot thread will disappear. thus, the call to start() should be the last thing
    this boot thread does */
    Ok(())
}

/* dump system information to the user */
fn describe_system()
{
    const KILOBYTE: usize = 1024;
    const MEGABYTE: usize = KILOBYTE * KILOBYTE;
    const GIGABYTE: usize = KILOBYTE * MEGABYTE;

    /* say hello via the debug port with some information */
    hvdebug!("Diosix {} :: Debug enabled. {} and {} RAM found",

        /* build version number */
        env!("CARGO_PKG_VERSION"),

        /* report number of CPU cores found */
        match hardware::get_nr_cpu_cores()
        {
            None | Some(0) => format!("no CPU cores"),
            Some(1) => format!("1 CPU core"),
            Some(c) => format!("{} CPU cores", c)
        },

        /* count up total system RAM using GiB / MiB / KiB */
        match hardware::get_phys_ram_total()
        {
            Some(t) => if t >= GIGABYTE
            {
                format!("{} GiB", t / GIGABYTE)
            }
            else if t >= MEGABYTE
            {
                format!("{} MiB", t / MEGABYTE)
            }
            else
            {
                format!("{} KiB", t / KILOBYTE)
            },

            None => format!("no")
    });
}

/* mandatory error handler for memory allocations */
#[alloc_error_handler]
fn hvalloc_error(attempt: core::alloc::Layout) -> !
{
    let heap = &(*<pcore::PhysicalCore>::this()).heap;
    hvalert!("hvalloc_error: Failed to allocate/free {} bytes. Heap: {:?}", attempt.size(), heap);
    debughousekeeper!();
    loop {} /* it would be nice to be able to not die here :( */
}

/* perform all unit tests required */
#[cfg(test)]
fn run_tests(unit_tests: &[&dyn Fn()])
{
    /* run each test one by one */
    for test in unit_tests
    {
        test();
    }

    /* exit cleanly once tests are complete */
    platform::test::end(Ok(0));
}

#[test_case]
fn test_assertion()
{
    assert_eq!(42, 42);
}<commit_msg>Reinstate support for non-S-mode CPU cores * Better to include checks for when non-S-mode CPU cores are fully supported<commit_after>/* diosix hypervisor main entry code
 *
 * (c) Chris Williams, 2019-2021.
 *
 * See LICENSE for usage and copying.
 */

/* let the compiler know we're on our own here in bare-metal world */
#![no_std]
#![no_main]
#![feature(asm)]

#![allow(dead_code)]
#![allow(unused_unsafe)]
#![allow(improper_ctypes)]
#![feature(type_ascription)]

/* provide a framework for unit testing */
#![feature(custom_test_frameworks)]
#![test_runner(crate::run_tests)]
#![reexport_test_harness_main = "hvtests"] /* entry point for tests */

/* plug our custom heap allocator into the Rust language: Box, etc */
#![feature(alloc_error_handler)]
#![feature(box_syntax)]
#[allow(unused_imports)]
#[macro_use]
extern crate alloc;

/* needed to convert raw dtb pointer into a slice */
use core::slice;

/* needed for fast lookup tables of stuff */
extern crate hashbrown;

/* needed for elf parsing */
extern crate xmas_elf;

/* needed for device tree parsing and manipulation */
extern crate devicetree;

/* needed for parsing diosix manifest file-system (DMFS) images bundled with the hypervisor */
extern crate dmfs;

/* needed for lazyily-allocated static variables */
#[macro_use]
extern crate lazy_static;

/* this will bring in all the hardware-specific code */
extern crate platform;

/* and now for all our non-hw specific code */
#[macro_use]
mod debug;      /* get us some kind of debug output, typically to a serial port */
#[macro_use]
mod capsule;    /* manage capsules */
#[macro_use]
mod heap;       /* per-CPU private heap management */
#[macro_use]
mod physmem;    /* manage host physical memory */
mod hardware;   /* parse device trees into hardware objects */
mod panic;      /* implement panic() handlers */
mod irq;        /* handle hw interrupts and sw exceptions, collectively known as IRQs */
mod virtmem;    /* manage capsule virtual memory */
mod pcore;      /* manage CPU cores */
mod vcore;      /* virtual CPU core management... */
mod scheduler;  /* ...and scheduling */
mod loader;     /* parse and load supervisor binaries */
mod message;    /* send messages between physical cores */
mod service;    /* allow capsules to register services */
mod manifest;   /* manage capsules loaded with the hypervisor */

/* needed for exclusive locks */
mod lock;
use lock::Mutex;

/* list of error codes */
mod error;
use error::Cause;

use pcore::{PhysicalCoreID, BOOT_PCORE_ID};

/* tell Rust to use our HVallocator to allocate and free heap memory.
although we'll keep track of physical memory, we'll let Rust perform essential
tasks, such as dropping objects when it's no longer needed, borrow checking, etc */
#[global_allocator]
static HV_HEAP: heap::HVallocator = heap::HVallocator;

lazy_static!
{
    /* set to true to allow physical CPU cores to start running supervisor code */
    static ref INIT_DONE: Mutex<bool> = Mutex::new("system bring-up", false);

    /* a physical CPU core obtaining this lock when it is false must walk the DMFS, create
    capsules required to run at boot time, and set the flag to true. any other core
    obtaining it as true must release the lock and move on */
    static ref MANIFEST_UNPACKED: Mutex<bool> = Mutex::new("dmfs unpacked", false);

    /* set to true if individual cores can sound off their presence and capabilities */
    static ref ROLL_CALL: Mutex<bool> = Mutex::new("CPU roll call", false);
}

/* pointer sizes: stick to usize as much as possible: don't always assume it's a 64-bit machine */

/* hventry
   This is the official entry point of the Rust-level hypervisor.
   Call hvmain, which is where all the real work happens, and catch any errors.
   => cpu_nr = this boot-assigned CPU ID number
      dtb_ptr = pointer to start of device tree blob structure
      dtb_len = 32-bit big-endian length of the device tree blob
   <= return to infinite loop, awaiting interrupts */
#[no_mangle]
pub extern "C" fn hventry(cpu_nr: PhysicalCoreID, dtb_ptr: *const u8, dtb_len: u32)
{
    /* carry out tests if that's what we're here for */
    #[cfg(test)]
    hvtests();

    /* if not performing tests, start the system as normal */
    match hvmain(cpu_nr, dtb_ptr, dtb_len)
    {
        Err(e) =>
        {
            hvalert!("Hypervisor failed to start. Reason: {:?}", e);
            debughousekeeper!(); /* attempt to flush queued debug to output */
        },
        _ => () /* continue waiting for an IRQ to come in */
    }
}

/* hvmain
   This code runs at the hypervisor level, with full physical memory access.
   Its job is to initialize physical CPU cores and other resources so that capsules can be
   created in which supervisors run that manage their own user spaces, in which
   applications run. The hypervisor ensures capsules are kept apart using
   hardware protections.

   Assumes all physical CPU cores enter this function during startup.
   The boot CPU is chosen to initialize the system in pre-SMP mode.
   If we're on a single CPU core then everything should still run OK.
   Assumes hardware and exception interrupts are enabled and handlers
   installed.

   Also assumes all CPU cores are compatible ISA-wise. There is provision
   for marking some cores as more powerful than others for systems with
   a mix of performance and efficiency CPU cores.

   => cpu_nr = arbitrary CPU core ID number assigned by boot code,
               separate from hardware ID number.
               BOOT_PCORE_ID = boot CPU core.
      dtb_ptr = pointer to device tree in memory from bootlaoder
      dtb_len = 32-bit big endian size of the device tree
   <= return to infinite loop, waiting for interrupts
*/
fn hvmain(cpu_nr: PhysicalCoreID, dtb_ptr: *const u8, dtb_len: u32) -> Result<(), Cause>
{
    /* set up each physical processor core with its own private heap pool and any other resources.
    each private pool uses physical memory assigned by the pre-hvmain boot code. init() should be called
    first thing to set up each processor core, including the boot CPU, which then sets up the global
    resources. all non-boot CPUs should wait until global resources are ready. */
    pcore::PhysicalCore::init(cpu_nr);

    /* note that pre-physmem::init(), CPU cores rely on their pre-hventry()-assigned
    heap space. after physmem::init(), CPU cores can extend their heaps using physical memory.
    the hypervisor will become stuck pre-physmem::init() if it goes beyond its assigned heap space. */

    match cpu_nr
    {
        /* delegate to boot CPU the welcome banner and set up global resources.
        note: the platform code should ensure whichever CPU core is assigned
        BOOT_PCORE_ID as its cpu_nr can initialize the hypervisor */
        BOOT_PCORE_ID =>
        {
            /* convert the dtb pointer into a rust byte slice. assumes dtb_len is valid */
            let dtb = unsafe { slice::from_raw_parts(dtb_ptr, u32::from_be(dtb_len) as usize) };

            /* process device tree to create data structures representing system hardware,
            allowing these peripherals to be accessed by subsequent routines. this should
            also initialize any found hardware */
            hardware::parse_and_init(dtb)?;

            /* register all the available physical RAM */
            physmem::init()?;
            describe_system();

            /* allow other cores to continue */
            *(INIT_DONE.lock()) = true;
        },

        /* non-boot cores must wait here for early initialization to complete */
        _ => while *(INIT_DONE.lock()) != true {}
    }

    /* Create capsules to run from the bundled DMFS image.
    the hypervisor can't make any assumptions about the underlying hardware.
    the device tree for these early capsules is derived from the host's device tree,
    modified to virtualize peripherals. the virtual CPU cores will based on the
    physical CPU core that creates it. this is more straightforward than the hypervisor
    trying to specify a hypothetical CPU core

    as such, only allow supervisor-mode capable CPU cores to build capasules */
    if pcore::PhysicalCore::smode_supported() == true
    {
        let mut flag = MANIFEST_UNPACKED.lock();

        if *flag == false
        {
            /* process the manifest and mark it as handled */
            manifest::unpack_at_boot()?;
            *flag = true;

            /* allow all working cores to join the roll call */
            *(ROLL_CALL.lock()) = true;
        }
    }

    /* once ROLL_CALL is set to true, acknowledge we're alive and well, and report CPU core features */
    while *(ROLL_CALL.lock()) != true {}
    hvdebug!("Physical CPU core {:?} ready to roll", pcore::PhysicalCore::describe());

    /* enable timer on this physical CPU core to start scheduling and running virtual cores */
    scheduler::start()?;

    /* initialization complete. fall through to infinite loop waiting for a timer interrupt
    to come in. when it does fire, this stack will be flattened, a virtual CPU loaded up to run,
    and this boot thread will disappear. thus, the call to start() should be the last thing
    this boot thread does */
    Ok(())
}

/* dump system information to the user */
fn describe_system()
{
    const KILOBYTE: usize = 1024;
    const MEGABYTE: usize = KILOBYTE * KILOBYTE;
    const GIGABYTE: usize = KILOBYTE * MEGABYTE;

    /* say hello via the debug port with some information */
    hvdebug!("Diosix {} :: Debug enabled. {} and {} RAM found",

        /* build version number */
        env!("CARGO_PKG_VERSION"),

        /* report number of CPU cores found */
        match hardware::get_nr_cpu_cores()
        {
            None | Some(0) => format!("no CPU cores"),
            Some(1) => format!("1 CPU core"),
            Some(c) => format!("{} CPU cores", c)
        },

        /* count up total system RAM using GiB / MiB / KiB */
        match hardware::get_phys_ram_total()
        {
            Some(t) => if t >= GIGABYTE
            {
                format!("{} GiB", t / GIGABYTE)
            }
            else if t >= MEGABYTE
            {
                format!("{} MiB", t / MEGABYTE)
            }
            else
            {
                format!("{} KiB", t / KILOBYTE)
            },

            None => format!("no")
    });
}

/* mandatory error handler for memory allocations */
#[alloc_error_handler]
fn hvalloc_error(attempt: core::alloc::Layout) -> !
{
    let heap = &(*<pcore::PhysicalCore>::this()).heap;
    hvalert!("hvalloc_error: Failed to allocate/free {} bytes. Heap: {:?}", attempt.size(), heap);
    debughousekeeper!();
    loop {} /* it would be nice to be able to not die here :( */
}

/* perform all unit tests required */
#[cfg(test)]
fn run_tests(unit_tests: &[&dyn Fn()])
{
    /* run each test one by one */
    for test in unit_tests
    {
        test();
    }

    /* exit cleanly once tests are complete */
    platform::test::end(Ok(0));
}

#[test_case]
fn test_assertion()
{
    assert_eq!(42, 42);
}<|endoftext|>


<commit_before>#[macro_use]
extern crate quote;
extern crate reqwest;
extern crate serde_json;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate lazy_static;
extern crate linked_hash_map;
extern crate inflector;
extern crate regex;

use std::fs::File;
use serde_json::{ Value };
use std::collections::{ HashSet };
use quote::{ Tokens, Ident, ToTokens };
use std::ascii::AsciiExt;
use linked_hash_map::LinkedHashMap;
use inflector::Inflector;
use regex::Regex;

mod util;
mod inference;

use util::*;
use inference::*;

mod errors {
    error_chain! {
        types {
            Error, ErrorKind, ResultExt, Result;
        }

        links {
        }

        foreign_links {
            ReqwestError(::reqwest::Error);
            JsonError(::serde_json::Error);
            IoError(::std::io::Error);
        }

        errors {
            MissingSource {
                description("No source for sample specified")
            }
            ExistingType(t: String) {
                display("No code generated, JSON matches existing type {}", t)
            }
        }
    }
}

pub use errors::*;

const RUST_KEYWORDS_ARR: &'static [&'static str] = &["abstract", "alignof", "as", "become", "box",
    "break", "const", "continue", "crate", "do", "else", "enum", "extern", "false", "final", "fn",
    "for", "if", "impl", "in", "let", "loop", "macro", "match", "mod", "move", "mut", "offsetof",
    "override", "priv", "proc", "pub", "pure", "ref", "return", "Self", "self", "sizeof", "static",
    "struct", "super", "trait", "true", "type", "typeof", "unsafe", "unsized", "use", "virtual",
    "where", "while", "yield"];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> = {
        RUST_KEYWORDS_ARR.iter().cloned().collect()
    };
}

pub enum SampleSource<'a> {
    Url(&'a str),
    File(&'a str),
    Text(&'a str),
}

#[derive(Debug, PartialEq)]
pub enum MissingFields {
    Fail,
    UseDefault,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Visibility {
    Private,
    Pub,
    PubRestricted(String)
}

impl ToTokens for Visibility {
    fn to_tokens(&self, tokens: &mut Tokens) {
        use Visibility::*;
        match *self {
            Private => {},
            Pub => {
                tokens.append("pub");
            }
            PubRestricted(ref path) => {
                tokens.append("pub(");
                tokens.append(path);
                tokens.append(")");
            }
        }
    }
}

pub enum FieldVisibility {
    Inherited,
    Specified(Visibility)
}

pub struct Options {
    pub extern_crate: bool,
    pub runnable: bool,
    pub missing_fields: MissingFields,
    pub deny_unknown_fields: bool,
    pub allow_option_vec: bool,
    pub type_visibility: Visibility,
    pub field_visibility: FieldVisibility,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            extern_crate: false,
            runnable: false,
            missing_fields: MissingFields::Fail,
            deny_unknown_fields: false,
            allow_option_vec: false,
            type_visibility: Visibility::Private,
            field_visibility: FieldVisibility::Inherited,
        }
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    types: Vec<Tokens>,
}

macro_rules! some_if {
    ($cond:expr, $then:expr) => ({
        if $cond {
            Some($then)
        } else {
            None
        }
    })
}

pub fn from_str_with_defaults(name: &str, json: &str) -> Result<Tokens> {
    codegen(name, &SampleSource::Text(json), Options::default())
}

pub fn codegen(name: &str, source: &SampleSource, mut options: Options) -> Result<Tokens> {
    let sample = get_and_parse_sample(source)?;
    let name = handle_pub_in_name(name, &mut options);

    let mut ctxt = Ctxt {
        options: options,
        type_names: HashSet::new(),
        types: Vec::new(),
    };
    let inferred = infer_type_from_value(&sample);
    let type_name = generate_type_from_inferred(&mut ctxt, name, &inferred);

    let example = some_if!(ctxt.options.runnable, {
        ctxt.options.extern_crate = true;
        usage_example(&type_name)
    });

    let crates = some_if!(ctxt.options.extern_crate, quote! {
        #[macro_use]
        extern crate serde_derive;
        extern crate serde_json;
    });

    let defs = ctxt.types.iter().rev();

    Ok(quote! {
        #crates

        #(#defs)*

        #example
    })
}

/// Parse "names" like `pub(crate) Foo` into a name and a visibility option
fn handle_pub_in_name<'a>(name: &'a str, options: &mut Options) -> &'a str {
    lazy_static! {
        static ref PUB_RE: Regex =
            Regex::new(r"(?x)
                pub ( \( (?P<restriction> [^)]+ ) \) )?
                \s+
                (?P<name> .+ )
            ").unwrap();
    }
    match PUB_RE.captures(name) {
        Some(captures) => {
            options.type_visibility = match captures.name("restriction") {
                Some(restriction) => Visibility::PubRestricted(restriction.as_str().to_owned()),
                None => Visibility::Pub,
            };
            captures.name("name").unwrap().as_str()
        }
        None => {
            // If there is no visibility specified here, we want to use whatever is set elsewhere
            name
        }
    }
}

fn usage_example(type_id: &Tokens) -> Tokens {
    let var_id = Ident::from(snake_case(type_id.as_str()));

    quote! {
        fn main() {
            let #var_id: #type_id = Default::default();

            let serialized = serde_json::to_string(&#var_id).unwrap();

            println!("serialized = {}", serialized);

            let deserialized: #type_id = serde_json::from_str(&serialized).unwrap();

            println!("deserialized = {:?}", deserialized);
        }
    }
}

fn generate_type_from_inferred(ctxt: &mut Ctxt, path: &str, inferred: &InferredType) -> Tokens {
    use InferredType::*;
    match *inferred {
        Null | Any => quote! { ::serde_json::Value },
        Bool => quote! { bool },
        StringT => quote! { String },
        Integer => quote! { i64 },
        Floating => quote! { f64 },
        EmptyVec => quote! { Vec<::serde_json::Value> },
        VecT { elem_type: ref e } => {
            let singular = path.to_singular();
            let inner = generate_type_from_inferred(ctxt, &singular, e);
            quote! { Vec<#inner> }
        }
        Struct { fields: ref map } => {
            generate_struct_from_inferred_fields(ctxt, path, map)
        }
        Optional(ref e) => {
            let inner = generate_type_from_inferred(ctxt, path, e);
            match ctxt.options.missing_fields {
                MissingFields::Fail => quote! { Option<#inner> },
                MissingFields::UseDefault => quote! { #inner },
            }
        }
    }
}

fn field_name(name: &str, used_names: &HashSet<String>) -> String {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> String {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

fn type_or_field_name(name: &str,
                      used_names: &HashSet<String>,
                      default_name: &str,
                      case_fn: fn(&str) -> String)
                      -> String {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii() && c.is_numeric() {
            let temp = String::from("n") + name;
            case_fn(&temp)
        } else {
            case_fn(name)
        }
    } else {
        case_fn(name)
    };
    if RUST_KEYWORDS.contains::<str>(&output_name) {
        output_name.push_str("_field");
    }
    if output_name == "" {
        output_name.push_str(default_name);
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn collapse_option_vec<'a>(ctxt: &mut Ctxt,
                           typ: &'a InferredType)
                           -> (Option<Tokens>, &'a InferredType) {
    if !ctxt.options.allow_option_vec && ctxt.options.missing_fields != MissingFields::UseDefault {
        if let InferredType::Optional(ref inner) = *typ {
            match **inner {
                InferredType::EmptyVec | InferredType::VecT { .. } => {
                    return (Some(quote! { #[serde(default)] }), &**inner);
                }
                _ => {}
            }
        }
    }
    (None, typ)
}

fn generate_struct_from_inferred_fields(
        ctxt: &mut Ctxt,
        path: &str,
        map: &LinkedHashMap<String, InferredType>) -> Tokens {
    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    let ident = Ident::from(type_name);
    let visibility = ctxt.options.type_visibility.clone();
    let field_visibility = match ctxt.options.field_visibility {
        FieldVisibility::Inherited => visibility.clone(),
        FieldVisibility::Specified(ref v) => v.clone(),
    };

    let mut field_names = HashSet::new();

    let fields: Vec<Tokens> = map.iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());
            let rename = some_if!(&field_name != name,
                quote! { #[serde(rename = #name)] });
            let field_ident = Ident::from(field_name);
            let (default, collapsed) = collapse_option_vec(ctxt, typ);
            let field_type = generate_type_from_inferred(ctxt, name, collapsed);
            quote! {
                #rename
                #default
                #field_visibility #field_ident: #field_type
            }
        })
        .collect();

    let derives = quote! { #[derive(Default, Debug, Clone, Serialize, Deserialize)] };

    let unknown_fields = some_if!(ctxt.options.deny_unknown_fields,
        quote! { #[serde(deny_unknown_fields)] });

    let use_defaults = some_if!(ctxt.options.missing_fields == MissingFields::UseDefault,
        quote! { #[serde(default)] });

    let code = quote! {
        #derives
        #unknown_fields
        #use_defaults
        #visibility struct #ident {
            #(#fields),*
        }
    };

    ctxt.types.push(code);
    quote! { #ident }
}

pub fn infer_source_type(s: &str) -> SampleSource {
    let s = s.trim();
    if s.starts_with('{') || s.starts_with('[') {
        return SampleSource::Text(s);
    }
    if s.starts_with("http://") || s.starts_with("https://") {
        return SampleSource::Url(s);
    }
    SampleSource::File(s)
}

fn get_and_parse_sample(source: &SampleSource) -> Result<Value> {
    let parse_result = match *source {
        SampleSource::Url(url) => serde_json::de::from_reader(reqwest::get(url)?),
        SampleSource::File(path) => serde_json::de::from_reader(File::open(path)?),
        SampleSource::Text(text) => serde_json::from_str(text),
    };
    Ok(parse_result.chain_err(|| "Unable to parse JSON sample")?)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_handle_pub_in_name() {
        let mut options = Options::default();
        let name = handle_pub_in_name("Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Options::default().type_visibility);
        let name = handle_pub_in_name("pub Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Visibility::Pub);
        let name = handle_pub_in_name("pub(crate) Foo Bar", &mut options);
        assert_eq!(name, "Foo Bar");
        assert_eq!(options.type_visibility, Visibility::PubRestricted("crate".to_string()));
        let name = handle_pub_in_name("pub(some::path) Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Visibility::PubRestricted("some::path".to_string()));
    }
}
<commit_msg>Make derive list configurable<commit_after>#[macro_use]
extern crate quote;
extern crate reqwest;
extern crate serde_json;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate lazy_static;
extern crate linked_hash_map;
extern crate inflector;
extern crate regex;

use std::fs::File;
use serde_json::{ Value };
use std::collections::{ HashSet };
use quote::{ Tokens, Ident, ToTokens };
use std::ascii::AsciiExt;
use linked_hash_map::LinkedHashMap;
use inflector::Inflector;
use regex::Regex;

mod util;
mod inference;

use util::*;
use inference::*;

mod errors {
    error_chain! {
        types {
            Error, ErrorKind, ResultExt, Result;
        }

        links {
        }

        foreign_links {
            ReqwestError(::reqwest::Error);
            JsonError(::serde_json::Error);
            IoError(::std::io::Error);
        }

        errors {
            MissingSource {
                description("No source for sample specified")
            }
            ExistingType(t: String) {
                display("No code generated, JSON matches existing type {}", t)
            }
        }
    }
}

pub use errors::*;

const RUST_KEYWORDS_ARR: &'static [&'static str] = &["abstract", "alignof", "as", "become", "box",
    "break", "const", "continue", "crate", "do", "else", "enum", "extern", "false", "final", "fn",
    "for", "if", "impl", "in", "let", "loop", "macro", "match", "mod", "move", "mut", "offsetof",
    "override", "priv", "proc", "pub", "pure", "ref", "return", "Self", "self", "sizeof", "static",
    "struct", "super", "trait", "true", "type", "typeof", "unsafe", "unsized", "use", "virtual",
    "where", "while", "yield"];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> = {
        RUST_KEYWORDS_ARR.iter().cloned().collect()
    };
}

pub enum SampleSource<'a> {
    Url(&'a str),
    File(&'a str),
    Text(&'a str),
}

#[derive(Debug, PartialEq)]
pub enum MissingFields {
    Fail,
    UseDefault,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Visibility {
    Private,
    Pub,
    PubRestricted(String)
}

impl ToTokens for Visibility {
    fn to_tokens(&self, tokens: &mut Tokens) {
        use Visibility::*;
        match *self {
            Private => {},
            Pub => {
                tokens.append("pub");
            }
            PubRestricted(ref path) => {
                tokens.append("pub(");
                tokens.append(path);
                tokens.append(")");
            }
        }
    }
}

pub enum FieldVisibility {
    Inherited,
    Specified(Visibility)
}

pub struct Options {
    pub extern_crate: bool,
    pub runnable: bool,
    pub missing_fields: MissingFields,
    pub deny_unknown_fields: bool,
    pub allow_option_vec: bool,
    pub type_visibility: Visibility,
    pub field_visibility: FieldVisibility,
    pub derives: String,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            extern_crate: false,
            runnable: false,
            missing_fields: MissingFields::Fail,
            deny_unknown_fields: false,
            allow_option_vec: false,
            type_visibility: Visibility::Private,
            field_visibility: FieldVisibility::Inherited,
            derives: "Default, Debug, Clone, PartialEq, Serialize, Deserialize".into(),
        }
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    types: Vec<Tokens>,
}

macro_rules! some_if {
    ($cond:expr, $then:expr) => ({
        if $cond {
            Some($then)
        } else {
            None
        }
    })
}

pub fn from_str_with_defaults(name: &str, json: &str) -> Result<Tokens> {
    codegen(name, &SampleSource::Text(json), Options::default())
}

pub fn codegen(name: &str, source: &SampleSource, mut options: Options) -> Result<Tokens> {
    let sample = get_and_parse_sample(source)?;
    let name = handle_pub_in_name(name, &mut options);

    let mut ctxt = Ctxt {
        options: options,
        type_names: HashSet::new(),
        types: Vec::new(),
    };
    let inferred = infer_type_from_value(&sample);
    let type_name = generate_type_from_inferred(&mut ctxt, name, &inferred);

    let example = some_if!(ctxt.options.runnable, {
        ctxt.options.extern_crate = true;
        usage_example(&type_name)
    });

    let crates = some_if!(ctxt.options.extern_crate, quote! {
        #[macro_use]
        extern crate serde_derive;
        extern crate serde_json;
    });

    let defs = ctxt.types.iter().rev();

    Ok(quote! {
        #crates

        #(#defs)*

        #example
    })
}

/// Parse "names" like `pub(crate) Foo` into a name and a visibility option
fn handle_pub_in_name<'a>(name: &'a str, options: &mut Options) -> &'a str {
    lazy_static! {
        static ref PUB_RE: Regex =
            Regex::new(r"(?x)
                pub ( \( (?P<restriction> [^)]+ ) \) )?
                \s+
                (?P<name> .+ )
            ").unwrap();
    }
    match PUB_RE.captures(name) {
        Some(captures) => {
            options.type_visibility = match captures.name("restriction") {
                Some(restriction) => Visibility::PubRestricted(restriction.as_str().to_owned()),
                None => Visibility::Pub,
            };
            captures.name("name").unwrap().as_str()
        }
        None => {
            // If there is no visibility specified here, we want to use whatever is set elsewhere
            name
        }
    }
}

fn usage_example(type_id: &Tokens) -> Tokens {
    let var_id = Ident::from(snake_case(type_id.as_str()));

    quote! {
        fn main() {
            let #var_id: #type_id = Default::default();

            let serialized = serde_json::to_string(&#var_id).unwrap();

            println!("serialized = {}", serialized);

            let deserialized: #type_id = serde_json::from_str(&serialized).unwrap();

            println!("deserialized = {:?}", deserialized);
        }
    }
}

fn generate_type_from_inferred(ctxt: &mut Ctxt, path: &str, inferred: &InferredType) -> Tokens {
    use InferredType::*;
    match *inferred {
        Null | Any => quote! { ::serde_json::Value },
        Bool => quote! { bool },
        StringT => quote! { String },
        Integer => quote! { i64 },
        Floating => quote! { f64 },
        EmptyVec => quote! { Vec<::serde_json::Value> },
        VecT { elem_type: ref e } => {
            let singular = path.to_singular();
            let inner = generate_type_from_inferred(ctxt, &singular, e);
            quote! { Vec<#inner> }
        }
        Struct { fields: ref map } => {
            generate_struct_from_inferred_fields(ctxt, path, map)
        }
        Optional(ref e) => {
            let inner = generate_type_from_inferred(ctxt, path, e);
            match ctxt.options.missing_fields {
                MissingFields::Fail => quote! { Option<#inner> },
                MissingFields::UseDefault => quote! { #inner },
            }
        }
    }
}

fn field_name(name: &str, used_names: &HashSet<String>) -> String {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> String {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

fn type_or_field_name(name: &str,
                      used_names: &HashSet<String>,
                      default_name: &str,
                      case_fn: fn(&str) -> String)
                      -> String {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii() && c.is_numeric() {
            let temp = String::from("n") + name;
            case_fn(&temp)
        } else {
            case_fn(name)
        }
    } else {
        case_fn(name)
    };
    if RUST_KEYWORDS.contains::<str>(&output_name) {
        output_name.push_str("_field");
    }
    if output_name == "" {
        output_name.push_str(default_name);
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn collapse_option_vec<'a>(ctxt: &mut Ctxt,
                           typ: &'a InferredType)
                           -> (Option<Tokens>, &'a InferredType) {
    if !ctxt.options.allow_option_vec && ctxt.options.missing_fields != MissingFields::UseDefault {
        if let InferredType::Optional(ref inner) = *typ {
            match **inner {
                InferredType::EmptyVec | InferredType::VecT { .. } => {
                    return (Some(quote! { #[serde(default)] }), &**inner);
                }
                _ => {}
            }
        }
    }
    (None, typ)
}

fn generate_struct_from_inferred_fields(
        ctxt: &mut Ctxt,
        path: &str,
        map: &LinkedHashMap<String, InferredType>) -> Tokens {
    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    let ident = Ident::from(type_name);
    let visibility = ctxt.options.type_visibility.clone();
    let field_visibility = match ctxt.options.field_visibility {
        FieldVisibility::Inherited => visibility.clone(),
        FieldVisibility::Specified(ref v) => v.clone(),
    };

    let mut field_names = HashSet::new();

    let fields: Vec<Tokens> = map.iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());
            let rename = some_if!(&field_name != name,
                quote! { #[serde(rename = #name)] });
            let field_ident = Ident::from(field_name);
            let (default, collapsed) = collapse_option_vec(ctxt, typ);
            let field_type = generate_type_from_inferred(ctxt, name, collapsed);
            quote! {
                #rename
                #default
                #field_visibility #field_ident: #field_type
            }
        })
        .collect();

    let derive_list = Ident::from(&*ctxt.options.derives);

    let unknown_fields = some_if!(ctxt.options.deny_unknown_fields,
        quote! { #[serde(deny_unknown_fields)] });

    let use_defaults = some_if!(ctxt.options.missing_fields == MissingFields::UseDefault,
        quote! { #[serde(default)] });

    let code = quote! {
        #[derive(#derive_list)]
        #unknown_fields
        #use_defaults
        #visibility struct #ident {
            #(#fields),*
        }
    };

    ctxt.types.push(code);
    quote! { #ident }
}

pub fn infer_source_type(s: &str) -> SampleSource {
    let s = s.trim();
    if s.starts_with('{') || s.starts_with('[') {
        return SampleSource::Text(s);
    }
    if s.starts_with("http://") || s.starts_with("https://") {
        return SampleSource::Url(s);
    }
    SampleSource::File(s)
}

fn get_and_parse_sample(source: &SampleSource) -> Result<Value> {
    let parse_result = match *source {
        SampleSource::Url(url) => serde_json::de::from_reader(reqwest::get(url)?),
        SampleSource::File(path) => serde_json::de::from_reader(File::open(path)?),
        SampleSource::Text(text) => serde_json::from_str(text),
    };
    Ok(parse_result.chain_err(|| "Unable to parse JSON sample")?)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_handle_pub_in_name() {
        let mut options = Options::default();
        let name = handle_pub_in_name("Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Options::default().type_visibility);
        let name = handle_pub_in_name("pub Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Visibility::Pub);
        let name = handle_pub_in_name("pub(crate) Foo Bar", &mut options);
        assert_eq!(name, "Foo Bar");
        assert_eq!(options.type_visibility, Visibility::PubRestricted("crate".to_string()));
        let name = handle_pub_in_name("pub(some::path) Foo", &mut options);
        assert_eq!(name, "Foo");
        assert_eq!(options.type_visibility, Visibility::PubRestricted("some::path".to_string()));
    }
}
<|endoftext|>


<commit_before>//! The `select` macro.

use proc_macro_hack::proc_macro_hack;

#[doc(hidden)]
#[macro_export]
macro_rules! document_select_macro {
    ($item:item) => {
        /// Polls multiple futures and streams simultaneously, executing the branch
        /// for the future that finishes first. Futures passed to
        /// `select!` must be `Unpin` and implement `FusedFuture`.
        /// Futures and streams which are not already fused can be fused using the
        /// `.fuse()` method. Note, though, that fusing a future or stream directly
        /// in the call to `select!` will not be enough to prevent it from being
        /// polled after completion if the `select!` call is in a loop, so when
        /// `select!`ing in a loop, users should take care to `fuse()` outside of
        /// the loop.
        ///
        /// `select!` can select over futures with different output types, but each
        /// branch has to have the same return type.
        ///
        /// This macro is only usable inside of async functions, closures, and blocks.
        /// It is also gated behind the `async-await` feature of this library, which is
        /// _not_ activated by default.
        ///
        /// # Examples
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::select;
        /// let mut a = future::ready(4);
        /// let mut b = future::empty::<()>();
        ///
        /// let res = select! {
        ///     a_res = a => a_res + 1,
        ///     _ = b => 0,
        /// };
        /// assert_eq!(res, 5);
        /// # });
        /// ```
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::stream::{self, StreamExt};
        /// use futures::select;
        /// let mut st = stream::iter(vec![2]).fuse();
        /// let mut fut = future::empty::<()>();
        ///
        /// select! {
        ///     x = st.next() => assert_eq!(Some(2), x),
        ///     _ = fut => panic!(),
        /// };
        /// # });
        /// ```
        ///
        /// `select` also accepts a `complete` branch and a `default` branch.
        /// `complete` will run if all futures and streams have already been
        /// exhausted. `default` will run if no futures or streams are
        /// immediately ready. `complete` takes priority over `default` in
        /// the case where all futures have completed.
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::select;
        /// let mut a_fut = future::ready(4);
        /// let mut b_fut = future::ready(6);
        /// let mut total = 0;
        ///
        /// loop {
        ///     select! {
        ///         a = a_fut => total += a,
        ///         b = b_fut => total += b,
        ///         complete => break,
        ///         default => panic!(), // never runs (futures run first, then complete)
        ///     };
        /// }
        /// assert_eq!(total, 10);
        /// # });
        /// ```
        ///
        /// Note that the futures that have been matched over can still be mutated
        /// from inside the `select!` block's branches. This can be used to implement
        /// more complex behavior such as timer resets or writing into the head of
        /// a stream.
        $item
    }
}

document_select_macro! {
    #[proc_macro_hack(support_nested)]
    pub use futures_select_macro::select;
}
<commit_msg>Document branch precedence of `select!`<commit_after>//! The `select` macro.

use proc_macro_hack::proc_macro_hack;

#[doc(hidden)]
#[macro_export]
macro_rules! document_select_macro {
    ($item:item) => {
        /// Polls multiple futures and streams simultaneously, executing the branch
        /// for the future that finishes first. If multiple futures are ready,
        /// one will be pseudo-randomly selected at runtime. Futures passed to
        /// `select!` must be `Unpin` and implement `FusedFuture`.
        /// Futures and streams which are not already fused can be fused using the
        /// `.fuse()` method. Note, though, that fusing a future or stream directly
        /// in the call to `select!` will not be enough to prevent it from being
        /// polled after completion if the `select!` call is in a loop, so when
        /// `select!`ing in a loop, users should take care to `fuse()` outside of
        /// the loop.
        ///
        /// `select!` can select over futures with different output types, but each
        /// branch has to have the same return type.
        ///
        /// This macro is only usable inside of async functions, closures, and blocks.
        /// It is also gated behind the `async-await` feature of this library, which is
        /// _not_ activated by default.
        ///
        /// # Examples
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::select;
        /// let mut a = future::ready(4);
        /// let mut b = future::empty::<()>();
        ///
        /// let res = select! {
        ///     a_res = a => a_res + 1,
        ///     _ = b => 0,
        /// };
        /// assert_eq!(res, 5);
        /// # });
        /// ```
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::stream::{self, StreamExt};
        /// use futures::select;
        /// let mut st = stream::iter(vec![2]).fuse();
        /// let mut fut = future::empty::<()>();
        ///
        /// select! {
        ///     x = st.next() => assert_eq!(Some(2), x),
        ///     _ = fut => panic!(),
        /// };
        /// # });
        /// ```
        ///
        /// `select` also accepts a `complete` branch and a `default` branch.
        /// `complete` will run if all futures and streams have already been
        /// exhausted. `default` will run if no futures or streams are
        /// immediately ready. `complete` takes priority over `default` in
        /// the case where all futures have completed.
        ///
        /// ```
        /// #![feature(async_await)]
        /// # futures::executor::block_on(async {
        /// use futures::future::{self, FutureExt};
        /// use futures::select;
        /// let mut a_fut = future::ready(4);
        /// let mut b_fut = future::ready(6);
        /// let mut total = 0;
        ///
        /// loop {
        ///     select! {
        ///         a = a_fut => total += a,
        ///         b = b_fut => total += b,
        ///         complete => break,
        ///         default => panic!(), // never runs (futures run first, then complete)
        ///     };
        /// }
        /// assert_eq!(total, 10);
        /// # });
        /// ```
        ///
        /// Note that the futures that have been matched over can still be mutated
        /// from inside the `select!` block's branches. This can be used to implement
        /// more complex behavior such as timer resets or writing into the head of
        /// a stream.
        $item
    }
}

document_select_macro! {
    #[proc_macro_hack(support_nested)]
    pub use futures_select_macro::select;
}
<|endoftext|>


<commit_before>use std::path::PathBuf;

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
pub struct Loc<T, P> {
    #[serde(rename = "value")]
    inner: T,
    pos: P,
}

impl<T, P> ::std::ops::Deref for Loc<T, P> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::borrow::Borrow<T> for Loc<T, P> {
    fn borrow(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::convert::AsRef<T> for Loc<T, P> {
    fn as_ref(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::convert::AsMut<T> for Loc<T, P> {
    fn as_mut(&mut self) -> &mut T {
        &mut self.inner
    }
}

impl<T, P> Loc<T, P>
    where P: Clone
{
    pub fn new(inner: T, pos: P) -> Loc<T, P> {
        Loc {
            inner: inner,
            pos: pos,
        }
    }

    pub fn move_inner(self) -> T {
        self.inner
    }

    pub fn pos(&self) -> &P {
        &self.pos
    }

    pub fn map_into<M, U>(self, map: M) -> Loc<U, P>
        where M: FnOnce(T) -> U
    {
        Loc::new(map(self.inner), self.pos)
    }

    pub fn map<'a, M, U>(&'a self, map: M) -> Loc<U, P>
        where M: FnOnce(&'a T) -> U,
              U: 'a
    {
        Loc::new(map(&self.inner), self.pos.clone())
    }

    pub fn both(self) -> (T, P) {
        (self.inner, self.pos)
    }

    pub fn ref_both(&self) -> (&T, &P) {
        (&self.inner, &self.pos)
    }
}

pub trait WithPrefix<T> {
    type Prefix;
    type Output;

    fn with_prefix(self, prefix: Self::Prefix) -> Loc<T, Self::Output>;
}

impl<T, B, C> WithPrefix<T> for Loc<T, (B, C)>
    where T: Clone,
          B: Clone,
          C: Clone
{
    type Prefix = PathBuf;
    type Output = (PathBuf, B, C);

    fn with_prefix(self, prefix: Self::Prefix) -> Loc<T, Self::Output> {
        let (b, c) = self.pos.clone();
        Loc::new(self.inner.clone(), (prefix, b, c))
    }
}

impl<T, P> ::std::fmt::Display for Loc<T, P>
    where T: ::std::fmt::Display
{
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{}", self.inner)
    }
}

impl<T, P> ::std::fmt::Debug for Loc<T, P>
    where T: ::std::fmt::Debug,
          P: ::std::fmt::Debug
{
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "<{:?}@{:?}>", self.inner, self.pos)
    }
}
<commit_msg>drop WithPrefix implementation<commit_after>use std::path::PathBuf;

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
pub struct Loc<T, P> {
    #[serde(rename = "value")]
    inner: T,
    pos: P,
}

impl<T, P> ::std::ops::Deref for Loc<T, P> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::borrow::Borrow<T> for Loc<T, P> {
    fn borrow(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::convert::AsRef<T> for Loc<T, P> {
    fn as_ref(&self) -> &T {
        &self.inner
    }
}

impl<T, P> ::std::convert::AsMut<T> for Loc<T, P> {
    fn as_mut(&mut self) -> &mut T {
        &mut self.inner
    }
}

impl<T, P> Loc<T, P>
    where P: Clone
{
    pub fn new(inner: T, pos: P) -> Loc<T, P> {
        Loc {
            inner: inner,
            pos: pos,
        }
    }

    pub fn move_inner(self) -> T {
        self.inner
    }

    pub fn pos(&self) -> &P {
        &self.pos
    }

    pub fn map_into<M, U>(self, map: M) -> Loc<U, P>
        where M: FnOnce(T) -> U
    {
        Loc::new(map(self.inner), self.pos)
    }

    pub fn map<'a, M, U>(&'a self, map: M) -> Loc<U, P>
        where M: FnOnce(&'a T) -> U,
              U: 'a
    {
        Loc::new(map(&self.inner), self.pos.clone())
    }

    pub fn both(self) -> (T, P) {
        (self.inner, self.pos)
    }

    pub fn ref_both(&self) -> (&T, &P) {
        (&self.inner, &self.pos)
    }
}

impl<T, P> ::std::fmt::Display for Loc<T, P>
    where T: ::std::fmt::Display
{
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{}", self.inner)
    }
}

impl<T, P> ::std::fmt::Debug for Loc<T, P>
    where T: ::std::fmt::Debug,
          P: ::std::fmt::Debug
{
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "<{:?}@{:?}>", self.inner, self.pos)
    }
}
<|endoftext|>


<commit_before>use std::slice::from_raw_parts;
use std::ffi::CStr;
use std::mem::copy_lifetime;
use libc;

use htslib;

#[derive(Debug)]
pub enum Aux<'a> {
    Integer(i32),
    String(&'a [u8]),
    Float(f64),
    Char(u8),
}

impl<'a> Aux<'a> {
    pub fn string(&self) -> &'a [u8] {
        match *self {
            Aux::String(x) => x,
            _ => panic!("not a string"),
        }
    }

    pub fn float(&self) -> f64 {
        match *self {
            Aux::Float(x) => x,
            _ => panic!("not a float"),
        }
    }

    pub fn integer(&self) -> i32 {
        match *self {
            Aux::Integer(x) => x,
            _ => panic!("not an integer"),
        }
    }

    pub fn char(&self) -> u8 {
        match *self {
            Aux::Char(x) => x,
            _ => panic!("not a character"),
        }
    }
}

pub struct Record {
    b: htslib::bam1_t,
}

impl Record {
    pub fn new() -> Self {
        let b = unsafe { *htslib::bam_init1() };
        Record { b: b }
    }
    // TODO ensure that owning b really destroys b.data when object goes out of scope

    fn data(&self) -> &[u8] {
        unsafe { from_raw_parts(self.b.data, self.b.l_data as usize) }
    }

    pub fn tid(&self) -> i32 {
        self.b.core.tid
    }

    pub fn pos(&self) -> i32 {
        self.b.core.pos
    }

    pub fn bin(&self) -> u16 {
        self.b.core.bin
    }

    pub fn map_qual(&self) -> u8 {
        self.b.core.qual
    }

    pub fn flag(&self) -> u16 {
        self.b.core.flag
    }

    pub fn mtid(&self) -> i32 {
        self.b.core.mtid
    }

    pub fn mpos(&self) -> i32 {
        self.b.core.mpos
    }

    pub fn isize(&self) -> i32 {
        self.b.core.isize
    }

    fn qname_len(&self) -> usize {
        self.b.core.l_qname as usize
    }

    pub fn qname(&self) -> &[u8] {
        &self.data()[..self.qname_len()-1] // -1 ignores the termination symbol

    }

    fn cigar_len(&self) -> usize {
        self.b.core.n_cigar as usize
    }

    pub fn cigar(&self) -> Vec<Cigar> {
        let raw = unsafe { from_raw_parts(self.data()[self.qname_len()..].as_ptr() as *const u32, self.cigar_len()) };
        raw.iter().map(|&c| {
            let len = c >> 4;
            match c & 0b1111 {
                0 => Cigar::Match(len),
                1 => Cigar::Ins(len),
                2 => Cigar::Del(len),
                3 => Cigar::RefSkip(len),
                4 => Cigar::SoftClip(len),
                5 => Cigar::HardClip(len),
                6 => Cigar::Pad(len),
                7 => Cigar::Equal(len),
                8 => Cigar::Diff(len),
                9 => Cigar::Back(len),
                _ => panic!("Unexpected cigar type"),
            }
        }).collect()
    }

    fn seq_len(&self) -> usize {
        self.b.core.l_qseq as usize
    }

    pub fn seq(&self) -> Seq {
        Seq {
            encoded: &self.data()
                        [self.qname_len() + self.cigar_len()*4..]
                        [..(self.seq_len() + 1) / 2]
        }
    }

    pub fn qual(&self) -> &[u8] {
        &self.data()[self.qname_len() + self.cigar_len()*4 + (self.seq_len()+1)/2..][..self.seq_len()]
    }

    pub fn aux(&self, name: &[u8]) -> Option<Aux> {
        let aux = unsafe { htslib::bam_aux_get(&self.b, name.as_ptr() as *mut i8 ) };

        unsafe {
            match *aux {
                b'c'|b'C'|b's'|b'S'|b'i'|b'I' => Some(Aux::Integer(htslib::bam_aux2i(aux))),
                b'f'|b'd' => Some(Aux::Float(htslib::bam_aux2f(aux))),
                b'A' => Some(Aux::Char(htslib::bam_aux2A(aux) as u8)),
                b'Z'|b'H' => {
                    let f = aux.offset(1) as *const i8;
                    let x = CStr::from_ptr(f).to_bytes();
                    Some(Aux::String(copy_lifetime(self, x)))
                },
                _ => None,
            }
        }
    }

    pub fn is_paired(&self) -> bool{
        (self.flag() & 1u16) > 0
    }

    pub fn is_proper_pair(&self) -> bool{
        (self.flag() & 2u16) > 0
    }

    pub fn is_unmapped(&self) -> bool{
        (self.flag() & 4u16) > 0
    }

    pub fn mate_is_unmapped(&self) -> bool{
        (self.flag() & 8u16) > 0
    }

    pub fn is_reverse(&self) -> bool{
        (self.flag() & 16u16) > 0
    }

    pub fn mate_is_reverse(&self) -> bool{
        (self.flag() & 32u16) > 0
    }

    pub fn is_first_in_pair(&self) -> bool{
        (self.flag() & 64u16) > 0
    }

    pub fn is_second_in_pair(&self) -> bool{
        (self.flag() & 128u16) > 0
    }

    pub fn is_secondary(&self) -> bool{
        (self.flag() & 256u16) > 0
    }

    pub fn is_quality_check_failed(&self) -> bool{
        (self.flag() & 512u16) > 0
    }

    pub fn is_duplicate(&self) -> bool{
        (self.flag() & 1024u16) > 0
    }

    pub fn is_supplementary(&self) -> bool{
        (self.flag() & 2048u16) > 0
    }
}

impl Drop for Record {
    fn drop(&mut self) {
        unsafe { libc::funcs::c95::stdlib::free(self.b.data as *mut libc::types::common::c95::c_void) };
    }
}

pub struct Bamfile {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl Bamfile{
     pub fn new(filename: &[u8]) -> Self {
        let f = unsafe { htslib::bgzf_open(filename.as_ptr() as *const i8, b"r\0".as_ptr() as *const i8) };
        let header = unsafe { htslib::bam_hdr_read(f) };
        Bamfile { f : f, header : header }
    }

    pub fn read(&self, record: &mut Record) -> Result<(), Error> {
        match unsafe { htslib::bam_read1(self.f, &mut record.b as *mut htslib::bam1_t) } {
            -1 => Err(Error::EOF),
            -2 => Err(Error::Truncated),
            -4 => Err(Error::Invalid),
            _  => Ok(())
        }
    }

    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl Drop for Bamfile {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

pub struct Records {
    bam: Bamfile
}

impl Iterator for Records {
    type Item = Result<Record, Error>;

    fn next(&mut self) -> Option<Result<Record, Error>> {
        let mut record = Record::new();
        match self.bam.read(&mut record) {
            Err(Error::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

static TRANSLATE_BASE: &'static [u8] = b"=ACMGRSVTWYHKDBN";

pub struct Seq<'a> {
    pub encoded: &'a [u8]
}

impl<'a> Seq<'a> {

    #[inline]
    pub fn encoded_base(&self, i: usize) -> u8 {
        (self.encoded[i / 2] >> ((! i & 1) << 2)) & 0b1111
    }

    #[inline]
    pub fn base(&self, i: usize) -> u8 {
        TRANSLATE_BASE[self.encoded_base(i) as usize]
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        (0..self.len()).map(|i| self.base(i)).collect()
    }

    pub fn len(&self) -> usize {
        self.encoded.len() * 2
    }
}

pub enum Cigar {
    Match(u32),  // M
    Ins(u32),  // I
    Del(u32),  // D
    RefSkip(u32),  // N
    SoftClip(u32),  // S
    HardClip(u32),  // H
    Pad(u32),  // P
    Equal(u32),  // =
    Diff(u32),  // X
    Back(u32)  // B
}

pub enum Error {
    Truncated,
    Invalid,
    EOF,
}

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

    #[test]
    fn test_record() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];

        let bam = Bamfile::new(b"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flag(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
        }

    }
}
<commit_msg>Test cigar parsing.<commit_after>use std::slice::from_raw_parts;
use std::ffi::CStr;
use std::mem::copy_lifetime;
use libc;

use htslib;

#[derive(Debug)]
pub enum Aux<'a> {
    Integer(i32),
    String(&'a [u8]),
    Float(f64),
    Char(u8),
}

impl<'a> Aux<'a> {
    pub fn string(&self) -> &'a [u8] {
        match *self {
            Aux::String(x) => x,
            _ => panic!("not a string"),
        }
    }

    pub fn float(&self) -> f64 {
        match *self {
            Aux::Float(x) => x,
            _ => panic!("not a float"),
        }
    }

    pub fn integer(&self) -> i32 {
        match *self {
            Aux::Integer(x) => x,
            _ => panic!("not an integer"),
        }
    }

    pub fn char(&self) -> u8 {
        match *self {
            Aux::Char(x) => x,
            _ => panic!("not a character"),
        }
    }
}

pub struct Record {
    b: htslib::bam1_t,
}

impl Record {
    pub fn new() -> Self {
        let b = unsafe { *htslib::bam_init1() };
        Record { b: b }
    }
    // TODO ensure that owning b really destroys b.data when object goes out of scope

    fn data(&self) -> &[u8] {
        unsafe { from_raw_parts(self.b.data, self.b.l_data as usize) }
    }

    pub fn tid(&self) -> i32 {
        self.b.core.tid
    }

    pub fn pos(&self) -> i32 {
        self.b.core.pos
    }

    pub fn bin(&self) -> u16 {
        self.b.core.bin
    }

    pub fn map_qual(&self) -> u8 {
        self.b.core.qual
    }

    pub fn flag(&self) -> u16 {
        self.b.core.flag
    }

    pub fn mtid(&self) -> i32 {
        self.b.core.mtid
    }

    pub fn mpos(&self) -> i32 {
        self.b.core.mpos
    }

    pub fn isize(&self) -> i32 {
        self.b.core.isize
    }

    fn qname_len(&self) -> usize {
        self.b.core.l_qname as usize
    }

    pub fn qname(&self) -> &[u8] {
        &self.data()[..self.qname_len()-1] // -1 ignores the termination symbol

    }

    fn cigar_len(&self) -> usize {
        self.b.core.n_cigar as usize
    }

    pub fn cigar(&self) -> Vec<Cigar> {
        let raw = unsafe { from_raw_parts(self.data()[self.qname_len()..].as_ptr() as *const u32, self.cigar_len()) };
        raw.iter().map(|&c| {
            let len = c >> 4;
            match c & 0b1111 {
                0 => Cigar::Match(len),
                1 => Cigar::Ins(len),
                2 => Cigar::Del(len),
                3 => Cigar::RefSkip(len),
                4 => Cigar::SoftClip(len),
                5 => Cigar::HardClip(len),
                6 => Cigar::Pad(len),
                7 => Cigar::Equal(len),
                8 => Cigar::Diff(len),
                9 => Cigar::Back(len),
                _ => panic!("Unexpected cigar type"),
            }
        }).collect()
    }

    fn seq_len(&self) -> usize {
        self.b.core.l_qseq as usize
    }

    pub fn seq(&self) -> Seq {
        Seq {
            encoded: &self.data()
                        [self.qname_len() + self.cigar_len()*4..]
                        [..(self.seq_len() + 1) / 2]
        }
    }

    pub fn qual(&self) -> &[u8] {
        &self.data()[self.qname_len() + self.cigar_len()*4 + (self.seq_len()+1)/2..][..self.seq_len()]
    }

    pub fn aux(&self, name: &[u8]) -> Option<Aux> {
        let aux = unsafe { htslib::bam_aux_get(&self.b, name.as_ptr() as *mut i8 ) };

        unsafe {
            match *aux {
                b'c'|b'C'|b's'|b'S'|b'i'|b'I' => Some(Aux::Integer(htslib::bam_aux2i(aux))),
                b'f'|b'd' => Some(Aux::Float(htslib::bam_aux2f(aux))),
                b'A' => Some(Aux::Char(htslib::bam_aux2A(aux) as u8)),
                b'Z'|b'H' => {
                    let f = aux.offset(1) as *const i8;
                    let x = CStr::from_ptr(f).to_bytes();
                    Some(Aux::String(copy_lifetime(self, x)))
                },
                _ => None,
            }
        }
    }

    pub fn is_paired(&self) -> bool{
        (self.flag() & 1u16) > 0
    }

    pub fn is_proper_pair(&self) -> bool{
        (self.flag() & 2u16) > 0
    }

    pub fn is_unmapped(&self) -> bool{
        (self.flag() & 4u16) > 0
    }

    pub fn mate_is_unmapped(&self) -> bool{
        (self.flag() & 8u16) > 0
    }

    pub fn is_reverse(&self) -> bool{
        (self.flag() & 16u16) > 0
    }

    pub fn mate_is_reverse(&self) -> bool{
        (self.flag() & 32u16) > 0
    }

    pub fn is_first_in_pair(&self) -> bool{
        (self.flag() & 64u16) > 0
    }

    pub fn is_second_in_pair(&self) -> bool{
        (self.flag() & 128u16) > 0
    }

    pub fn is_secondary(&self) -> bool{
        (self.flag() & 256u16) > 0
    }

    pub fn is_quality_check_failed(&self) -> bool{
        (self.flag() & 512u16) > 0
    }

    pub fn is_duplicate(&self) -> bool{
        (self.flag() & 1024u16) > 0
    }

    pub fn is_supplementary(&self) -> bool{
        (self.flag() & 2048u16) > 0
    }
}

impl Drop for Record {
    fn drop(&mut self) {
        unsafe { libc::funcs::c95::stdlib::free(self.b.data as *mut libc::types::common::c95::c_void) };
    }
}

pub struct Bamfile {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl Bamfile{
     pub fn new(filename: &[u8]) -> Self {
        let f = unsafe { htslib::bgzf_open(filename.as_ptr() as *const i8, b"r\0".as_ptr() as *const i8) };
        let header = unsafe { htslib::bam_hdr_read(f) };
        Bamfile { f : f, header : header }
    }

    pub fn read(&self, record: &mut Record) -> Result<(), Error> {
        match unsafe { htslib::bam_read1(self.f, &mut record.b as *mut htslib::bam1_t) } {
            -1 => Err(Error::EOF),
            -2 => Err(Error::Truncated),
            -4 => Err(Error::Invalid),
            _  => Ok(())
        }
    }

    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl Drop for Bamfile {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

pub struct Records {
    bam: Bamfile
}

impl Iterator for Records {
    type Item = Result<Record, Error>;

    fn next(&mut self) -> Option<Result<Record, Error>> {
        let mut record = Record::new();
        match self.bam.read(&mut record) {
            Err(Error::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

static TRANSLATE_BASE: &'static [u8] = b"=ACMGRSVTWYHKDBN";

pub struct Seq<'a> {
    pub encoded: &'a [u8]
}

impl<'a> Seq<'a> {

    #[inline]
    pub fn encoded_base(&self, i: usize) -> u8 {
        (self.encoded[i / 2] >> ((! i & 1) << 2)) & 0b1111
    }

    #[inline]
    pub fn base(&self, i: usize) -> u8 {
        TRANSLATE_BASE[self.encoded_base(i) as usize]
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        (0..self.len()).map(|i| self.base(i)).collect()
    }

    pub fn len(&self) -> usize {
        self.encoded.len() * 2
    }
}

#[derive(Debug)]
#[derive(PartialEq)]
pub enum Cigar {
    Match(u32),  // M
    Ins(u32),  // I
    Del(u32),  // D
    RefSkip(u32),  // N
    SoftClip(u32),  // S
    HardClip(u32),  // H
    Pad(u32),  // P
    Equal(u32),  // =
    Diff(u32),  // X
    Back(u32)  // B
}

pub enum Error {
    Truncated,
    Invalid,
    EOF,
}

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

    #[test]
    fn test_record() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = Bamfile::new(b"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flag(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }

    }
}
<|endoftext|>


<commit_before>use item_state::ItemState;
use bar_builder::{
    UpdateStream,
};
use std::rc::Rc;
use xcb::{
    Connection,
    Rectangle,
    Window,
};
use futures::stream::{Merge, MergedItem};
use futures::{Future, Stream};
use xcb_event_stream::XcbEventStream;
use error::*;
use std::error::Error;
use xcb;
use component::Slot;

type UpdateAndEventStream = Merge<UpdateStream, XcbEventStream>;

pub struct Bar {
    pub center_items: Vec<ItemState>,
    pub conn: Rc<Connection>,
    pub foreground: u32,
    pub geometry: Rectangle,
    pub item_positions: Vec<(i16, u16)>,
    pub left_items: Vec<ItemState>,
    pub right_items: Vec<ItemState>,
    pub stream: Option<UpdateAndEventStream>,
    pub window: Window,
}

impl Bar {
    fn get_stream(&mut self) -> UpdateAndEventStream {
        ::std::mem::replace(&mut self.stream, None).unwrap()
    }

    pub fn run(mut self) -> Box<Future<Item=(), Error=()>> {
        let future = self.get_stream()
            .map_err(|_| ErrorKind::ItemError.into())
            .for_each(move |item| -> Result<()> {
                let (_, update) = match item {
                    MergedItem::First(update) => (None, Some(update)),
                    MergedItem::Second(event) => (Some(event), None),
                    MergedItem::Both(update, event) => (Some(event), Some(update)),
                };

                if let Some(update) = update {
                    let size_changed;
                    let width;

                    {
                        let slot_items = match update.slot {
                            Slot::Left => &mut self.left_items,
                            Slot::Center => &mut self.center_items,
                            Slot::Right => &mut self.right_items,
                        };
                        slot_items[update.index].update(update.value)?;

                        width = slot_items[update.index].get_content_width();
                        size_changed = self.item_positions[update.id].1 != width;
                    }

                    if size_changed {
                        self.paint_bg(update.id)?;
                    }

                    self.item_positions[update.id].1 = width;

                    match update.slot {
                        Slot::Center => self.redraw_center()?,
                        Slot::Left => self.redraw_end(false, size_changed, update.index)?,
                        Slot::Right => self.redraw_end(true, size_changed, update.index)?,
                    }
                }
                Ok(())
            })
            .map_err(|err| {
                println!("Error occurred: {}", err);
                let mut cause = err.cause();
                while let Some(err) = cause {
                    println!("  Caused by: {}", err);
                    cause = err.cause();
                }
            });

        Box::new(future)
    }

    fn redraw_center(&mut self) -> Result<()> {
        let width_all: i16 = self.center_items.iter()
            .map(|item| item.get_content_width() as i16)
            .sum();

        let mut pos = (self.geometry.width() as i16)/2-width_all/2;

        for item in self.center_items.iter() {
            self.item_positions[item.get_id()].0 = pos;
            self.draw_item(item, pos)?;
            pos += item.get_content_width() as i16;
        }

        Ok(())
    }

    fn redraw_end(&mut self, mut side: bool, size_changed: bool, index: usize) -> Result<()> {
        let (start, direction, items) = match side {
            false => (0, 1, &self.left_items),
            true => (self.geometry.width() as i16, -1, &self.right_items),
        };

        let (before, mut after) = items.split_at(index);

        if !size_changed {
            after = &after[0..1];
        }

        let mut pos: i16 = before.iter()
            .map(|item| item.get_content_width() as i16).sum();

        for item in after.iter() {
            if side {
                pos += item.get_content_width() as i16;
            }
            side = true;
            self.item_positions[item.get_id()].0 = pos;
            self.draw_item(item, start+(pos*direction))?;
        }

        self.conn.flush();

        Ok(())
    }

    fn draw_item(&self, item: &ItemState, pos: i16) -> Result<()> {
        if !item.is_ready() {
            return Ok(());
        }
        self.paint_bg(item.get_id())?;
        try_xcb!(xcb::copy_area_checked, "failed to copy pixmap",
            &self.conn,
            item.get_pixmap(),
            self.window,
            self.foreground,
            0, 0,
            pos, 0,
            item.get_content_width() as u16,
            self.geometry.height());
        Ok(())
    }

    fn paint_bg(&self, index: usize) -> Result<()> {
        let x = self.item_positions[index].0;
        let w = self.item_positions[index].1;

        try_xcb!(xcb::poly_fill_rectangle, "failed to draw background",
            &self.conn,
            self.window,
            self.foreground,
            &[Rectangle::new(x, 0, w, self.geometry.height())]);

        Ok(())
    }
}
<commit_msg>Carry errors returned from Components.<commit_after>use item_state::ItemState;
use bar_builder::{
    UpdateStream,
};
use std::rc::Rc;
use xcb::{
    Connection,
    Rectangle,
    Window,
};
use futures::stream::{Merge, MergedItem};
use futures::{Future, Stream};
use xcb_event_stream::XcbEventStream;
use error::*;
use std::error::Error;
use xcb;
use component::Slot;

type UpdateAndEventStream = Merge<UpdateStream, XcbEventStream>;

pub struct Bar {
    pub center_items: Vec<ItemState>,
    pub conn: Rc<Connection>,
    pub foreground: u32,
    pub geometry: Rectangle,
    pub item_positions: Vec<(i16, u16)>,
    pub left_items: Vec<ItemState>,
    pub right_items: Vec<ItemState>,
    pub stream: Option<UpdateAndEventStream>,
    pub window: Window,
}

impl Bar {
    fn get_stream(&mut self) -> UpdateAndEventStream {
        ::std::mem::replace(&mut self.stream, None).unwrap()
    }

    pub fn run(mut self) -> Box<Future<Item=(), Error=()>> {
        let future = self.get_stream()
            .map_err(|e| ::error::Error::with_chain(e, ErrorKind::ItemError))
            .for_each(move |item| -> Result<()> {
                let (_, update) = match item {
                    MergedItem::First(update) => (None, Some(update)),
                    MergedItem::Second(event) => (Some(event), None),
                    MergedItem::Both(update, event) => (Some(event), Some(update)),
                };

                if let Some(update) = update {
                    let size_changed;
                    let width;

                    {
                        let slot_items = match update.slot {
                            Slot::Left => &mut self.left_items,
                            Slot::Center => &mut self.center_items,
                            Slot::Right => &mut self.right_items,
                        };
                        slot_items[update.index].update(update.value)?;

                        width = slot_items[update.index].get_content_width();
                        size_changed = self.item_positions[update.id].1 != width;
                    }

                    if size_changed {
                        self.paint_bg(update.id)?;
                    }

                    self.item_positions[update.id].1 = width;

                    match update.slot {
                        Slot::Center => self.redraw_center()?,
                        Slot::Left => self.redraw_end(false, size_changed, update.index)?,
                        Slot::Right => self.redraw_end(true, size_changed, update.index)?,
                    }
                }
                Ok(())
            })
            .map_err(|err| {
                println!("Error occurred: {}", err);
                let mut cause = err.cause();
                while let Some(err) = cause {
                    println!("  Caused by: {}", err);
                    cause = err.cause();
                }
            });

        Box::new(future)
    }

    fn redraw_center(&mut self) -> Result<()> {
        let width_all: i16 = self.center_items.iter()
            .map(|item| item.get_content_width() as i16)
            .sum();

        let mut pos = (self.geometry.width() as i16)/2-width_all/2;

        for item in self.center_items.iter() {
            self.item_positions[item.get_id()].0 = pos;
            self.draw_item(item, pos)?;
            pos += item.get_content_width() as i16;
        }

        Ok(())
    }

    fn redraw_end(&mut self, mut side: bool, size_changed: bool, index: usize) -> Result<()> {
        let (start, direction, items) = match side {
            false => (0, 1, &self.left_items),
            true => (self.geometry.width() as i16, -1, &self.right_items),
        };

        let (before, mut after) = items.split_at(index);

        if !size_changed {
            after = &after[0..1];
        }

        let mut pos: i16 = before.iter()
            .map(|item| item.get_content_width() as i16).sum();

        for item in after.iter() {
            if side {
                pos += item.get_content_width() as i16;
            }
            side = true;
            self.item_positions[item.get_id()].0 = pos;
            self.draw_item(item, start+(pos*direction))?;
        }

        self.conn.flush();

        Ok(())
    }

    fn draw_item(&self, item: &ItemState, pos: i16) -> Result<()> {
        if !item.is_ready() {
            return Ok(());
        }
        self.paint_bg(item.get_id())?;
        try_xcb!(xcb::copy_area_checked, "failed to copy pixmap",
            &self.conn,
            item.get_pixmap(),
            self.window,
            self.foreground,
            0, 0,
            pos, 0,
            item.get_content_width() as u16,
            self.geometry.height());
        Ok(())
    }

    fn paint_bg(&self, index: usize) -> Result<()> {
        let x = self.item_positions[index].0;
        let w = self.item_positions[index].1;

        try_xcb!(xcb::poly_fill_rectangle, "failed to draw background",
            &self.conn,
            self.window,
            self.foreground,
            &[Rectangle::new(x, 0, w, self.geometry.height())]);

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicUsize, Ordering};

use actix::prelude::*;
use futures::{compat::*, prelude::*};
use log::*;
use typemap::{Key, SendMap};

pub struct Bus {
    map: SendMap,
}

impl Bus {
    pub fn subscribe<A, M>(addr: Addr<A>) -> impl Future<Output = Result<BusId, MailboxError>>
    where
        A: Actor + Handler<M>,
        A::Context: actix::dev::ToEnvelope<A, M>,
        M: Message + Send + Clone + 'static,
        M::Result: Send,
    {
        let bus = Bus::from_registry();
        bus.send(Subscribe::new(addr.recipient())).compat()
    }
}

impl Default for Bus {
    fn default() -> Self {
        Bus { map: SendMap::custom() }
    }
}

impl Actor for Bus {
    type Context = Context<Self>;
}

impl Supervised for Bus {}
impl SystemService for Bus {}

static LAST_ID: AtomicUsize = AtomicUsize::new(1);

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct BusId(NonZeroUsize);

impl BusId {
    pub fn new() -> Self {
        let id = LAST_ID.fetch_add(1, Ordering::SeqCst);
        assert_ne!(id, 0);
        unsafe { BusId(NonZeroUsize::new_unchecked(id)) }
    }

    pub fn publish<M>(&self, message: M)
    where
        M: Message + Send + Clone + 'static,
        M::Result: Send,
    {
        let bus = Bus::from_registry();
        bus.do_send(Publish { message, sender: Some(*self) });
    }
}

pub struct Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    pub recipient: Recipient<M>,
}

impl<M> Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    pub fn new(recipient: Recipient<M>) -> Self {
        Subscribe { recipient }
    }
}

impl<M> Message for Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Result = BusId;
}

impl<M> Handler<Subscribe<M>> for Bus
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Result = MessageResult<Subscribe<M>>;

    fn handle(&mut self, msg: Subscribe<M>, _: &mut Self::Context) -> Self::Result {
        let list = self.map.entry::<SubscriptionList<M>>().or_insert_with(Default::default);
        let id = BusId::new();
        list.add(msg.recipient);
        MessageResult(id)
    }
}

pub struct Publish<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    pub message: M,
    sender: Option<BusId>,
}

impl<M> Message for Publish<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Result = ();
}

impl<M> Handler<Publish<M>> for Bus
where
    M: Message + Send + Clone + 'static,
    M::Result: Send,
{
    type Result = ();

    fn handle(&mut self, msg: Publish<M>, _: &mut Self::Context) -> Self::Result {
        let list = self.map.entry::<SubscriptionList<M>>().or_insert_with(Default::default);
        list.send(msg.message, msg.sender);
    }
}

pub struct SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    subscribers: Vec<(BusId, Recipient<M>)>,
}

impl<M> SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    pub fn new() -> Self {
        SubscriptionList { subscribers: Vec::new() }
    }

    pub fn add(&mut self, recipient: Recipient<M>) -> BusId {
        let id = BusId::new();
        self.subscribers.push((id, recipient));
        id
    }
}

impl<M> SubscriptionList<M>
where
    M: actix::Message + Send + Clone,
    M::Result: Send,
{
    pub fn send(&mut self, msg: M, sender: Option<BusId>) {
        self.subscribers.retain(|(id, s)| {
            if sender == Some(*id) {
                return true;
            }
            match s.do_send(msg.clone()) {
                Ok(_) => { true }
                Err(SendError::Full(_)) => {
                    warn!("mailbox is full");
                    true
                }
                Err(SendError::Closed(_)) => {
                    false
                }
            }
        });
    }
}

impl<M> Default for SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    fn default() -> Self { Self::new() }
}

impl<M> Key for SubscriptionList<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Value = Self;
}

#[cfg(test)]
mod test {
    use std::collections::BTreeSet;
    use std::sync::mpsc::channel;
    use std::thread;
    use std::time::Duration;

    use rand::{self, Rng};

    use super::*;

    #[test]
    fn bus_id_uniqueness() {
        let n = 100;
        let mut rng = rand::thread_rng();
        let mut handle = vec![];
        let rx = {
            let (tx, rx) = channel();
            for _ in 0..n {
                let d = Duration::from_millis(rng.gen_range(0, 1000));
                let tx = tx.clone();
                let h = thread::spawn(move || {
                    thread::sleep(d);
                    let id = BusId::new();
                    tx.send(id).unwrap();
                });
                handle.push(h);
            }
            rx
        };
        for h in handle {
            h.join().unwrap();
        }
        let ids: Vec<_> = rx.into_iter().collect();
        assert_eq!(ids.len(), n);
        let set: BTreeSet<_> = ids.iter().map(|i| i.0).collect();
        assert_eq!(set.len(), ids.len());
    }
}
<commit_msg>Fix broken sender check<commit_after>use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicUsize, Ordering};

use actix::prelude::*;
use futures::{compat::*, prelude::*};
use log::*;
use typemap::{Key, SendMap};

pub struct Bus {
    map: SendMap,
}

impl Bus {
    pub fn subscribe<A, M>(addr: Addr<A>) -> impl Future<Output = Result<BusId, MailboxError>>
    where
        A: Actor + Handler<M>,
        A::Context: actix::dev::ToEnvelope<A, M>,
        M: Message + Send + Clone + 'static,
        M::Result: Send,
    {
        let bus = Bus::from_registry();
        bus.send(Subscribe::new(addr.recipient())).compat()
    }
}

impl Default for Bus {
    fn default() -> Self {
        Bus { map: SendMap::custom() }
    }
}

impl Actor for Bus {
    type Context = Context<Self>;

    fn started(&mut self, _: &mut Self::Context) {
        debug!("Bus::started");
    }

    fn stopping(&mut self, _: &mut Self::Context) -> Running {
        debug!("Bus::stopping");
        Running::Stop
    }

    fn stopped(&mut self, _: &mut Self::Context) {
        debug!("Bus::stopped");
    }
}

impl Supervised for Bus {
    fn restarting(&mut self, _: &mut Self::Context) {
        debug!("Bus::restarting");
    }
}

impl SystemService for Bus {
    fn service_started(&mut self, _: &mut Self::Context) {
        debug!("Bus::service_started");
    }
}

static LAST_ID: AtomicUsize = AtomicUsize::new(1);

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct BusId(NonZeroUsize);

impl BusId {
    pub fn new() -> Self {
        let id = LAST_ID.fetch_add(1, Ordering::SeqCst);
        assert_ne!(id, 0);
        unsafe { BusId(NonZeroUsize::new_unchecked(id)) }
    }

    pub fn publish<M>(&self, message: M)
    where
        M: Message + Send + Clone + 'static,
        M::Result: Send,
    {
        let bus = Bus::from_registry();
        bus.do_send(Publish { message, sender: Some(*self) });
    }
}

pub struct Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    pub recipient: Recipient<M>,
}

impl<M> Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    pub fn new(recipient: Recipient<M>) -> Self {
        Subscribe { recipient }
    }
}

impl<M> Message for Subscribe<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Result = BusId;
}

impl<M> Handler<Subscribe<M>> for Bus
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Result = MessageResult<Subscribe<M>>;

    fn handle(&mut self, msg: Subscribe<M>, _: &mut Self::Context) -> Self::Result {
        debug!("Bus received Subscribe<M>");
        let list = self.map.entry::<SubscriptionList<M>>().or_insert_with(Default::default);
        let id = list.add(msg.recipient);
        MessageResult(id)
    }
}

pub struct Publish<M> {
    pub message: M,
    sender: Option<BusId>,
}

impl<M> Message for Publish<M> {
    type Result = ();
}

impl<M> Handler<Publish<M>> for Bus
where
    M: Message + Send + Clone + 'static,
    M::Result: Send,
{
    type Result = ();

    fn handle(&mut self, msg: Publish<M>, _: &mut Self::Context) -> Self::Result {
        debug!("Bus received Publish<M>");
        let list = self.map.entry::<SubscriptionList<M>>().or_insert_with(Default::default);
        list.send(msg.message, msg.sender);
    }
}

pub struct SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    subscribers: Vec<(BusId, Recipient<M>)>,
}

impl<M> SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    pub fn new() -> Self {
        SubscriptionList { subscribers: Vec::new() }
    }

    pub fn add(&mut self, recipient: Recipient<M>) -> BusId {
        let id = BusId::new();
        self.subscribers.push((id, recipient));
        id
    }
}

impl<M> SubscriptionList<M>
where
    M: actix::Message + Send + Clone,
    M::Result: Send,
{
    pub fn send(&mut self, msg: M, sender: Option<BusId>) {
        self.subscribers.retain(|(id, s)| {
            if sender == Some(*id) {
                return true;
            }
            match s.do_send(msg.clone()) {
                Ok(_) => { true }
                Err(SendError::Full(_)) => {
                    warn!("mailbox is full");
                    true
                }
                Err(SendError::Closed(_)) => {
                    false
                }
            }
        });
    }
}

impl<M> Default for SubscriptionList<M>
where
    M: actix::Message + Send,
    M::Result: Send,
{
    fn default() -> Self { Self::new() }
}

impl<M> Key for SubscriptionList<M>
where
    M: Message + Send + 'static,
    M::Result: Send,
{
    type Value = Self;
}

#[cfg(test)]
mod test {
    use std::collections::BTreeSet;
    use std::sync::mpsc::channel;
    use std::thread;
    use std::time::Duration;

    use rand::{self, Rng};

    use super::*;

    #[test]
    fn bus_id_uniqueness() {
        let n = 100;
        let mut rng = rand::thread_rng();
        let mut handle = vec![];
        let rx = {
            let (tx, rx) = channel();
            for _ in 0..n {
                let d = Duration::from_millis(rng.gen_range(0, 1000));
                let tx = tx.clone();
                let h = thread::spawn(move || {
                    thread::sleep(d);
                    let id = BusId::new();
                    tx.send(id).unwrap();
                });
                handle.push(h);
            }
            rx
        };
        for h in handle {
            h.join().unwrap();
        }
        let ids: Vec<_> = rx.into_iter().collect();
        assert_eq!(ids.len(), n);
        let set: BTreeSet<_> = ids.iter().map(|i| i.0).collect();
        assert_eq!(set.len(), ids.len());
    }
}
<|endoftext|>