<commit_before>//! Trait solving using Chalk.

use std::env::var;

use chalk_ir::GoalData;
use chalk_recursive::Cache;
use chalk_solve::{logging_db::LoggingRustIrDatabase, Solver};

use base_db::CrateId;
use hir_def::{lang_item::LangItemTarget, TraitId};
use stdx::panic_context;
use syntax::SmolStr;

use crate::{
    db::HirDatabase, AliasEq, AliasTy, Canonical, DomainGoal, Goal, Guidance, InEnvironment,
    Interner, Solution, TraitRefExt, Ty, TyKind, WhereClause,
};

/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
pub(crate) struct ChalkContext<'a> {
    pub(crate) db: &'a dyn HirDatabase,
    pub(crate) krate: CrateId,
}

fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
    let overflow_depth =
        var("CHALK_OVERFLOW_DEPTH").ok().and_then(|s| s.parse().ok()).unwrap_or(300);
    let max_size = var("CHALK_SOLVER_MAX_SIZE").ok().and_then(|s| s.parse().ok()).unwrap_or(30);
    chalk_recursive::RecursiveSolver::new(overflow_depth, max_size, Some(Cache::new()))
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub krate: CrateId,
    // When we're using Chalk's Ty we can make this a BTreeMap since it's Ord,
    // but for now it's too annoying...
    pub(crate) traits_from_clauses: Vec<(Ty, TraitId)>,
    pub env: chalk_ir::Environment<Interner>,
}

impl TraitEnvironment {
    pub fn empty(krate: CrateId) -> Self {
        TraitEnvironment {
            krate,
            traits_from_clauses: Vec::new(),
            env: chalk_ir::Environment::new(Interner),
        }
    }

    pub fn traits_in_scope_from_clauses<'a>(
        &'a self,
        ty: Ty,
    ) -> impl Iterator<Item = TraitId> + 'a {
        self.traits_from_clauses
            .iter()
            .filter_map(move |(self_ty, trait_id)| (*self_ty == ty).then(|| *trait_id))
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Goal>>,
) -> Option<Solution> {
    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.goal.data(Interner) {
        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::Implemented(it))) => {
            db.trait_data(it.hir_trait_id()).name.to_string()
        }
        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(_))) => "alias_eq".to_string(),
        _ => "??".to_string(),
    });
    tracing::info!("trait_solve_query({:?})", goal.value.goal);

    if let GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(AliasEq {
        alias: AliasTy::Projection(projection_ty),
        ..
    }))) = &goal.value.goal.data(Interner)
    {
        if let TyKind::BoundVar(_) = projection_ty.self_type_parameter(Interner).kind(Interner) {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical: goal, universes: 1 };
    solve(db, krate, &u_canonical)
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    tracing::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let should_continue = || {
        db.unwind_if_cancelled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            tracing::debug!("fuel exhausted");
        }
        remaining > 0
    };

    let mut solve = || {
        let _ctx = if is_chalk_debug() || is_chalk_print() {
            Some(panic_context::enter(format!("solving {:?}", goal)))
        } else {
            None
        };
        let solution = if is_chalk_print() {
            let logging_db =
                LoggingRustIrDatabaseLoggingOnDrop(LoggingRustIrDatabase::new(context));
            solver.solve_limited(&logging_db.0, goal, &should_continue)
        } else {
            solver.solve_limited(&context, goal, &should_continue)
        };

        tracing::debug!("solve({:?}) => {:?}", goal, solution);

        solution
    };

    // don't set the TLS for Chalk unless Chalk debugging is active, to make
    // extra sure we only use it for debugging
    if is_chalk_debug() {
        crate::tls::set_current_program(db, solve)
    } else {
        solve()
    }
}

struct LoggingRustIrDatabaseLoggingOnDrop<'a>(LoggingRustIrDatabase<Interner, ChalkContext<'a>>);

impl<'a> Drop for LoggingRustIrDatabaseLoggingOnDrop<'a> {
    fn drop(&mut self) {
        eprintln!("chalk program:\n{}", self.0);
    }
}

fn is_chalk_debug() -> bool {
    std::env::var("CHALK_DEBUG").is_ok()
}

fn is_chalk_print() -> bool {
    std::env::var("CHALK_PRINT").is_ok()
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    const fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }

    pub fn get_id(&self, db: &dyn HirDatabase, krate: CrateId) -> Option<TraitId> {
        let target = db.lang_item(krate, SmolStr::new_inline(self.lang_item_name()))?;
        match target {
            LangItemTarget::TraitId(t) => Some(t),
            _ => None,
        }
    }
}
<commit_msg>Bump default CHALK_SOLVER_MAX_SIZE to 150<commit_after>//! Trait solving using Chalk.

use std::env::var;

use chalk_ir::GoalData;
use chalk_recursive::Cache;
use chalk_solve::{logging_db::LoggingRustIrDatabase, Solver};

use base_db::CrateId;
use hir_def::{lang_item::LangItemTarget, TraitId};
use stdx::panic_context;
use syntax::SmolStr;

use crate::{
    db::HirDatabase, AliasEq, AliasTy, Canonical, DomainGoal, Goal, Guidance, InEnvironment,
    Interner, Solution, TraitRefExt, Ty, TyKind, WhereClause,
};

/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
pub(crate) struct ChalkContext<'a> {
    pub(crate) db: &'a dyn HirDatabase,
    pub(crate) krate: CrateId,
}

fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
    let overflow_depth =
        var("CHALK_OVERFLOW_DEPTH").ok().and_then(|s| s.parse().ok()).unwrap_or(300);
    let max_size = var("CHALK_SOLVER_MAX_SIZE").ok().and_then(|s| s.parse().ok()).unwrap_or(150);
    chalk_recursive::RecursiveSolver::new(overflow_depth, max_size, Some(Cache::new()))
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub krate: CrateId,
    // When we're using Chalk's Ty we can make this a BTreeMap since it's Ord,
    // but for now it's too annoying...
    pub(crate) traits_from_clauses: Vec<(Ty, TraitId)>,
    pub env: chalk_ir::Environment<Interner>,
}

impl TraitEnvironment {
    pub fn empty(krate: CrateId) -> Self {
        TraitEnvironment {
            krate,
            traits_from_clauses: Vec::new(),
            env: chalk_ir::Environment::new(Interner),
        }
    }

    pub fn traits_in_scope_from_clauses<'a>(
        &'a self,
        ty: Ty,
    ) -> impl Iterator<Item = TraitId> + 'a {
        self.traits_from_clauses
            .iter()
            .filter_map(move |(self_ty, trait_id)| (*self_ty == ty).then(|| *trait_id))
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Goal>>,
) -> Option<Solution> {
    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.goal.data(Interner) {
        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::Implemented(it))) => {
            db.trait_data(it.hir_trait_id()).name.to_string()
        }
        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(_))) => "alias_eq".to_string(),
        _ => "??".to_string(),
    });
    tracing::info!("trait_solve_query({:?})", goal.value.goal);

    if let GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(AliasEq {
        alias: AliasTy::Projection(projection_ty),
        ..
    }))) = &goal.value.goal.data(Interner)
    {
        if let TyKind::BoundVar(_) = projection_ty.self_type_parameter(Interner).kind(Interner) {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical: goal, universes: 1 };
    solve(db, krate, &u_canonical)
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    tracing::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let should_continue = || {
        db.unwind_if_cancelled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            tracing::debug!("fuel exhausted");
        }
        remaining > 0
    };

    let mut solve = || {
        let _ctx = if is_chalk_debug() || is_chalk_print() {
            Some(panic_context::enter(format!("solving {:?}", goal)))
        } else {
            None
        };
        let solution = if is_chalk_print() {
            let logging_db =
                LoggingRustIrDatabaseLoggingOnDrop(LoggingRustIrDatabase::new(context));
            solver.solve_limited(&logging_db.0, goal, &should_continue)
        } else {
            solver.solve_limited(&context, goal, &should_continue)
        };

        tracing::debug!("solve({:?}) => {:?}", goal, solution);

        solution
    };

    // don't set the TLS for Chalk unless Chalk debugging is active, to make
    // extra sure we only use it for debugging
    if is_chalk_debug() {
        crate::tls::set_current_program(db, solve)
    } else {
        solve()
    }
}

struct LoggingRustIrDatabaseLoggingOnDrop<'a>(LoggingRustIrDatabase<Interner, ChalkContext<'a>>);

impl<'a> Drop for LoggingRustIrDatabaseLoggingOnDrop<'a> {
    fn drop(&mut self) {
        eprintln!("chalk program:\n{}", self.0);
    }
}

fn is_chalk_debug() -> bool {
    std::env::var("CHALK_DEBUG").is_ok()
}

fn is_chalk_print() -> bool {
    std::env::var("CHALK_PRINT").is_ok()
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    const fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }

    pub fn get_id(&self, db: &dyn HirDatabase, krate: CrateId) -> Option<TraitId> {
        let target = db.lang_item(krate, SmolStr::new_inline(self.lang_item_name()))?;
        match target {
            LangItemTarget::TraitId(t) => Some(t),
            _ => None,
        }
    }
}
<|endoftext|>


<commit_before>use std::num::Real;
use std::hashmap::HashMap;
use std::str::StrSlice;

fn shannon_entropy(s: &str) -> f32 {
    let mut map = HashMap::<char, uint>::new();

    for c in s.chars() {
        map.insert_or_update_with(c, 1, |_,v| *v += 1);
    }

    map.iter().fold(0f32, |mut acc, (_, nb)| {
        let p = (*nb as f32)/(s.len() as f32);
        acc -= p * p.log2();
        acc
    })
}

fn main() {
    println!("{:f}", shannon_entropy("1223334444"));
}<commit_msg>Made the shannon_entropy() method public for reuse, and upgraded precision to f64<commit_after>use std::num::Real;
use std::hashmap::HashMap;
use std::str::StrSlice;

pub fn shannon_entropy(s: &str) -> f64 {
    let mut map = HashMap::<char, uint>::new();

    for c in s.chars() {
        map.insert_or_update_with(c, 1, |_,v| *v += 1);
    }

    map.iter().fold(0f64, |mut acc, (_, nb)| {
        let p = (*nb as f64)/(s.len() as f64);
        acc -= p * p.log2();
        acc
    })
}

fn main() {
    println!("{:f}", shannon_entropy("1223334444"));
}<|endoftext|>


<commit_before>use std::io;
use std::net;
use std::time::Duration;
use std::os::unix::io::AsRawFd;

extern crate libc;
extern crate time;
extern crate mio;
extern crate crslmnl as mnl;

use mio::{ Ready, Poll, PollOpt, Token };

use mnl::linux::netlink as netlink;
use mnl::linux::netfilter::nfnetlink as nfnl;
use mnl::linux::netfilter::nfnetlink_conntrack as nfct;
use mnl::linux::netfilter::nf_conntrack_common as nfct_common;
use mnl::linux::netfilter::nf_conntrack_tcp as nfct_tcp;

mod timerfd;

fn put_msg(nlh: &mut mnl::Nlmsg, i: u16, seq: u32) {
    nlh.put_header().unwrap();
    *nlh.nlmsg_type = (nfnl::NFNL_SUBSYS_CTNETLINK << 8) | nfct::IPCTNL_MSG_CT_NEW;
    *nlh.nlmsg_flags = netlink::NLM_F_REQUEST | netlink::NLM_F_CREATE
        | netlink::NLM_F_EXCL | netlink::NLM_F_ACK;
    *nlh.nlmsg_seq = seq;

    let nfh = nlh.put_sized_header::<nfnl::Nfgenmsg>().unwrap();
    nfh.nfgen_family = libc::AF_INET as u8;
    nfh.version = nfnl::NFNETLINK_V0;
    nfh.res_id = 0;

    let mut nest1 = nlh.nest_start(nfct::CTA_TUPLE_ORIG).unwrap();
    let mut nest2 = nlh.nest_start(nfct::CTA_TUPLE_IP).unwrap();
    nlh.put_u32(nfct::CTA_IP_V4_SRC, u32::from(net::Ipv4Addr::new(1, 1, 1, 1))).unwrap();
    nlh.put_u32(nfct::CTA_IP_V4_DST, u32::from(net::Ipv4Addr::new(2, 2, 2, 2))).unwrap();
    nlh.nest_end(nest2);

    nest2 = nlh.nest_start(nfct::CTA_TUPLE_PROTO).unwrap();
    nlh.put_u8(nfct::CTA_PROTO_NUM, libc::IPPROTO_TCP as u8).unwrap();
    nlh.put_u16(nfct::CTA_PROTO_SRC_PORT, u16::to_be(i)).unwrap();
    nlh.put_u16(nfct::CTA_PROTO_DST_PORT, u16::to_be(1025)).unwrap();
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nest1 = nlh.nest_start(nfct::CTA_TUPLE_REPLY).unwrap();
    nest2 = nlh.nest_start(nfct::CTA_TUPLE_IP).unwrap();
    nlh.put_u32(nfct::CTA_IP_V4_SRC, u32::from(net::Ipv4Addr::new(2, 2, 2, 2))).unwrap();
    nlh.put_u32(nfct::CTA_IP_V4_DST, u32::from(net::Ipv4Addr::new(1, 1, 1, 1))).unwrap();
    nlh.nest_end(nest2);

    nest2 = nlh.nest_start(nfct::CTA_TUPLE_PROTO).unwrap();
    nlh.put_u8(nfct::CTA_PROTO_NUM, libc::IPPROTO_TCP as u8).unwrap();
    nlh.put_u16(nfct::CTA_PROTO_SRC_PORT, u16::to_be(1025)).unwrap();
    nlh.put_u16(nfct::CTA_PROTO_DST_PORT, u16::to_be(i)).unwrap();
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nest1 = nlh.nest_start(nfct::CTA_PROTOINFO).unwrap();
    nest2 = nlh.nest_start(nfct::CTA_PROTOINFO_TCP).unwrap();
    nlh.put_u8(nfct::CTA_PROTOINFO_TCP_STATE, nfct_tcp::TCP_CONNTRACK_SYN_SENT).unwrap();
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nlh.put_u32(nfct::CTA_STATUS, u32::to_be(nfct_common::IPS_CONFIRMED)).unwrap();
    nlh.put_u32(nfct::CTA_TIMEOUT, u32::to_be(1000)).unwrap();
}

fn error_cb(nlh: mnl::Nlmsg, _: &mut u8) -> mnl::CbRet {
    let err = nlh.payload::<netlink::Nlmsgerr>();
    if err.error != 0 {
        println!("message with seq {} has failed: {}",
                 nlh.nlmsg_seq, io::Error::from_raw_os_error(-err.error));
    }
    mnl::CbRet::OK
}

fn send_batch(nl: &mut mnl::Socket, b: &mut mnl::NlmsgBatch, portid: u32) {
    nl.send_batch(b)
        .unwrap_or_else(|errno| panic!("mnl_socket_sendto: {}", errno));

    let poll = Poll::new().unwrap();
    let timer = timerfd::Timerfd::create(libc::CLOCK_MONOTONIC, 0).unwrap();
    poll.register(&timer, Token(0),
                  Ready::readable(), PollOpt::edge()).unwrap();

    let rawfd = nl.as_raw_fd();
    let listener = mio::unix::EventedFd(&rawfd);
    poll.register(&listener, Token(rawfd as usize),
                  Ready::readable(), PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(256);
    let mut rcv_buf = vec![0u8; mnl::SOCKET_BUFFER_SIZE()];
    let ctl_cbs = [None,
                   None,		// NLMSG_NOOP
                   Some(error_cb as mnl::Cb<u8>),];	// NLMSG_ERROR
    loop {
        timer.settime(
            0,
            &timerfd::Itimerspec {
                it_interval: Duration::new(0, 0),
                it_value: Duration::new(0, 1),
            }).unwrap();
        poll.poll(&mut events, None).unwrap();

        // handle only the first event - for event in events.iter() {
        match events.get(0) {
            Some(event) => {
                if event.token() == Token(0) {
                    return;
                }
            },
            None => continue, // this happened.
        }

        let nrecv = nl.recvfrom(&mut rcv_buf)
            .unwrap_or_else(|errno| panic!("mnl_socket_recvfrom: {}", errno));
        let rc = mnl::cb_run2(&rcv_buf[0..nrecv], 0, portid,
                              None, &mut 0, &ctl_cbs[..])
            .unwrap_or_else(|errno| panic!("mnl_cb_run2: {}", errno));
        if rc == mnl::CbRet::STOP {
            return;
        }
    }
}

fn main() {
    let nl = mnl::Socket::open(netlink::Family::NETFILTER)
        .unwrap_or_else(|errno| panic!("mnl_socket_open: {}", errno));
    nl.bind(0, mnl::SOCKET_AUTOPID)
        .unwrap_or_else(|errno| panic!("mnl_socket_bind: {}", errno));
    let portid = nl.portid();

    // The buffer that we use to batch messages is MNL_SOCKET_BUFFER_SIZE
    // multiplied by 2 bytes long, but we limit the batch to half of it
    // since the last message that does not fit the batch goes over the
    // upper boundary, if you break this rule, expect memory corruptions.
    let mut snd_buf = vec![0u8; mnl::SOCKET_BUFFER_SIZE() * 2];
    let b = mnl::NlmsgBatch::start(&mut snd_buf, mnl::SOCKET_BUFFER_SIZE())
        .unwrap_or_else(|errno| panic!("mnl_nlmsg_batch_start: {}", errno));

    let seq = time::now().to_timespec().sec as u32;
    for i in 1024u16..65535 {
        put_msg(&mut b.current_nlmsg().unwrap(), i, seq + i as u32 - 1024);
	// is there room for more messages in this batch?
	// if so, continue.
        if b.proceed_next() {
            continue;
        }
        send_batch(nl, b, portid);
	// this moves the last message that did not fit into the
	// batch to the head of it.
        b.reset();
    }

    // check if there is any message in the batch not sent yet.
    if !b.is_empty() {
        send_batch(nl, b, portid);
    }

    let _ = nl.close();
}
<commit_msg>examples: use _raw functions in nfct-create-batch.rs<commit_after>use std::io;
use std::net;
use std::time::Duration;
use std::os::unix::io::AsRawFd;

extern crate libc;
extern crate time;
extern crate mio;
extern crate crslmnl as mnl;

use mio::{ Ready, Poll, PollOpt, Token };

use mnl::linux::netlink as netlink;
use mnl::linux::netfilter::nfnetlink as nfnl;
use mnl::linux::netfilter::nfnetlink_conntrack as nfct;
use mnl::linux::netfilter::nf_conntrack_common as nfct_common;
use mnl::linux::netfilter::nf_conntrack_tcp as nfct_tcp;

mod timerfd;

fn put_msg(nlh: &mut mnl::Nlmsg, i: u16, seq: u32) {
    nlh.put_header_raw();
    *nlh.nlmsg_type = (nfnl::NFNL_SUBSYS_CTNETLINK << 8) | nfct::IPCTNL_MSG_CT_NEW;
    *nlh.nlmsg_flags = netlink::NLM_F_REQUEST | netlink::NLM_F_CREATE
        | netlink::NLM_F_EXCL | netlink::NLM_F_ACK;
    *nlh.nlmsg_seq = seq;

    let nfh = nlh.put_sized_header_raw::<nfnl::Nfgenmsg>();
    nfh.nfgen_family = libc::AF_INET as u8;
    nfh.version = nfnl::NFNETLINK_V0;
    nfh.res_id = 0;

    let mut nest1 = nlh.nest_start_raw(nfct::CTA_TUPLE_ORIG);
    let mut nest2 = nlh.nest_start_raw(nfct::CTA_TUPLE_IP);
    nlh.put_u32_raw(nfct::CTA_IP_V4_SRC, u32::from(net::Ipv4Addr::new(1, 1, 1, 1)));
    nlh.put_u32_raw(nfct::CTA_IP_V4_DST, u32::from(net::Ipv4Addr::new(2, 2, 2, 2)));
    nlh.nest_end(nest2);

    nest2 = nlh.nest_start_raw(nfct::CTA_TUPLE_PROTO);
    nlh.put_u8_raw(nfct::CTA_PROTO_NUM, libc::IPPROTO_TCP as u8);
    nlh.put_u16_raw(nfct::CTA_PROTO_SRC_PORT, u16::to_be(i));
    nlh.put_u16_raw(nfct::CTA_PROTO_DST_PORT, u16::to_be(1025));
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nest1 = nlh.nest_start_raw(nfct::CTA_TUPLE_REPLY);
    nest2 = nlh.nest_start_raw(nfct::CTA_TUPLE_IP);
    nlh.put_u32_raw(nfct::CTA_IP_V4_SRC, u32::from(net::Ipv4Addr::new(2, 2, 2, 2)));
    nlh.put_u32_raw(nfct::CTA_IP_V4_DST, u32::from(net::Ipv4Addr::new(1, 1, 1, 1)));
    nlh.nest_end(nest2);

    nest2 = nlh.nest_start_raw(nfct::CTA_TUPLE_PROTO);
    nlh.put_u8_raw(nfct::CTA_PROTO_NUM, libc::IPPROTO_TCP as u8);
    nlh.put_u16_raw(nfct::CTA_PROTO_SRC_PORT, u16::to_be(1025));
    nlh.put_u16_raw(nfct::CTA_PROTO_DST_PORT, u16::to_be(i));
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nest1 = nlh.nest_start_raw(nfct::CTA_PROTOINFO);
    nest2 = nlh.nest_start_raw(nfct::CTA_PROTOINFO_TCP);
    nlh.put_u8_raw(nfct::CTA_PROTOINFO_TCP_STATE, nfct_tcp::TCP_CONNTRACK_SYN_SENT);
    nlh.nest_end(nest2);
    nlh.nest_end(nest1);

    nlh.put_u32_raw(nfct::CTA_STATUS, u32::to_be(nfct_common::IPS_CONFIRMED));
    nlh.put_u32_raw(nfct::CTA_TIMEOUT, u32::to_be(1000));
}

fn error_cb(nlh: mnl::Nlmsg, _: &mut u8) -> mnl::CbRet {
    let err = nlh.payload::<netlink::Nlmsgerr>();
    if err.error != 0 {
        println!("message with seq {} has failed: {}",
                 nlh.nlmsg_seq, io::Error::from_raw_os_error(-err.error));
    }
    mnl::CbRet::OK
}

fn send_batch(nl: &mut mnl::Socket, b: &mut mnl::NlmsgBatch, portid: u32) {
    nl.send_batch(b)
        .unwrap_or_else(|errno| panic!("mnl_socket_sendto: {}", errno));

    let poll = Poll::new().unwrap();
    let timer = timerfd::Timerfd::create(libc::CLOCK_MONOTONIC, 0).unwrap();
    poll.register(&timer, Token(0),
                  Ready::readable(), PollOpt::edge()).unwrap();

    let rawfd = nl.as_raw_fd();
    let listener = mio::unix::EventedFd(&rawfd);
    poll.register(&listener, Token(rawfd as usize),
                  Ready::readable(), PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(256);
    let mut rcv_buf = vec![0u8; mnl::SOCKET_BUFFER_SIZE()];
    let ctl_cbs = [None,
                   None,		// NLMSG_NOOP
                   Some(error_cb as mnl::Cb<u8>),];	// NLMSG_ERROR
    loop {
        timer.settime(
            0,
            &timerfd::Itimerspec {
                it_interval: Duration::new(0, 0),
                it_value: Duration::new(0, 1),
            }).unwrap();
        poll.poll(&mut events, None).unwrap();

        // handle only the first event - for event in events.iter() {
        match events.get(0) {
            Some(event) => {
                if event.token() == Token(0) {
                    return;
                }
            },
            None => continue, // this happened.
        }

        let nrecv = nl.recvfrom(&mut rcv_buf)
            .unwrap_or_else(|errno| panic!("mnl_socket_recvfrom: {}", errno));
        let rc = mnl::cb_run2(&rcv_buf[0..nrecv], 0, portid,
                              None, &mut 0, &ctl_cbs[..])
            .unwrap_or_else(|errno| panic!("mnl_cb_run2: {}", errno));
        if rc == mnl::CbRet::STOP {
            return;
        }
    }
}

fn main() {
    let nl = mnl::Socket::open(netlink::Family::NETFILTER)
        .unwrap_or_else(|errno| panic!("mnl_socket_open: {}", errno));
    nl.bind(0, mnl::SOCKET_AUTOPID)
        .unwrap_or_else(|errno| panic!("mnl_socket_bind: {}", errno));
    let portid = nl.portid();

    // The buffer that we use to batch messages is MNL_SOCKET_BUFFER_SIZE
    // multiplied by 2 bytes long, but we limit the batch to half of it
    // since the last message that does not fit the batch goes over the
    // upper boundary, if you break this rule, expect memory corruptions.
    let mut snd_buf = vec![0u8; mnl::SOCKET_BUFFER_SIZE() * 2];
    let b = mnl::NlmsgBatch::start(&mut snd_buf, mnl::SOCKET_BUFFER_SIZE())
        .unwrap_or_else(|errno| panic!("mnl_nlmsg_batch_start: {}", errno));

    let seq = time::now().to_timespec().sec as u32;
    for i in 1024u16..65535 {
        put_msg(&mut b.current_nlmsg().unwrap(), i, seq + i as u32 - 1024);
	// is there room for more messages in this batch?
	// if so, continue.
        if b.proceed_next() {
            continue;
        }
        send_batch(nl, b, portid);
	// this moves the last message that did not fit into the
	// batch to the head of it.
        b.reset();
    }

    // check if there is any message in the batch not sent yet.
    if !b.is_empty() {
        send_batch(nl, b, portid);
    }

    let _ = nl.close();
}
<|endoftext|>


<commit_before>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[allow(uppercase_variables)];
#[feature(macro_rules)];

extern crate getopts;
extern crate sync;
extern crate time;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use time::Timespec;
use sync::{Arc, Future};

#[path = "../common/util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    } else if matches.opt_present("version") {
        println!("du version: {}", VERSION);
        return
    }

    let options = Options{
        all: matches.opt_present("all"),
        max_depth: match (matches.opt_present("summarize"), matches.opt_str("max-depth")) {
            (true, Some(s)) => match from_str::<uint>(s) {
                Some(_) => {
                    println!("du: warning: summarizing conflicts with --max-depth={:s}", s);
                    return
                },
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (true, None) => Some(0),
            (false, Some(s)) => match from_str::<uint>(s) {
                Some(u) => Some(u),
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (false, None) => None
        },
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = matches.free.clone();
    let strs = match strs.is_empty() {
        true => ~[~"./"],
        false => strs
    };

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<commit_msg>du: be dryer when extracting the --max-depth arg.<commit_after>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[allow(uppercase_variables)];
#[feature(macro_rules)];

extern crate getopts;
extern crate sync;
extern crate time;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use time::Timespec;
use sync::{Arc, Future};

#[path = "../common/util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    } else if matches.opt_present("version") {
        println!("du version: {}", VERSION);
        return
    }

    let summarize = matches.opt_present("summarize");

    let max_depth_str = matches.opt_str("max-depth");
    let max_depth = max_depth_str.as_ref().and_then(|s| from_str::<uint>(*s));
    match (max_depth_str, max_depth) {
        (Some(ref s), _) if summarize => {
            println!("{}: warning: summarizing conflicts with --max-depth={:s}", program, *s);
            return
        }
        (Some(ref s), None) => {
            println!("{}: invalid maximum depth '{:s}'", program, *s);
            return
        }
        (Some(_), Some(_)) | (None, _) => { /* valid */ }
    }

    let options = Options {
        all: matches.opt_present("all"),
        max_depth: max_depth,
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = if matches.free.is_empty() {~[~"./"]} else {matches.free.clone()};

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<|endoftext|>


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

use crate::U64Le;
use bytemuck::{Pod, Zeroable};
use core::fmt::{self, Display, Formatter};
use core::num::{NonZeroU32, TryFromIntError};
use core::ops::RangeInclusive;

/// Logical block address.
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(transparent)]
pub struct Lba(pub u64);

impl Lba {
    /// Convert to a plain [`u64`].
    #[must_use]
    pub fn to_u64(self) -> u64 {
        self.0
    }
}

impl PartialEq<u64> for Lba {
    fn eq(&self, other: &u64) -> bool {
        self.0 == *other
    }
}

impl Display for Lba {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl TryFrom<Lba> for usize {
    type Error = TryFromIntError;

    fn try_from(lba: Lba) -> Result<Self, Self::Error> {
        lba.0.try_into()
    }
}

impl From<LbaLe> for Lba {
    fn from(lba: LbaLe) -> Self {
        Self(lba.to_u64())
    }
}

/// Logical block address stored as a [`U64Le`].
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(transparent)]
pub struct LbaLe(pub U64Le);

impl LbaLe {
    /// Create a logical block address from a [`u64`].
    #[must_use]
    pub const fn from_u64(v: u64) -> Self {
        Self(U64Le::from_u64(v))
    }

    /// Get the logical block address as a [`u64`].
    #[must_use]
    pub const fn to_u64(self) -> u64 {
        self.0.to_u64()
    }
}

impl Display for LbaLe {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.to_u64().fmt(f)
    }
}

impl From<Lba> for LbaLe {
    fn from(lba: Lba) -> Self {
        Self::from_u64(lba.0)
    }
}

/// Inclusive range of logical block addresses.
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(C)]
pub struct LbaRangeInclusive {
    start: Lba,
    end: Lba,
}

impl LbaRangeInclusive {
    /// Create an LBA range. The end address must be greater than or
    /// equal to the start address.
    #[must_use]
    pub const fn new(start: Lba, end: Lba) -> Option<LbaRangeInclusive> {
        if end.0 >= start.0 {
            Some(LbaRangeInclusive { start, end })
        } else {
            None
        }
    }

    /// Starting LBA (inclusive).
    #[must_use]
    pub const fn start(self) -> Lba {
        self.start
    }

    /// Ending LBA (inclusive).
    #[must_use]
    pub const fn end(self) -> Lba {
        self.end
    }

    /// Create an LBA range from the corresponding byte range for the
    /// given block size.
    ///
    /// The byte range must correspond precisely to block bounds, with
    /// the start byte at the beginning of a block and the end byte at
    /// the end of a block
    ///
    /// # Examples
    ///
    /// ```
    /// use gpt_disk_types::{BlockSize, LbaRangeInclusive};
    ///
    /// let bs = BlockSize::BS_512;
    /// let r = LbaRangeInclusive::from_byte_range(512..=1535, bs).unwrap();
    /// assert_eq!(r.start().0, 1);
    /// assert_eq!(r.end().0, 2);
    #[must_use]
    pub fn from_byte_range(
        byte_range: RangeInclusive<u64>,
        block_size: BlockSize,
    ) -> Option<Self> {
        let start_byte = byte_range.start();
        let end_byte_plus_1 = byte_range.end().checked_add(1)?;
        let block_size = block_size.to_u64();

        if (start_byte % block_size) != 0 {
            return None;
        }
        if (end_byte_plus_1 % block_size) != 0 {
            return None;
        }

        let end_lba = (end_byte_plus_1 / block_size).checked_sub(1)?;

        LbaRangeInclusive::new(Lba(start_byte / block_size), Lba(end_lba))
    }

    /// Convert the LBA range to the corresponding byte range for the
    /// given block size.
    ///
    /// # Examples
    ///
    /// ```
    /// use gpt_disk_types::{BlockSize, Lba, LbaRangeInclusive};
    ///
    /// let r = LbaRangeInclusive::new(Lba(1), Lba(2)).unwrap();
    /// let bs = BlockSize::BS_512;
    /// assert_eq!(r.to_byte_range(bs).unwrap(), 512..=1535);
    /// ```
    #[must_use]
    pub fn to_byte_range(
        self,
        block_size: BlockSize,
    ) -> Option<RangeInclusive<u64>> {
        let block_size = block_size.to_u64();
        let start_byte = self.start.0.checked_mul(block_size)?;
        let end_byte = self
            .end
            .0
            .checked_mul(block_size)?
            .checked_add(block_size - 1)?;
        Some(start_byte..=end_byte)
    }
}

impl Display for LbaRangeInclusive {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}..={}", self.start, self.end)
    }
}

/// Size of a block in bytes.
///
/// This type enforces some restrictions on the block size: it must be
/// at least 512 bytes and fit within a [`u32`].
///
/// # Minimum size
///
/// The [`MasterBootRecord`] size is 512 bytes and must fit within a
/// block, so the block size must be at least that large.
///
/// [`MasterBootRecord`]: crate::MasterBootRecord
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
#[repr(transparent)]
pub struct BlockSize(NonZeroU32);

impl BlockSize {
    /// 512-byte block size.
    pub const BS_512: Self = Self(if let Some(nz) = NonZeroU32::new(512) {
        nz
    } else {
        unreachable!()
    });

    /// 4096-byte block size.
    pub const BS_4096: Self = Self(if let Some(nz) = NonZeroU32::new(4096) {
        nz
    } else {
        unreachable!()
    });

    /// Create a `BlockSize`.
    #[must_use]
    pub const fn new(num_bytes: u32) -> Option<Self> {
        if let Some(nz) = NonZeroU32::new(num_bytes) {
            if num_bytes >= 512 {
                Some(Self(nz))
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Create a `BlockSize`.
    #[must_use]
    pub fn from_usize(num_bytes: usize) -> Option<Self> {
        Self::new(u32::try_from(num_bytes).ok()?)
    }

    /// Get the size in bytes as a [`u32`].
    #[must_use]
    pub const fn to_u32(self) -> u32 {
        self.0.get()
    }

    /// Get the size in bytes as a [`u64`].
    #[allow(clippy::as_conversions)]
    #[must_use]
    pub const fn to_u64(self) -> u64 {
        self.0.get() as u64
    }

    /// Get the size in bytes as a [`usize`].
    #[must_use]
    pub fn to_usize(self) -> Option<usize> {
        self.0.get().try_into().ok()
    }
}

impl Default for BlockSize {
    fn default() -> Self {
        BlockSize::BS_512
    }
}

impl Display for BlockSize {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}
<commit_msg>Add LbaRangeInclusive::num_bytes<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::U64Le;
use bytemuck::{Pod, Zeroable};
use core::fmt::{self, Display, Formatter};
use core::num::{NonZeroU32, TryFromIntError};
use core::ops::RangeInclusive;

/// Logical block address.
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(transparent)]
pub struct Lba(pub u64);

impl Lba {
    /// Convert to a plain [`u64`].
    #[must_use]
    pub fn to_u64(self) -> u64 {
        self.0
    }
}

impl PartialEq<u64> for Lba {
    fn eq(&self, other: &u64) -> bool {
        self.0 == *other
    }
}

impl Display for Lba {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl TryFrom<Lba> for usize {
    type Error = TryFromIntError;

    fn try_from(lba: Lba) -> Result<Self, Self::Error> {
        lba.0.try_into()
    }
}

impl From<LbaLe> for Lba {
    fn from(lba: LbaLe) -> Self {
        Self(lba.to_u64())
    }
}

/// Logical block address stored as a [`U64Le`].
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(transparent)]
pub struct LbaLe(pub U64Le);

impl LbaLe {
    /// Create a logical block address from a [`u64`].
    #[must_use]
    pub const fn from_u64(v: u64) -> Self {
        Self(U64Le::from_u64(v))
    }

    /// Get the logical block address as a [`u64`].
    #[must_use]
    pub const fn to_u64(self) -> u64 {
        self.0.to_u64()
    }
}

impl Display for LbaLe {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.to_u64().fmt(f)
    }
}

impl From<Lba> for LbaLe {
    fn from(lba: Lba) -> Self {
        Self::from_u64(lba.0)
    }
}

/// Inclusive range of logical block addresses.
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    PartialEq,
    Hash,
    Ord,
    PartialOrd,
    Pod,
    Zeroable,
)]
#[repr(C)]
pub struct LbaRangeInclusive {
    start: Lba,
    end: Lba,
}

impl LbaRangeInclusive {
    /// Create an LBA range. The end address must be greater than or
    /// equal to the start address.
    #[must_use]
    pub const fn new(start: Lba, end: Lba) -> Option<LbaRangeInclusive> {
        if end.0 >= start.0 {
            Some(LbaRangeInclusive { start, end })
        } else {
            None
        }
    }

    /// Starting LBA (inclusive).
    #[must_use]
    pub const fn start(self) -> Lba {
        self.start
    }

    /// Ending LBA (inclusive).
    #[must_use]
    pub const fn end(self) -> Lba {
        self.end
    }

    /// Create an LBA range from the corresponding byte range for the
    /// given block size.
    ///
    /// The byte range must correspond precisely to block bounds, with
    /// the start byte at the beginning of a block and the end byte at
    /// the end of a block
    ///
    /// # Examples
    ///
    /// ```
    /// use gpt_disk_types::{BlockSize, LbaRangeInclusive};
    ///
    /// let bs = BlockSize::BS_512;
    /// let r = LbaRangeInclusive::from_byte_range(512..=1535, bs).unwrap();
    /// assert_eq!(r.start().0, 1);
    /// assert_eq!(r.end().0, 2);
    #[must_use]
    pub fn from_byte_range(
        byte_range: RangeInclusive<u64>,
        block_size: BlockSize,
    ) -> Option<Self> {
        let start_byte = byte_range.start();
        let end_byte_plus_1 = byte_range.end().checked_add(1)?;
        let block_size = block_size.to_u64();

        if (start_byte % block_size) != 0 {
            return None;
        }
        if (end_byte_plus_1 % block_size) != 0 {
            return None;
        }

        let end_lba = (end_byte_plus_1 / block_size).checked_sub(1)?;

        LbaRangeInclusive::new(Lba(start_byte / block_size), Lba(end_lba))
    }

    /// Convert the LBA range to the corresponding byte range for the
    /// given block size.
    ///
    /// # Examples
    ///
    /// ```
    /// use gpt_disk_types::{BlockSize, Lba, LbaRangeInclusive};
    ///
    /// let r = LbaRangeInclusive::new(Lba(1), Lba(2)).unwrap();
    /// let bs = BlockSize::BS_512;
    /// assert_eq!(r.to_byte_range(bs).unwrap(), 512..=1535);
    /// ```
    #[must_use]
    pub fn to_byte_range(
        self,
        block_size: BlockSize,
    ) -> Option<RangeInclusive<u64>> {
        let block_size = block_size.to_u64();
        let start_byte = self.start.0.checked_mul(block_size)?;
        let end_byte = self
            .end
            .0
            .checked_mul(block_size)?
            .checked_add(block_size - 1)?;
        Some(start_byte..=end_byte)
    }

    /// Get the number of bytes in the LBA range for the given block
    /// size.
    ///
    /// # Examples
    ///
    /// ```
    /// use gpt_disk_types::{BlockSize, Lba, LbaRangeInclusive};
    ///
    /// let r = LbaRangeInclusive::new(Lba(1), Lba(2)).unwrap();
    /// let bs = BlockSize::BS_512;
    /// assert_eq!(r.num_bytes(bs).unwrap(), 1024);
    /// ```
    #[must_use]
    pub fn num_bytes(self, block_size: BlockSize) -> Option<u64> {
        let r = self.to_byte_range(block_size)?;
        r.end().checked_sub(*r.start())?.checked_add(1)
    }
}

impl Display for LbaRangeInclusive {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}..={}", self.start, self.end)
    }
}

/// Size of a block in bytes.
///
/// This type enforces some restrictions on the block size: it must be
/// at least 512 bytes and fit within a [`u32`].
///
/// # Minimum size
///
/// The [`MasterBootRecord`] size is 512 bytes and must fit within a
/// block, so the block size must be at least that large.
///
/// [`MasterBootRecord`]: crate::MasterBootRecord
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
#[repr(transparent)]
pub struct BlockSize(NonZeroU32);

impl BlockSize {
    /// 512-byte block size.
    pub const BS_512: Self = Self(if let Some(nz) = NonZeroU32::new(512) {
        nz
    } else {
        unreachable!()
    });

    /// 4096-byte block size.
    pub const BS_4096: Self = Self(if let Some(nz) = NonZeroU32::new(4096) {
        nz
    } else {
        unreachable!()
    });

    /// Create a `BlockSize`.
    #[must_use]
    pub const fn new(num_bytes: u32) -> Option<Self> {
        if let Some(nz) = NonZeroU32::new(num_bytes) {
            if num_bytes >= 512 {
                Some(Self(nz))
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Create a `BlockSize`.
    #[must_use]
    pub fn from_usize(num_bytes: usize) -> Option<Self> {
        Self::new(u32::try_from(num_bytes).ok()?)
    }

    /// Get the size in bytes as a [`u32`].
    #[must_use]
    pub const fn to_u32(self) -> u32 {
        self.0.get()
    }

    /// Get the size in bytes as a [`u64`].
    #[allow(clippy::as_conversions)]
    #[must_use]
    pub const fn to_u64(self) -> u64 {
        self.0.get() as u64
    }

    /// Get the size in bytes as a [`usize`].
    #[must_use]
    pub fn to_usize(self) -> Option<usize> {
        self.0.get().try_into().ok()
    }
}

impl Default for BlockSize {
    fn default() -> Self {
        BlockSize::BS_512
    }
}

impl Display for BlockSize {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}
<|endoftext|>


<commit_before>//ZZZ TODO Move imports & code into their own module.
pub mod schema;
pub mod models;

use diesel::prelude::*;
use diesel::sqlite::SqliteConnection; //sqlite specific.
use dotenv::dotenv;
use std::env;

extern crate diesel;

pub fn establish_connection() -> SqliteConnection {
    dotenv().ok();

    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL must be set in the .env file.");

    SqliteConnection::establish(&database_url)
        .expect(&format!("Error connecting to {}", database_url))
}

use self::models::{Image, NewImage};
pub fn create_image(conn: &SqliteConnection, image: Vec<u8>) {
    use self::schema::images;

    let new_image = NewImage {
        image: image,
    };

    diesel::insert(&new_image).into(images::table);

}
<commit_msg>Forgot to execute statement...<commit_after>//ZZZ TODO Move imports & code into their own module.
pub mod schema;
pub mod models;

use diesel::prelude::*;
use diesel::sqlite::SqliteConnection; //sqlite specific.
use dotenv::dotenv;
use std::env;

extern crate diesel;

pub fn establish_connection() -> SqliteConnection {
    dotenv().ok();

    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL must be set in the .env file.");

    SqliteConnection::establish(&database_url)
        .expect(&format!("Error connecting to {}", database_url))
}

use self::models::{Image, NewImage};
pub fn create_image(conn: &SqliteConnection, image: Vec<u8>) {
    use self::schema::images;

    let new_image = NewImage {
        image: image,
    };

    diesel::insert(&new_image).into(images::table).execute(conn);

}
<|endoftext|>


<commit_before>/*
 * OpenGL Shaders
 */
pub const VERTEX_SHADER: &str = r#"
#version 100
uniform mat3 matrix;
uniform bool invert_y;

attribute vec2 vert;
attribute vec2 tex_coords;
attribute vec4 position;

varying vec2 v_tex_coords;

mat2 scale(vec2 scale_vec){
    return mat2(
        scale_vec.x, 0.0,
        0.0, scale_vec.y
    );
}

void main() {
    if (invert_y) {
        v_tex_coords = vec2(tex_coords.x, 1.0 - tex_coords.y);
    } else {
        v_tex_coords = tex_coords;
    }

    vec2 transform_translation = position.xy;
    vec2 transform_scale = position.zw;
    v_tex_coords = (vec3((tex_coords * scale(transform_scale)) + transform_translation, 1.0)).xy;
    gl_Position = vec4(matrix * vec3((vert * scale(transform_scale)) + transform_translation, 1.0), 1.0);
}
"#;

pub const FRAGMENT_COUNT: usize = 3;

pub const FRAGMENT_SHADER_ABGR: &str = r#"
#version 100

precision mediump float;
uniform sampler2D tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = texture2D(tex, v_tex_coords) * alpha;
}
"#;

pub const FRAGMENT_SHADER_XBGR: &str = r#"
#version 100

precision mediump float;
uniform sampler2D tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = vec4(texture2D(tex, v_tex_coords).rgb, 1.0) * alpha;
}
"#;

pub const FRAGMENT_SHADER_EXTERNAL: &str = r#"
#version 100
#extension GL_OES_EGL_image_external : require

precision mediump float;
uniform samplerExternalOES tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = texture2D(tex, v_tex_coords) * alpha;
}
"#;

pub const VERTEX_SHADER_SOLID: &str = r#"
#version 100

uniform mat3 matrix;
attribute vec2 vert;
attribute vec4 position;

mat2 scale(vec2 scale_vec){
    return mat2(
        scale_vec.x, 0.0,
        0.0, scale_vec.y
    );
}

void main() {
    vec2 transform_translation = position.xy;
    vec2 transform_scale = position.zw;
    gl_Position = vec4(matrix * vec3((vert * scale(transform_scale)) + transform_translation, 1.0), 1.0);
}
"#;

pub const FRAGMENT_SHADER_SOLID: &str = r#"
#version 100

precision mediump float;
uniform vec4 color;

void main() {
    gl_FragColor = color;
}
"#;
<commit_msg>gles2: Fix rendering of inverted textures<commit_after>/*
 * OpenGL Shaders
 */
pub const VERTEX_SHADER: &str = r#"
#version 100
uniform mat3 matrix;
uniform bool invert_y;

attribute vec2 vert;
attribute vec2 tex_coords;
attribute vec4 position;

varying vec2 v_tex_coords;

mat2 scale(vec2 scale_vec){
    return mat2(
        scale_vec.x, 0.0,
        0.0, scale_vec.y
    );
}

void main() {
    if (invert_y) {
        v_tex_coords = vec2(tex_coords.x, 1.0 - tex_coords.y);
    } else {
        v_tex_coords = tex_coords;
    }

    vec2 transform_translation = position.xy;
    vec2 transform_scale = position.zw;
    v_tex_coords = (vec3((v_tex_coords * scale(transform_scale)) + transform_translation, 1.0)).xy;
    gl_Position = vec4(matrix * vec3((vert * scale(transform_scale)) + transform_translation, 1.0), 1.0);
}
"#;

pub const FRAGMENT_COUNT: usize = 3;

pub const FRAGMENT_SHADER_ABGR: &str = r#"
#version 100

precision mediump float;
uniform sampler2D tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = texture2D(tex, v_tex_coords) * alpha;
}
"#;

pub const FRAGMENT_SHADER_XBGR: &str = r#"
#version 100

precision mediump float;
uniform sampler2D tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = vec4(texture2D(tex, v_tex_coords).rgb, 1.0) * alpha;
}
"#;

pub const FRAGMENT_SHADER_EXTERNAL: &str = r#"
#version 100
#extension GL_OES_EGL_image_external : require

precision mediump float;
uniform samplerExternalOES tex;
uniform float alpha;
varying vec2 v_tex_coords;

void main() {
    gl_FragColor = texture2D(tex, v_tex_coords) * alpha;
}
"#;

pub const VERTEX_SHADER_SOLID: &str = r#"
#version 100

uniform mat3 matrix;
attribute vec2 vert;
attribute vec4 position;

mat2 scale(vec2 scale_vec){
    return mat2(
        scale_vec.x, 0.0,
        0.0, scale_vec.y
    );
}

void main() {
    vec2 transform_translation = position.xy;
    vec2 transform_scale = position.zw;
    gl_Position = vec4(matrix * vec3((vert * scale(transform_scale)) + transform_translation, 1.0), 1.0);
}
"#;

pub const FRAGMENT_SHADER_SOLID: &str = r#"
#version 100

precision mediump float;
uniform vec4 color;

void main() {
    gl_FragColor = color;
}
"#;
<|endoftext|>


<commit_before>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//! #[macro_use]
//! extern crate enum_map_derive;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

mod enummap_impls;
mod implementations;
mod iter;
mod serde;

pub use iter::{Iter, IterMut};

// this is not stable, do not depend on this
#[doc(hidden)]
pub trait Internal<V>: Sized {
    type Array;
    fn slice(&Self::Array) -> &[V];
    fn slice_mut(&mut Self::Array) -> &mut [V];
    fn from_usize(usize) -> Self;
    fn to_usize(self) -> usize;
    fn from_function<F: FnMut(Self) -> V>(F) -> Self::Array;
}

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// extern crate enum_map;
/// #[macro_use]
/// extern crate enum_map_derive;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
    where K::Array: Default
{
    /// Creates an enum map with default values.
    ///
    /// ```
    /// extern crate enum_map;
    /// #[macro_use]
    /// extern crate enum_map_derive;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        K::slice(&self.array).len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate enum_map;
    /// #[macro_use]
    /// extern crate enum_map_derive;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        K::slice(&self.array).is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
/// #[macro_use]
/// extern crate enum_map_derive;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        ::enum_map::EnumMap::from(|k| match k { $($t)* })
    };
}
<commit_msg>Introduce enum map swap functionality<commit_after>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//! #[macro_use]
//! extern crate enum_map_derive;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

mod enummap_impls;
mod implementations;
mod iter;
mod serde;

pub use iter::{Iter, IterMut};

// this is not stable, do not depend on this
#[doc(hidden)]
pub trait Internal<V>: Sized {
    type Array;
    fn slice(&Self::Array) -> &[V];
    fn slice_mut(&mut Self::Array) -> &mut [V];
    fn from_usize(usize) -> Self;
    fn to_usize(self) -> usize;
    fn from_function<F: FnMut(Self) -> V>(F) -> Self::Array;
}

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// extern crate enum_map;
/// #[macro_use]
/// extern crate enum_map_derive;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
    where K::Array: Default
{
    /// Creates an enum map with default values.
    ///
    /// ```
    /// extern crate enum_map;
    /// #[macro_use]
    /// extern crate enum_map_derive;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        K::slice(&self.array).len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate enum_map;
    /// #[macro_use]
    /// extern crate enum_map_derive;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        K::slice(&self.array).is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }

    /// Swaps two indexes.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { false => 0, true => 1 };
    ///     map.swap(false, true);
    ///     assert_eq!(map[false], 1);
    ///     assert_eq!(map[true], 0);
    /// }
    /// ```
    pub fn swap(&mut self, a: K, b: K) {
        K::slice_mut(&mut self.array).swap(a.to_usize(), b.to_usize())
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
/// #[macro_use]
/// extern crate enum_map_derive;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        ::enum_map::EnumMap::from(|k| match k { $($t)* })
    };
}
<|endoftext|>