<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::rustc::hir::intravisit::FnKind;
use crate::rustc::hir::{def_id, Body, FnDecl};
use crate::rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use crate::rustc::mir::{
    self, traversal,
    visit::{MutatingUseContext, NonUseContext, PlaceContext, Visitor},
    TerminatorKind,
};
use crate::rustc::ty;
use crate::rustc::{declare_tool_lint, lint_array};
use crate::rustc_errors::Applicability;
use crate::syntax::{
    ast::NodeId,
    source_map::{BytePos, Span},
};
use crate::utils::{
    in_macro, is_copy, match_def_path, match_type, paths, snippet_opt, span_lint_node, span_lint_node_and_then,
    walk_ptrs_ty_depth,
};
use if_chain::if_chain;
use std::convert::TryFrom;

macro_rules! unwrap_or_continue {
    ($x:expr) => {
        match $x {
            Some(x) => x,
            None => continue,
        }
    };
}

/// **What it does:** Checks for a redudant `clone()` (and its relatives) which clones an owned
/// value that is going to be dropped without further use.
///
/// **Why is this bad?** It is not always possible for the compiler to eliminate useless
/// allocations and deallocations generated by redundant `clone()`s.
///
/// **Known problems:**
///
/// * Suggestions made by this lint could require NLL to be enabled.
/// * False-positive if there is a borrow preventing the value from moving out.
///
/// ```rust
/// let x = String::new();
///
/// let y = &x;
///
/// foo(x.clone()); // This lint suggests to remove this `clone()`
/// ```
///
/// **Example:**
/// ```rust
/// {
///     let x = Foo::new();
///     call(x.clone());
///     call(x.clone()); // this can just pass `x`
/// }
///
/// ["lorem", "ipsum"].join(" ").to_string()
///
/// Path::new("/a/b").join("c").to_path_buf()
/// ```
declare_clippy_lint! {
    pub REDUNDANT_CLONE,
    nursery,
    "`clone()` of an owned value that is going to be dropped immediately"
}

pub struct RedundantClone;

impl LintPass for RedundantClone {
    fn get_lints(&self) -> LintArray {
        lint_array!(REDUNDANT_CLONE)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for RedundantClone {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        _: NodeId,
    ) {
        let def_id = cx.tcx.hir.body_owner_def_id(body.id());
        let mir = cx.tcx.optimized_mir(def_id);

        for (bb, bbdata) in mir.basic_blocks().iter_enumerated() {
            let terminator = bbdata.terminator();

            if in_macro(terminator.source_info.span) {
                continue;
            }

            // Give up on loops
            if terminator.successors().any(|s| *s == bb) {
                continue;
            }

            let (fn_def_id, arg, arg_ty, _) = unwrap_or_continue!(is_call_with_ref_arg(cx, mir, &terminator.kind));

            let from_borrow = match_def_path(cx.tcx, fn_def_id, &paths::CLONE_TRAIT_METHOD)
                || match_def_path(cx.tcx, fn_def_id, &paths::TO_OWNED_METHOD)
                || (match_def_path(cx.tcx, fn_def_id, &paths::TO_STRING_METHOD)
                    && match_type(cx, arg_ty, &paths::STRING));

            let from_deref = !from_borrow
                && (match_def_path(cx.tcx, fn_def_id, &paths::PATH_TO_PATH_BUF)
                    || match_def_path(cx.tcx, fn_def_id, &paths::OS_STR_TO_OS_STRING));

            if !from_borrow && !from_deref {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; clone(move _2); }` or `{ _2 = _1; to_path_buf(_2); } (from_deref)
            // In case of `from_deref`, `arg` is already a reference since it is `deref`ed in the previous
            // block.
            let cloned = unwrap_or_continue!(find_stmt_assigns_to(arg, from_borrow, bbdata.statements.iter().rev()));

            // _1 in MIR `{ _2 = &_1; _3 = deref(move _2); } -> { _4 = _3; to_path_buf(move _4); }`
            let referent = if from_deref {
                let ps = mir.predecessors_for(bb);
                if ps.len() != 1 {
                    continue;
                }
                let pred_terminator = mir[ps[0]].terminator();

                let pred_arg = if_chain! {
                    if let Some((pred_fn_def_id, pred_arg, pred_arg_ty, Some(res))) =
                        is_call_with_ref_arg(cx, mir, &pred_terminator.kind);
                    if *res == mir::Place::Local(cloned);
                    if match_def_path(cx.tcx, pred_fn_def_id, &paths::DEREF_TRAIT_METHOD);
                    if match_type(cx, pred_arg_ty, &paths::PATH_BUF)
                        || match_type(cx, pred_arg_ty, &paths::OS_STRING);
                    then {
                        pred_arg
                    } else {
                        continue;
                    }
                };

                unwrap_or_continue!(find_stmt_assigns_to(pred_arg, true, mir[ps[0]].statements.iter().rev()))
            } else {
                cloned
            };

            let used_later = traversal::ReversePostorder::new(&mir, bb).skip(1).any(|(tbb, tdata)| {
                // Give up on loops
                if tdata.terminator().successors().any(|s| *s == bb) {
                    return true;
                }

                let mut vis = LocalUseVisitor {
                    local: referent,
                    used_other_than_drop: false,
                };
                vis.visit_basic_block_data(tbb, tdata);
                vis.used_other_than_drop
            });

            if !used_later {
                let span = terminator.source_info.span;
                let node = if let mir::ClearCrossCrate::Set(scope_local_data) = &mir.source_scope_local_data {
                    scope_local_data[terminator.source_info.scope].lint_root
                } else {
                    unreachable!()
                };

                if_chain! {
                    if let Some(snip) = snippet_opt(cx, span);
                    if let Some(dot) = snip.rfind('.');
                    then {
                        let sugg_span = span.with_lo(
                            span.lo() + BytePos(u32::try_from(dot).unwrap())
                        );

                        span_lint_node_and_then(cx, REDUNDANT_CLONE, node, sugg_span, "redundant clone", |db| {
                            db.span_suggestion_with_applicability(
                                sugg_span,
                                "remove this",
                                String::new(),
                                Applicability::MaybeIncorrect,
                            );
                            db.span_note(
                                span.with_hi(span.lo() + BytePos(u32::try_from(dot).unwrap())),
                                "this value is dropped without further use",
                            );
                        });
                    } else {
                        span_lint_node(cx, REDUNDANT_CLONE, node, span, "redundant clone");
                    }
                }
            }
        }
    }
}

/// If `kind` is `y = func(x: &T)` where `T: !Copy`, returns `(DefId of func, x, T, y)`.
fn is_call_with_ref_arg<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &'tcx mir::Mir<'tcx>,
    kind: &'tcx mir::TerminatorKind<'tcx>,
) -> Option<(def_id::DefId, mir::Local, ty::Ty<'tcx>, Option<&'tcx mir::Place<'tcx>>)> {
    if_chain! {
        if let TerminatorKind::Call { func, args, destination, .. } = kind;
        if args.len() == 1;
        if let mir::Operand::Move(mir::Place::Local(local)) = &args[0];
        if let ty::FnDef(def_id, _) = func.ty(&*mir, cx.tcx).sty;
        if let (inner_ty, 1) = walk_ptrs_ty_depth(args[0].ty(&*mir, cx.tcx));
        if !is_copy(cx, inner_ty);
        then {
            Some((def_id, *local, inner_ty, destination.as_ref().map(|(dest, _)| dest)))
        } else {
            None
        }
    }
}

/// Finds the first `to = (&)from`, and returns `Some(from)`.
fn find_stmt_assigns_to<'a, 'tcx: 'a>(
    to: mir::Local,
    by_ref: bool,
    mut stmts: impl Iterator<Item = &'a mir::Statement<'tcx>>,
) -> Option<mir::Local> {
    stmts.find_map(|stmt| {
        if let mir::StatementKind::Assign(mir::Place::Local(local), v) = &stmt.kind {
            if *local == to {
                if by_ref {
                    if let mir::Rvalue::Ref(_, _, mir::Place::Local(r)) = **v {
                        return Some(r);
                    }
                } else if let mir::Rvalue::Use(mir::Operand::Copy(mir::Place::Local(r))) = **v {
                    return Some(r);
                }
            }
        }

        None
    })
}

struct LocalUseVisitor {
    local: mir::Local,
    used_other_than_drop: bool,
}

impl<'tcx> mir::visit::Visitor<'tcx> for LocalUseVisitor {
    fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
        let statements = &data.statements;
        for (statement_index, statement) in statements.iter().enumerate() {
            self.visit_statement(block, statement, mir::Location { block, statement_index });

            // Once flagged, skip remaining statements
            if self.used_other_than_drop {
                return;
            }
        }

        self.visit_terminator(
            block,
            data.terminator(),
            mir::Location {
                block,
                statement_index: statements.len(),
            },
        );
    }

    fn visit_local(&mut self, local: &mir::Local, ctx: PlaceContext<'tcx>, _: mir::Location) {
        match ctx {
            PlaceContext::MutatingUse(MutatingUseContext::Drop) | PlaceContext::NonUse(NonUseContext::StorageDead) => {
                return
            },
            _ => {},
        }

        if *local == self.local {
            self.used_other_than_drop = true;
        }
    }
}
<commit_msg>Fix rustfmt format<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::rustc::hir::intravisit::FnKind;
use crate::rustc::hir::{def_id, Body, FnDecl};
use crate::rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use crate::rustc::mir::{
    self, traversal,
    visit::{MutatingUseContext, NonUseContext, PlaceContext, Visitor},
    TerminatorKind,
};
use crate::rustc::ty;
use crate::rustc::{declare_tool_lint, lint_array};
use crate::rustc_errors::Applicability;
use crate::syntax::{
    ast::NodeId,
    source_map::{BytePos, Span},
};
use crate::utils::{
    in_macro, is_copy, match_def_path, match_type, paths, snippet_opt, span_lint_node, span_lint_node_and_then,
    walk_ptrs_ty_depth,
};
use if_chain::if_chain;
use std::convert::TryFrom;

macro_rules! unwrap_or_continue {
    ($x:expr) => {
        match $x {
            Some(x) => x,
            None => continue,
        }
    };
}

/// **What it does:** Checks for a redudant `clone()` (and its relatives) which clones an owned
/// value that is going to be dropped without further use.
///
/// **Why is this bad?** It is not always possible for the compiler to eliminate useless
/// allocations and deallocations generated by redundant `clone()`s.
///
/// **Known problems:**
///
/// * Suggestions made by this lint could require NLL to be enabled.
/// * False-positive if there is a borrow preventing the value from moving out.
///
/// ```rust
/// let x = String::new();
///
/// let y = &x;
///
/// foo(x.clone()); // This lint suggests to remove this `clone()`
/// ```
///
/// **Example:**
/// ```rust
/// {
///     let x = Foo::new();
///     call(x.clone());
///     call(x.clone()); // this can just pass `x`
/// }
///
/// ["lorem", "ipsum"].join(" ").to_string()
///
/// Path::new("/a/b").join("c").to_path_buf()
/// ```
declare_clippy_lint! {
    pub REDUNDANT_CLONE,
    nursery,
    "`clone()` of an owned value that is going to be dropped immediately"
}

pub struct RedundantClone;

impl LintPass for RedundantClone {
    fn get_lints(&self) -> LintArray {
        lint_array!(REDUNDANT_CLONE)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for RedundantClone {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        _: NodeId,
    ) {
        let def_id = cx.tcx.hir.body_owner_def_id(body.id());
        let mir = cx.tcx.optimized_mir(def_id);

        for (bb, bbdata) in mir.basic_blocks().iter_enumerated() {
            let terminator = bbdata.terminator();

            if in_macro(terminator.source_info.span) {
                continue;
            }

            // Give up on loops
            if terminator.successors().any(|s| *s == bb) {
                continue;
            }

            let (fn_def_id, arg, arg_ty, _) = unwrap_or_continue!(is_call_with_ref_arg(cx, mir, &terminator.kind));

            let from_borrow = match_def_path(cx.tcx, fn_def_id, &paths::CLONE_TRAIT_METHOD)
                || match_def_path(cx.tcx, fn_def_id, &paths::TO_OWNED_METHOD)
                || (match_def_path(cx.tcx, fn_def_id, &paths::TO_STRING_METHOD)
                    && match_type(cx, arg_ty, &paths::STRING));

            let from_deref = !from_borrow
                && (match_def_path(cx.tcx, fn_def_id, &paths::PATH_TO_PATH_BUF)
                    || match_def_path(cx.tcx, fn_def_id, &paths::OS_STR_TO_OS_STRING));

            if !from_borrow && !from_deref {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; clone(move _2); }` or `{ _2 = _1; to_path_buf(_2); } (from_deref)
            // In case of `from_deref`, `arg` is already a reference since it is `deref`ed in the previous
            // block.
            let cloned = unwrap_or_continue!(find_stmt_assigns_to(arg, from_borrow, bbdata.statements.iter().rev()));

            // _1 in MIR `{ _2 = &_1; _3 = deref(move _2); } -> { _4 = _3; to_path_buf(move _4); }`
            let referent = if from_deref {
                let ps = mir.predecessors_for(bb);
                if ps.len() != 1 {
                    continue;
                }
                let pred_terminator = mir[ps[0]].terminator();

                let pred_arg = if_chain! {
                    if let Some((pred_fn_def_id, pred_arg, pred_arg_ty, Some(res))) =
                        is_call_with_ref_arg(cx, mir, &pred_terminator.kind);
                    if *res == mir::Place::Local(cloned);
                    if match_def_path(cx.tcx, pred_fn_def_id, &paths::DEREF_TRAIT_METHOD);
                    if match_type(cx, pred_arg_ty, &paths::PATH_BUF)
                        || match_type(cx, pred_arg_ty, &paths::OS_STRING);
                    then {
                        pred_arg
                    } else {
                        continue;
                    }
                };

                unwrap_or_continue!(find_stmt_assigns_to(pred_arg, true, mir[ps[0]].statements.iter().rev()))
            } else {
                cloned
            };

            let used_later = traversal::ReversePostorder::new(&mir, bb).skip(1).any(|(tbb, tdata)| {
                // Give up on loops
                if tdata.terminator().successors().any(|s| *s == bb) {
                    return true;
                }

                let mut vis = LocalUseVisitor {
                    local: referent,
                    used_other_than_drop: false,
                };
                vis.visit_basic_block_data(tbb, tdata);
                vis.used_other_than_drop
            });

            if !used_later {
                let span = terminator.source_info.span;
                let node = if let mir::ClearCrossCrate::Set(scope_local_data) = &mir.source_scope_local_data {
                    scope_local_data[terminator.source_info.scope].lint_root
                } else {
                    unreachable!()
                };

                if_chain! {
                    if let Some(snip) = snippet_opt(cx, span);
                    if let Some(dot) = snip.rfind('.');
                    then {
                        let sugg_span = span.with_lo(
                            span.lo() + BytePos(u32::try_from(dot).unwrap())
                        );

                        span_lint_node_and_then(cx, REDUNDANT_CLONE, node, sugg_span, "redundant clone", |db| {
                            db.span_suggestion_with_applicability(
                                sugg_span,
                                "remove this",
                                String::new(),
                                Applicability::MaybeIncorrect,
                            );
                            db.span_note(
                                span.with_hi(span.lo() + BytePos(u32::try_from(dot).unwrap())),
                                "this value is dropped without further use",
                            );
                        });
                    } else {
                        span_lint_node(cx, REDUNDANT_CLONE, node, span, "redundant clone");
                    }
                }
            }
        }
    }
}

/// If `kind` is `y = func(x: &T)` where `T: !Copy`, returns `(DefId of func, x, T, y)`.
fn is_call_with_ref_arg<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &'tcx mir::Mir<'tcx>,
    kind: &'tcx mir::TerminatorKind<'tcx>,
) -> Option<(def_id::DefId, mir::Local, ty::Ty<'tcx>, Option<&'tcx mir::Place<'tcx>>)> {
    if_chain! {
        if let TerminatorKind::Call { func, args, destination, .. } = kind;
        if args.len() == 1;
        if let mir::Operand::Move(mir::Place::Local(local)) = &args[0];
        if let ty::FnDef(def_id, _) = func.ty(&*mir, cx.tcx).sty;
        if let (inner_ty, 1) = walk_ptrs_ty_depth(args[0].ty(&*mir, cx.tcx));
        if !is_copy(cx, inner_ty);
        then {
            Some((def_id, *local, inner_ty, destination.as_ref().map(|(dest, _)| dest)))
        } else {
            None
        }
    }
}

/// Finds the first `to = (&)from`, and returns `Some(from)`.
fn find_stmt_assigns_to<'a, 'tcx: 'a>(
    to: mir::Local,
    by_ref: bool,
    mut stmts: impl Iterator<Item = &'a mir::Statement<'tcx>>,
) -> Option<mir::Local> {
    stmts.find_map(|stmt| {
        if let mir::StatementKind::Assign(mir::Place::Local(local), v) = &stmt.kind {
            if *local == to {
                if by_ref {
                    if let mir::Rvalue::Ref(_, _, mir::Place::Local(r)) = **v {
                        return Some(r);
                    }
                } else if let mir::Rvalue::Use(mir::Operand::Copy(mir::Place::Local(r))) = **v {
                    return Some(r);
                }
            }
        }

        None
    })
}

struct LocalUseVisitor {
    local: mir::Local,
    used_other_than_drop: bool,
}

impl<'tcx> mir::visit::Visitor<'tcx> for LocalUseVisitor {
    fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
        let statements = &data.statements;
        for (statement_index, statement) in statements.iter().enumerate() {
            self.visit_statement(block, statement, mir::Location { block, statement_index });

            // Once flagged, skip remaining statements
            if self.used_other_than_drop {
                return;
            }
        }

        self.visit_terminator(
            block,
            data.terminator(),
            mir::Location {
                block,
                statement_index: statements.len(),
            },
        );
    }

    fn visit_local(&mut self, local: &mir::Local, ctx: PlaceContext<'tcx>, _: mir::Location) {
        match ctx {
            PlaceContext::MutatingUse(MutatingUseContext::Drop) | PlaceContext::NonUse(NonUseContext::StorageDead) => {
                return;
            },
            _ => {},
        }

        if *local == self.local {
            self.used_other_than_drop = true;
        }
    }
}
<|endoftext|>


<commit_before>//! Check running processes

use std::collections::HashSet;
use std::str::FromStr;

use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(
    name = "check-procs (part of tabin-plugins)",
    raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
    after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states. Note that since there can be
    multiple states the regex must come before the `state` flag:

        check-procs 'java.*MyMainClass' --crit-over 5 --state zombie waiting

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state=running --state=waiting 'cassandra|postgres'"
)]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<Regex>,
    #[structopt(
        long = "crit-under",
        name = "N",
        help = "Error if there are fewer than <N> procs matching <pattern>"
    )]
    crit_under: Option<usize>,
    #[structopt(
        long = "crit-over",
        name = "M",
        help = "Error if there are more than <M> procs matching <pattern>"
    )]
    crit_over: Option<usize>,

    #[structopt(
        long = "state",
        help = "Filter to only processes in these states. \
                If passed multiple times, processes matching any state are included.\n\
                Choices: running sleeping uninterruptible-sleep waiting stopped zombie"
    )]
    states: Vec<State>,

    #[structopt(
        long = "allow-unparseable-procs",
        help = "In combination with --crit-over M this will not alert if any \
                processes cannot be parsed"
    )]
    allow_unparseable_procs: bool,

    #[structopt(
        long = "kill-matching",
        name = "SIGNAL",
        help = "If *any* processes match, then kill them with the provided signal \
                which can be either an integer or a name like KILL or SIGTERM. \
                This option does not affect the exit status, all matches are always \
                killed, and if --crit-under/over are violated then then this will \
                still exit critical."
    )]
    kill_matching: Option<Signal>,

    #[structopt(
        long = "kill-parents-of-matching",
        name = "PARENT_SIGNAL",
        help = "If *any* processes match, then kill their parents with the provided \
                signal which can be either an integer or a name like KILL or SIGTERM. \
                This has the same exit status behavior as kill-matching."
    )]
    kill_matching_parents: Option<Signal>,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => {
                Ok(Signal(NixSignal::from_c_int(integer).map_err(|_| {
                    format!("Not a valid signal integer: {}", s)
                })?))
            }
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> Args {
    let args = Args::from_args();
    if args.pattern.is_none() && args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    args
}

fn main() {
    let args = parse_args();
    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);
    let re = args.pattern.as_ref().map(|r| r.to_string());

    let matches = filter_procs(args.pattern, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (greater than {})",
                matches.len(),
                re,
                crit_over
            );
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (less than {})",
                matches.len(),
                re,
                crit_under
            );
        }
    }

    if status == Status::Ok {
        match (args.crit_over, args.crit_under) {
            (Some(o), Some(u)) => println!(
                "OKAY: There are {} matching procs (between {} and {})",
                matches.len(),
                o,
                u
            ),
            (Some(o), None) => println!(
                "OKAY: There are {} matching procs (less than {})",
                matches.len(),
                o
            ),
            (None, Some(u)) => println!(
                "OKAY: There are {} matching procs (greater than {})",
                matches.len(),
                u
            ),
            (None, None) => unreachable!(),
        }
    }
    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();

    procs
        .iter()
        .filter(|&(pid, _process)| *pid != me && *pid != parent)
        .filter(|&(_pid, process)| states.is_empty() || states.contains(&process.stat.state))
        .filter(|&(_pid, process)| {
            re.as_ref()
                .map(|re| re.is_match(&process.useful_cmdline()))
                .unwrap_or(true)
        })
        .collect()
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn validate_parse_zombies_and_pattern() {
        let args = Args::from_iter(["c-p", "--state", "zombie", "--", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "--state=zombie", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "so.*proc", "--state", "zombie", "waiting"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Waiting]);
    }

    // Waiting for structopt 0.2.8 to be released with the from_iter_safe method
    // #[test]
    // #[should_panic]
    // fn regexes_must_be_valid() {
    //     Args::from_iter_safe(["c-p", "--crit-over=5", "["].into_iter()).unwrap();
    // }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs.clone());

        let filtered = filter_procs(regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].1, &procs[2]);

        let filtered = filter_procs(regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<commit_msg>[check-procs] add debug/verbose output<commit_after>//! Check running processes

use std::collections::HashSet;
use std::str::FromStr;

use log::LevelFilter::{Debug, Trace, Warn};
use log::{debug, trace};
use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

const LOG_VAR: &str = "TABIN_LOG";

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(
    name = "check-procs (part of tabin-plugins)",
    raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
    after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states. Note that since there can be
    multiple states the regex must come before the `state` flag:

        check-procs 'java.*MyMainClass' --crit-over 5 --state zombie waiting

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state=running --state=waiting 'cassandra|postgres'"
)]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<Regex>,
    #[structopt(
        long = "crit-under",
        name = "N",
        help = "Error if there are fewer than <N> procs matching <pattern>"
    )]
    crit_under: Option<usize>,
    #[structopt(
        long = "crit-over",
        name = "M",
        help = "Error if there are more than <M> procs matching <pattern>"
    )]
    crit_over: Option<usize>,

    #[structopt(
        long = "state",
        help = "Filter to only processes in these states. \
                If passed multiple times, processes matching any state are included.\n\
                Choices: running sleeping uninterruptible-sleep waiting stopped zombie"
    )]
    states: Vec<State>,

    #[structopt(
        long = "allow-unparseable-procs",
        help = "In combination with --crit-over M this will not alert if any \
                processes cannot be parsed"
    )]
    allow_unparseable_procs: bool,

    #[structopt(
        long = "kill-matching",
        name = "SIGNAL",
        help = "If *any* processes match, then kill them with the provided signal \
                which can be either an integer or a name like KILL or SIGTERM. \
                This option does not affect the exit status, all matches are always \
                killed, and if --crit-under/over are violated then then this will \
                still exit critical."
    )]
    kill_matching: Option<Signal>,

    #[structopt(
        long = "kill-parents-of-matching",
        name = "PARENT_SIGNAL",
        help = "If *any* processes match, then kill their parents with the provided \
                signal which can be either an integer or a name like KILL or SIGTERM. \
                This has the same exit status behavior as kill-matching."
    )]
    kill_matching_parents: Option<Signal>,

    /// print debug logs, use multiple times to make it more verbose
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: u8,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => {
                Ok(Signal(NixSignal::from_c_int(integer).map_err(|_| {
                    format!("Not a valid signal integer: {}", s)
                })?))
            }
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> Args {
    let args = Args::from_args();
    if args.pattern.is_none() && args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    args
}

fn main() {
    let args = parse_args();
    env_logger::Builder::from_env(LOG_VAR)
        .filter_level(match args.verbose {
            0 => Warn,
            1 => Debug,
            _ => Trace,
        })
        .init();

    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);
    let re = args.pattern.as_ref().map(|r| r.to_string());

    let matches = filter_procs(args.pattern, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (greater than {})",
                matches.len(),
                re,
                crit_over
            );
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (less than {})",
                matches.len(),
                re,
                crit_under
            );
        }
    }

    if status == Status::Ok {
        match (args.crit_over, args.crit_under) {
            (Some(o), Some(u)) => println!(
                "OKAY: There are {} matching procs (between {} and {})",
                matches.len(),
                o,
                u
            ),
            (Some(o), None) => println!(
                "OKAY: There are {} matching procs (less than {})",
                matches.len(),
                o
            ),
            (None, Some(u)) => println!(
                "OKAY: There are {} matching procs (greater than {})",
                matches.len(),
                u
            ),
            (None, None) => unreachable!(),
        }
    }
    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();

    debug!("filtering for procs that match: {:?} {:?}", re, states);

    procs
        .iter()
        .filter(|&(pid, _process)| *pid != me && *pid != parent)
        .filter(|&(_pid, process)| states.is_empty() || states.contains(&process.stat.state))
        .filter(|&(_pid, process)| {
            re.as_ref()
                .map(|re| {
                    let cmdline = process.useful_cmdline();
                    let is_match = re.is_match(&cmdline);
                    trace!("is_match={} process={}", is_match, cmdline);
                    is_match
                })
                .unwrap_or(true)
        })
        .collect()
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn validate_parse_zombies_and_pattern() {
        let args = Args::from_iter(["c-p", "--state", "zombie", "--", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "--state=zombie", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "so.*proc", "--state", "zombie", "waiting"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Waiting]);
    }

    // Waiting for structopt 0.2.8 to be released with the from_iter_safe method
    // #[test]
    // #[should_panic]
    // fn regexes_must_be_valid() {
    //     Args::from_iter_safe(["c-p", "--crit-over=5", "["].into_iter()).unwrap();
    // }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs.clone());

        let filtered = filter_procs(regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].1, &procs[2]);

        let filtered = filter_procs(regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<|endoftext|>


<commit_before>use crate::io::{Decode, DecodeError, Encode};
use crate::visit::Visit;

impl Encode for f32 {
    fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
        w.write_all(&self.to_bits().to_le_bytes())
    }
}

impl Decode for f32 {
    fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
        let mut bytes = [0; 4];
        r.read_exact(&mut bytes)?;
        Ok(f32::from_bits(u32::from_le_bytes(bytes)))
    }
}

impl Visit for f32 {}

impl Encode for f64 {
    fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
        w.write_all(&self.to_bits().to_le_bytes())
    }
}

impl Decode for f64 {
    fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
        let mut bytes = [0; std::mem::size_of::<f64>()];
        r.read_exact(&mut bytes)?;
        Ok(f64::from_bits(u64::from_le_bytes(bytes)))
    }
}

impl Visit for f64 {}
<commit_msg>Use macro for floats<commit_after>use crate::io::{Decode, DecodeError, Encode};
use crate::visit::Visit;

macro_rules! def_float {
    ($ty:ident) => {
        impl Encode for $ty {
            fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
                w.write_all(&self.to_le_bytes())
            }
        }

        impl Decode for $ty {
            fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                let mut bytes = [0; std::mem::size_of::<Self>()];
                r.read_exact(&mut bytes)?;
                Ok($ty::from_le_bytes(bytes))
            }
        }

        impl Visit for $ty {}
    };
}

def_float!(f32);
def_float!(f64);
<|endoftext|>


<commit_before>//! byt - io
//!
//! Implements a piece table to abstract over accessing and
//! modifying a file on disk.

// EXTERNS

// LIBRARY INCLUDES
use std::fmt;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Empty;
use std::io::BufReader;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io;

// SUBMODULES
mod tests;

// LOCAL INCLUDES

#[derive(Debug, Copy, Clone, PartialEq)]
enum SourceFile {
    /// The original file on disk.
    Original,
    /// The append file, held in memory.
    Append
}

/// A Piece stores some information about a part of the
/// modified file.
#[derive(Clone)]
struct Piece {
    /// The file this piece refers to.
    file : SourceFile,
    /// The offset (in bytes) of the text in this Piece in its file.
    file_offset : u64,
    /// The length of the text in this Piece.
    length : u64,
    /// The logical offset of the Piece.
    logical_offset : u64,
}

impl fmt::Display for Piece {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}:{} ({}:{})", self.file, self.file_offset, self.logical_offset, self.length)
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Operation {
    Insert,
    Delete
}

/// Records a particular modification of the text.
struct Action {
    op     : Operation,
    offset : u64,
    pieces : Vec<Piece>,
    length : u64
}

/// Implements logical operations on a file that are not written
/// until asked to.
pub struct PieceFile {
    /// Stores all mutation operations of the file so that we can
    /// undo.
    actions : Vec<Action>,
    /// The in-memory append string. All edits will refer to bytes
    /// stored here.
    append_file : String,
    /// The total size (in bytes) of the PieceFile.
    length : u64,
    /// The current offset for reads.
    offset : u64,
    /// The path of the file this PieceFile refers to.
    path : String,
    /// Stores all current Pieces.
    piece_table : Vec<Piece>,
    /// The seekable file reader.
    reader : Option<BufReader<File>>,
}

impl PieceFile {
    // #################################
    // P R I V A T E  F U N C T I O N S
    // #################################
    /// Get a Piece at a particular offset.
    fn get_at_offset(&self, offset : u64) -> Option<usize> {
        let mut logical_length = 0;

        // This has to be O(N) for the time being unless we want
        // to update all logical offsets when the table changes.
        // On modern hardware this will hardly be a bottleneck,
        // but we'll see.
        for index in 0..self.piece_table.len() {
            let piece = &self.piece_table[index];
            logical_length += piece.length;

            if offset > logical_length {
                continue;
            }

            return Some(index);
        }

        None
    }

    /// Update the logical offsets starting at a certain index.
    fn update_offsets(&mut self, start_index : usize) {
        // Don't do anything if this is the last index
        if start_index == self.piece_table.len() - 1 {
            return;
        }

        let mut offset = self.piece_table[start_index].logical_offset +
                         self.piece_table[start_index].length;

        for index in start_index + 1 .. self.piece_table.len() {
            let piece = &mut self.piece_table[index];
            piece.logical_offset = offset;
            offset += piece.length;
        }
    }

    // ###############################
    // P U B L I C  F U N C T I O N S
    // ###############################
    /// Open a new PieceFile. If the file doesn't exist, it is created.
    pub fn open(path : String) -> io::Result<PieceFile> {
        let file = OpenOptions::new()
            .read(true)
            .open(path.clone()).unwrap();

        // Get the length of the file to initialize the first Piece
        let size = file.metadata().unwrap().len() as u64;

        let mut piece_file = PieceFile {
            actions     : Vec::new(),
            append_file : String::new(),
            length      : size,
            offset      : 0,
            path,
            piece_table : Vec::new(),
            reader      : Some(BufReader::new(file)),
        };

        piece_file.piece_table.push(Piece {
            file : SourceFile::Original,
            file_offset : 0,
            length : size as u64,
            logical_offset : 0,
        });

        Ok(piece_file)
    }

    /// Create a new empty PieceFile.
    pub fn empty() -> io::Result<PieceFile> {
        let mut piece_file = PieceFile {
            actions     : Vec::new(),
            append_file : String::new(),
            length      : 0,
            offset      : 0,
            path        : String::from(""),
            piece_table : Vec::new(),
            reader      : None,
        };

        Ok(piece_file)
    }

    /// Insert some text.
    pub fn insert(&mut self, text : &str, offset : u64) {
        let length = text.len() as u64;
        let append_offset = self.append_file.len() as u64;

        self.append_file += text;
        self.length      += length;

        let mut piece = Piece {
            file   : SourceFile::Append,
            file_offset : append_offset,
            logical_offset : 0, // Unknown right now
            length,
        };

        let action = Action {
            op     : Operation::Insert,
            pieces : Vec::new(),
            offset,
            length,
        };

        self.actions.push(action);

        // There are edge cases if you do an insert
        // at the beginning or end of the file.
        if offset == 0 {
            self.piece_table.insert(0, piece);
            self.update_offsets(0);
            return;
        }
        else if offset + length == self.length {
            piece.logical_offset = self.length - length;
            self.piece_table.push(piece);
            // No need to update offsets, this is the last piece
            return;
        }

        // The insertion may create as many as three pieces. It can
        // split a piece that already exists into two parts and then
        // goes in between them.
        let split_index = match self.get_at_offset(offset) {
            Some(v) => v,
            None => return // TODO: Handle this better
        };
        let split_piece = self.piece_table.remove(split_index);

        let lower_size  = offset - split_piece.logical_offset;
        let upper_size  = (split_piece.logical_offset + split_piece.length) - offset;

        if upper_size > 0 {
            self.piece_table.insert(split_index, Piece {
                file           : split_piece.file,
                file_offset    : split_piece.file_offset + lower_size,
                length         : upper_size,
                logical_offset : split_piece.logical_offset + lower_size + length,
            });
        }

        self.piece_table.insert(split_index, piece);

        if lower_size > 0 {
            self.piece_table.insert(split_index, Piece {
                file           : split_piece.file,
                file_offset    : split_piece.file_offset,
                length         : lower_size,
                logical_offset : split_piece.logical_offset,
            });
        }
        self.update_offsets(split_index);
    }

    /// Delete some text.
    pub fn delete(&mut self, offset : u64, length : u64) {
        let start_offset = offset;
        let end_offset   = offset + length;
        let start_index  = self.get_at_offset(start_offset).unwrap();
        let end_index    = self.get_at_offset(end_offset).unwrap();
        let delete_size  = ((end_index - start_index) as u64) + 1;
        let num_pieces   = (end_index - start_index) + 1;

        let mut action = Action {
            length,
            offset,
            op     : Operation::Delete,
            pieces : Vec::new()
        };

        // TODO: ensure we don't overflow
        self.length -= length;

        // Deletes can affect multiple pieces if the user wants to delete
        // across piece boundaries. We have to start at the bottom index
        // and move up. The code below doesn't result in a net positive number
        // of pieces in the piece table. We do any deletion necessary in a second
        // pass.
        for index in start_index .. end_index + 1 {
            let piece              = self.piece_table[index].clone();
            let piece_start_offset = piece.logical_offset;
            let piece_end_offset   = piece.logical_offset + piece.length;

            // Edge case : delete is embedded WITHIN a single piece
            if start_offset > piece_start_offset && end_offset < piece_end_offset {
                self.piece_table.remove(index);
                let upper_size = piece_end_offset - end_offset;
                let lower_size = start_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : delete_size,
                    logical_offset : start_offset,
                });

                // We insert the upper part first
                self.piece_table.insert(index, Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size + delete_size,
                    length         : upper_size,
                    logical_offset : start_offset,
                });

                // Then the lower part
                self.piece_table.insert(index, Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset,
                    length         : lower_size,
                    logical_offset : piece.logical_offset,
                });
                break
            }

            // The deletion area starts in this piece and goes on to another
            if start_offset > piece_start_offset && end_offset >= piece_end_offset {
                let piece      = &mut self.piece_table[index];
                let upper_size = piece_end_offset - start_offset;
                let lower_size = start_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : upper_size,
                    logical_offset : piece_start_offset + lower_size,
                });

                piece.length = lower_size;
                continue;
            }

            // The deletion area finishes in this piece
            if start_offset < piece_start_offset && end_offset <= piece_end_offset {
                let piece      = &mut self.piece_table[index];
                let upper_size = piece_end_offset - end_offset;
                let lower_size = end_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : lower_size,
                    logical_offset : piece_start_offset - lower_size,
                });

                piece.file_offset += lower_size;
                piece.length       = upper_size;
                continue;
            }

            // The deletion area is outside the bounds of this piece
            if start_offset <= piece_start_offset && end_offset >= piece_end_offset {
                let piece = &mut self.piece_table[index];
                action.pieces.push(piece.clone());
                piece.length = 0;
            }
        }

        self.actions.push(action);

        // It's possible that we left zero-length pieces above. We should delete them.
        self.piece_table.retain(|ref v| v.length > 0);
        self.update_offsets(start_index);
    }
}

impl Read for PieceFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        Ok(0)
    }
}

impl Seek for PieceFile {
    fn seek(&mut self, pos : SeekFrom) -> io::Result<u64> {
        let newOffset = match pos {
            SeekFrom::Start(v)   => v as i64,
            SeekFrom::End(v)     => (self.length as i64) + v,
            SeekFrom::Current(v) => (self.offset as i64) + v,
        };

        if newOffset < 0 {
            panic!("Seek offset less than 0");
        }

        Ok(newOffset as u64)
    }
}
<commit_msg>feat(PieceFile): implement read_piece<commit_after>//! byt - io
//!
//! Implements a piece table to abstract over accessing and
//! modifying a file on disk.

// EXTERNS

// LIBRARY INCLUDES
use std::fmt;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Empty;
use std::io::BufReader;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io;
use std::str;

// SUBMODULES
mod tests;

// LOCAL INCLUDES

#[derive(Debug, Copy, Clone, PartialEq)]
enum SourceFile {
    /// The original file on disk.
    Original,
    /// The append file, held in memory.
    Append
}

/// A Piece stores some information about a part of the
/// modified file.
#[derive(Clone)]
struct Piece {
    /// The file this piece refers to.
    file : SourceFile,
    /// The offset (in bytes) of the text in this Piece in its file.
    file_offset : u64,
    /// The length of the text in this Piece.
    length : u64,
    /// The logical offset of the Piece.
    logical_offset : u64,
}

impl Piece {
    /// Convert a logical offset, which is piece-table global, to an offset inside
    /// the Piece's file.
    pub fn logical_to_file(&self, offset : u64) -> u64 {
        return (offset - self.logical_offset) + self.file_offset;
    }
}

impl fmt::Display for Piece {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}:{} ({}:{})", self.file, self.file_offset, self.logical_offset, self.length)
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Operation {
    Insert,
    Delete
}

/// Records a particular modification of the text.
struct Action {
    op     : Operation,
    offset : u64,
    pieces : Vec<Piece>,
    length : u64
}

/// Implements logical operations on a file that are not written
/// until asked to.
pub struct PieceFile {
    /// Stores all mutation operations of the file so that we can
    /// undo.
    actions : Vec<Action>,
    /// The in-memory append string. All edits will refer to bytes
    /// stored here.
    append_file : String,
    /// The total size (in bytes) of the PieceFile.
    length : u64,
    /// The current offset for reads.
    offset : u64,
    /// The path of the file this PieceFile refers to.
    path : String,
    /// Stores all current Pieces.
    piece_table : Vec<Piece>,
    /// The seekable file reader.
    reader : Option<BufReader<File>>,
}

impl PieceFile {
    // #################################
    // P R I V A T E  F U N C T I O N S
    // #################################
    /// Get a Piece at a particular offset.
    fn get_at_offset(&self, offset : u64) -> Option<usize> {
        let mut logical_length = 0;

        // This has to be O(N) for the time being unless we want
        // to update all logical offsets when the table changes.
        // On modern hardware this will hardly be a bottleneck,
        // but we'll see.
        for index in 0..self.piece_table.len() {
            let piece = &self.piece_table[index];
            logical_length += piece.length;

            if offset > logical_length {
                continue;
            }

            return Some(index);
        }

        None
    }

    /// Update the logical offsets starting at a certain index.
    fn update_offsets(&mut self, start_index : usize) {
        // Don't do anything if this is the last index
        if start_index == self.piece_table.len() - 1 {
            return;
        }

        let mut offset = self.piece_table[start_index].logical_offset +
                         self.piece_table[start_index].length;

        for index in start_index + 1 .. self.piece_table.len() {
            let piece = &mut self.piece_table[index];
            piece.logical_offset = offset;
            offset += piece.length;
        }
    }

    // ###############################
    // P U B L I C  F U N C T I O N S
    // ###############################
    /// Open a new PieceFile. If the file doesn't exist, it is created.
    pub fn open(path : String) -> io::Result<PieceFile> {
        let file = OpenOptions::new()
            .read(true)
            .open(path.clone()).unwrap();

        // Get the length of the file to initialize the first Piece
        let size = file.metadata().unwrap().len() as u64;

        let mut piece_file = PieceFile {
            actions     : Vec::new(),
            append_file : String::new(),
            length      : size,
            offset      : 0,
            path,
            piece_table : Vec::new(),
            reader      : Some(BufReader::new(file)),
        };

        piece_file.piece_table.push(Piece {
            file : SourceFile::Original,
            file_offset : 0,
            length : size as u64,
            logical_offset : 0,
        });

        Ok(piece_file)
    }

    /// Create a new empty PieceFile.
    pub fn empty() -> io::Result<PieceFile> {
        let mut piece_file = PieceFile {
            actions     : Vec::new(),
            append_file : String::new(),
            length      : 0,
            offset      : 0,
            path        : String::from(""),
            piece_table : Vec::new(),
            reader      : None,
        };

        Ok(piece_file)
    }

    /// Insert some text.
    pub fn insert(&mut self, text : &str, offset : u64) {
        let length = text.len() as u64;
        let append_offset = self.append_file.len() as u64;

        self.append_file += text;
        self.length      += length;

        let mut piece = Piece {
            file   : SourceFile::Append,
            file_offset : append_offset,
            logical_offset : 0, // Unknown right now
            length,
        };

        let action = Action {
            op     : Operation::Insert,
            pieces : Vec::new(),
            offset,
            length,
        };

        self.actions.push(action);

        // There are edge cases if you do an insert
        // at the beginning or end of the file.
        if offset == 0 {
            self.piece_table.insert(0, piece);
            self.update_offsets(0);
            return;
        }
        else if offset + length == self.length {
            piece.logical_offset = self.length - length;
            self.piece_table.push(piece);
            // No need to update offsets, this is the last piece
            return;
        }

        // The insertion may create as many as three pieces. It can
        // split a piece that already exists into two parts and then
        // goes in between them.
        let split_index = match self.get_at_offset(offset) {
            Some(v) => v,
            None => return // TODO: Handle this better
        };
        let split_piece = self.piece_table.remove(split_index);

        let lower_size  = offset - split_piece.logical_offset;
        let upper_size  = (split_piece.logical_offset + split_piece.length) - offset;

        if upper_size > 0 {
            self.piece_table.insert(split_index, Piece {
                file           : split_piece.file,
                file_offset    : split_piece.file_offset + lower_size,
                length         : upper_size,
                logical_offset : split_piece.logical_offset + lower_size + length,
            });
        }

        self.piece_table.insert(split_index, piece);

        if lower_size > 0 {
            self.piece_table.insert(split_index, Piece {
                file           : split_piece.file,
                file_offset    : split_piece.file_offset,
                length         : lower_size,
                logical_offset : split_piece.logical_offset,
            });
        }
        self.update_offsets(split_index);
    }

    /// Delete some text.
    pub fn delete(&mut self, offset : u64, length : u64) {
        let start_offset = offset;
        let end_offset   = offset + length;
        let start_index  = self.get_at_offset(start_offset).unwrap();
        let end_index    = self.get_at_offset(end_offset).unwrap();
        let delete_size  = ((end_index - start_index) as u64) + 1;
        let num_pieces   = (end_index - start_index) + 1;

        let mut action = Action {
            length,
            offset,
            op     : Operation::Delete,
            pieces : Vec::new()
        };

        // TODO: ensure we don't overflow
        self.length -= length;

        // Deletes can affect multiple pieces if the user wants to delete
        // across piece boundaries. We have to start at the bottom index
        // and move up. The code below doesn't result in a net positive number
        // of pieces in the piece table. We do any deletion necessary in a second
        // pass.
        for index in start_index .. end_index + 1 {
            let piece              = self.piece_table[index].clone();
            let piece_start_offset = piece.logical_offset;
            let piece_end_offset   = piece.logical_offset + piece.length;

            // Edge case : delete is embedded WITHIN a single piece
            if start_offset > piece_start_offset && end_offset < piece_end_offset {
                self.piece_table.remove(index);
                let upper_size = piece_end_offset - end_offset;
                let lower_size = start_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : delete_size,
                    logical_offset : start_offset,
                });

                // We insert the upper part first
                self.piece_table.insert(index, Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size + delete_size,
                    length         : upper_size,
                    logical_offset : start_offset,
                });

                // Then the lower part
                self.piece_table.insert(index, Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset,
                    length         : lower_size,
                    logical_offset : piece.logical_offset,
                });
                break
            }

            // The deletion area starts in this piece and goes on to another
            if start_offset > piece_start_offset && end_offset >= piece_end_offset {
                let piece      = &mut self.piece_table[index];
                let upper_size = piece_end_offset - start_offset;
                let lower_size = start_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : upper_size,
                    logical_offset : piece_start_offset + lower_size,
                });

                piece.length = lower_size;
                continue;
            }

            // The deletion area finishes in this piece
            if start_offset < piece_start_offset && end_offset <= piece_end_offset {
                let piece      = &mut self.piece_table[index];
                let upper_size = piece_end_offset - end_offset;
                let lower_size = end_offset - piece_start_offset;

                action.pieces.push(Piece {
                    file           : piece.file,
                    file_offset    : piece.file_offset + lower_size,
                    length         : lower_size,
                    logical_offset : piece_start_offset - lower_size,
                });

                piece.file_offset += lower_size;
                piece.length       = upper_size;
                continue;
            }

            // The deletion area is outside the bounds of this piece
            if start_offset <= piece_start_offset && end_offset >= piece_end_offset {
                let piece = &mut self.piece_table[index];
                action.pieces.push(piece.clone());
                piece.length = 0;
            }
        }

        self.actions.push(action);

        // It's possible that we left zero-length pieces above. We should delete them.
        self.piece_table.retain(|ref v| v.length > 0);
        self.update_offsets(start_index);
    }

    /// Read characters from the buffer. The result is guaranteed to contain only
    /// whole UTF-8 graphemes, but the result is guaranteed to have at MOST
    /// `num_bytes`. If taking the provided number of bytes would result in a panic
    /// (i.e if it falls on a UTF8 grapheme boundary) then it takes one grapheme
    /// fewer.
    pub fn read(&self, num_bytes : u64) -> io::Result<Box<String>> {
        let mut result = Box::new(String::new());

        let start_offset = self.offset;
        let start_index  = self.get_at_offset(start_offset).unwrap();
        let end_offset   = self.offset + num_bytes - 1;
        let end_index    = self.get_at_offset(end_offset).unwrap();
        let num_pieces   = end_index - start_index + 1;

        Ok(result)
    }

    /// Reads characters from a piece into a destination string. The `offset` refers to
    /// logical offset in the whole piece table, not file-specific offset.
    fn read_piece(&mut self, piece : &Piece, offset : u64, num_bytes : u64, dest : &mut String) {
        let piece_start_offset = piece.logical_offset;
        let piece_end_offset   = piece_start_offset + piece.length - 1;
        let mut buf = vec![0 as u8; num_bytes as usize];

        match piece.file {
            SourceFile::Append => {
                let append_start_offset   = piece.logical_to_file(offset);
                let mut append_end_offset = (append_start_offset + num_bytes - 1);
                let append_bytes          = self.append_file.as_bytes();

                let append_slice = append_bytes
                    .get(append_start_offset as usize ..
                         append_end_offset as usize)
                    .unwrap();

                buf.clone_from_slice(append_slice);
            },
            SourceFile::Original => {
                let mut reader = self.reader.as_mut().unwrap();
                reader.read(buf.as_mut_slice());
            },
        }

        let mut slice = buf.as_slice();
        let mut index = buf.len() - 1;
        let mut converted = str::from_utf8(&buf[0..index]);

        // It's possible we'll stumble upon a boundary. If we do, reduce the number of
        // bytes until all is well.
        while converted.is_err() {
            index -= 1;
            converted = str::from_utf8(&buf[0..index]);
        }

        dest.push_str(converted.unwrap());
    }
}

impl Seek for PieceFile {
    fn seek(&mut self, pos : SeekFrom) -> io::Result<u64> {
        let newOffset = match pos {
            SeekFrom::Start(v)   => v as i64,
            SeekFrom::End(v)     => (self.length as i64) + v,
            SeekFrom::Current(v) => (self.offset as i64) + v,
        };

        if newOffset < 0 {
            panic!("Seek offset less than 0");
        }

        Ok(newOffset as u64)
    }
}
<|endoftext|>


<commit_before>//! TODO Fill in

use std::{ptr, mem};
use xcb::{xkb, ffi, Connection};
use nix::{self, libc};
use render;
use std::ffi::CStr;
use gdk_pixbuf::Pixbuf;
use glib::translate::ToGlibPtr;
use std::fmt::{self, Display, Formatter};
use std::process::{Command, Stdio};
use std::thread;
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua, Value, LightUserData};
use super::{XCB_CONNECTION_HANDLE, signal};

// TODO this isn't defined in the xcb crate, even though it should be.
// A patch should be open adding this to its generation scheme
extern "C" {
    fn xcb_xkb_get_names_value_list_unpack(buffer: *mut libc::c_void,
                                           nTypes: u8,
                                           indicators: u32,
                                           virtualMods: u16,
                                           groupNames: u8,
                                           nKeys: u8,
                                           nKeyAliases: u8,
                                           nRadioGroups: u8,
                                           which: u32,
                                           aux: *mut ffi::xkb::xcb_xkb_get_names_value_list_t) -> libc::c_int;
}

#[derive(Clone, Debug)]
pub struct AwesomeState {
    preferred_icon_size: u32
}

impl Default for AwesomeState {
    fn default() -> Self {
        AwesomeState {
            preferred_icon_size: 0
        }
    }
}

impl Display for AwesomeState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Awesome: {:p}", self)
    }
}

impl UserData for AwesomeState {}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let awesome_table = lua.create_table();
    state_setup(lua, &awesome_table)?;
    meta_setup(lua, &awesome_table)?;
    method_setup(lua, &awesome_table)?;
    property_setup(lua, &awesome_table)?;
    let globals = lua.globals();
    globals.set("awesome", awesome_table)
}

fn state_setup(lua: &Lua, awesome_table: &Table) -> rlua::Result<()> {
    awesome_table.set("__data", AwesomeState::default().to_lua(lua)?)
}

fn meta_setup(lua: &Lua, awesome_table: &Table) -> rlua::Result<()> {
    let meta_table = awesome_table.get_metatable().unwrap_or_else(|| {
        let table = lua.create_table();
        awesome_table.set_metatable(Some(table.clone()));
        table
    });
    meta_table.set("__tostring", lua.create_function(|_, val: Table| {
        Ok(format!("{}", val.get::<_, AwesomeState>("__data")?))
    }))
}

fn method_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Fill in rest
    awesome_table.set("connect_signal", lua.create_function(signal::global_connect_signal))?;
    awesome_table.set("disconnect_signal", lua.create_function(signal::global_disconnect_signal))?;
    awesome_table.set("emit_signal", lua.create_function(signal::global_emit_signal))?;
    awesome_table.set("xrdb_get_value", lua.create_function(xrdb_get_value))?;
    awesome_table.set("xkb_set_layout_group", lua.create_function(xkb_set_layout_group))?;
    awesome_table.set("xkb_get_layout_group", lua.create_function(xkb_get_layout_group))?;
    awesome_table.set("set_preferred_icon_size", lua.create_function(set_preferred_icon_size))?;
    awesome_table.set("register_xproperty", lua.create_function(register_xproperty))?;
    awesome_table.set("xkb_get_group_names", lua.create_function(xkb_get_group_names))?;
    awesome_table.set("set_xproperty", lua.create_function(set_xproperty))?;
    awesome_table.set("get_xproperty", lua.create_function(get_xproperty))?;
    awesome_table.set("restart", lua.create_function(restart))?;
    awesome_table.set("load_image", lua.create_function(load_image))?;
    awesome_table.set("sync", lua.create_function(sync))?;
    awesome_table.set("exec", lua.create_function(exec))?;
    awesome_table.set("kill", lua.create_function(kill))?;
    awesome_table.set("quit", lua.create_function(quit))
}

fn property_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Do properly
    awesome_table.set("version", "0".to_lua(lua)?)?;
    awesome_table.set("themes_path", "/usr/share/awesome/themes".to_lua(lua)?)?;
    awesome_table.set("conffile", "".to_lua(lua)?)
}

/// Registers a new X property
/// This actually does nothing, since this is Wayland.
fn register_xproperty<'lua>(_: &'lua Lua, _: Value<'lua>) -> rlua::Result<()> {
    warn!("register_xproperty not supported");
    Ok(())
}

/// Get layout short names
fn xkb_get_group_names<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    let xcb_con = lua.globals().get::<_, LightUserData>(XCB_CONNECTION_HANDLE)?.0;
    // FIXME
    // hmm surely we can do this safely
    unsafe {
        let con = Connection::from_raw_conn(xcb_con as _);
        let raw_con = con.get_raw_conn();
        let id = xkb::ID_USE_CORE_KBD as _;
        // The structure here looks weird because we need to ensure Connection
        // isn't cleaned up even in the event of an error...
        let names_r = {
            let names_cookie = xkb::get_names_unchecked(&con, id, xkb::NAME_DETAIL_SYMBOLS);
            names_cookie.get_reply()
        };
        mem::forget(con);
        let names_r = match names_r {
            Ok(names_r) => {
                if names_r.ptr.is_null() {
                    warn!("Failed to get xkb symbols name");
                    return Ok(Value::Nil)
                }
                names_r
            },
            Err(err) => {
                warn!("Failed to get xkb symbols name {:?}", err);
                return Ok(Value::Nil)
            }
        };
        let buffer = ffi::xkb::xcb_xkb_get_names_value_list(names_r.ptr);
        if buffer.is_null() {
            warn!("Returned buffer was NULL");
            return Ok(Value::Nil)
        }
        let names_r_ptr = names_r.ptr;
        if names_r_ptr.is_null() {
            warn!("Name reply pointer was null");
            return Ok(Value::Nil)
        }
        let mut names_list: ffi::xkb::xcb_xkb_get_names_value_list_t = mem::uninitialized();
        xcb_xkb_get_names_value_list_unpack(buffer, (*names_r_ptr).nTypes, (*names_r_ptr).indicators, (*names_r_ptr).virtualMods,
                                            (*names_r_ptr).groupNames, (*names_r_ptr).nKeys, (*names_r_ptr).nKeyAliases,
                                            (*names_r_ptr).nRadioGroups, (*names_r_ptr).which, &mut names_list);
        let atom_name_c = ffi::xproto::xcb_get_atom_name_unchecked(raw_con, names_list.symbolsName);
        let atom_name_r = ffi::xproto::xcb_get_atom_name_reply(raw_con, atom_name_c, ptr::null_mut());
        if atom_name_r.is_null() {
            warn!("Failed to get atom symbols name");
            return Ok(Value::Nil)
        }
        let name_c = ffi::xproto::xcb_get_atom_name_name(atom_name_r);
        CStr::from_ptr(name_c).to_string_lossy().into_owned().to_lua(lua)
    }
}

/// Restart Awesome by restarting the Lua thread
fn restart<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<()> {
    use lua::{self, LuaQuery};
    if let Err(err) = lua::send(LuaQuery::Restart) {
        warn!("Could not restart Lua thread {:#?}", err);
    }
    Ok(())
}

/// Load an image from the given path
/// Returns either a cairo surface as light user data, nil and an error message
fn load_image<'lua>(lua: &'lua Lua, file_path: String) -> rlua::Result<Value<'lua>> {
    let pixbuf = Pixbuf::new_from_file(file_path.as_str())
        .map_err(|err| rlua::Error::RuntimeError(format!("{}", err)))?;
    let surface = render::load_surface_from_pixbuf(pixbuf);
    // UGH, I wanted to do to_glib_full, but that isn't defined apparently
    // So now I have to ignore the lifetime completely and just forget about the surface.
    let surface_ptr = surface.to_glib_none().0;
    ::std::mem::forget(surface);
    rlua::LightUserData(surface_ptr as _).to_lua(lua)
}

fn exec(_: &Lua, command: String) -> rlua::Result<()> {
    trace!("exec: \"{}\"", command);
    thread::Builder::new().name(command.clone()).spawn(|| {
        Command::new(command)
            .stdout(Stdio::null())
            .spawn()
            .expect("Could not spawn command")
    }).expect("Unable to spawn thread");
    Ok(())
}

/// Kills a PID with the given signal
///
/// Returns false if it could not send the signal to that process
fn kill(_: &Lua, (pid, sig): (libc::pid_t, libc::c_int)) -> rlua::Result<bool> {
    Ok(nix::sys::signal::kill(pid, sig).is_ok())
}

fn set_preferred_icon_size(lua: &Lua, val: u32) -> rlua::Result<()> {
    let mut awesome_state: AwesomeState = lua.globals().get::<_, Table>("awesome")?.get("__data")?;
    awesome_state.preferred_icon_size = val;
    lua.globals().get::<_, Table>("awesome")?.set("__data", awesome_state.to_lua(lua)?)

}

fn quit(_: &Lua, _: ()) -> rlua::Result<()> {
    ::rustwlc::terminate();
    Ok(())
}

/// No need to sync in Wayland
fn sync(_: &Lua, _: ()) -> rlua::Result<()> { Ok(()) }

fn set_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("set_xproperty not supported");
    Ok(())
}

fn get_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("get_xproperty not supported");
    Ok(())
}

fn xkb_set_layout_group(_: &Lua, _group: i32) -> rlua::Result<()> {
    warn!("xkb_set_layout_group not supported; Wait until wlroots");
    Ok(())
}

fn xkb_get_layout_group(_: &Lua, _: ()) -> rlua::Result<i32> {
    warn!("xkb_get_layout_group not supported; Wait until wlroots");
    warn!("Returning dummy data so I can move forward");
    warn!("Replace me before release!");
    Ok(0)
}

fn xrdb_get_value(lua: &Lua, (_resource_class, resource_name): (String, String))
                  -> rlua::Result<Value> {
    if &resource_name == "Xft.dpi" {
        return Ok("1".to_lua(lua)?)
    }
    warn!("xrdb_get_value not supported");
    Ok(Value::Nil)
}
<commit_msg>Implemented xkb_get_layout_group<commit_after>//! TODO Fill in

use std::{ptr, mem};
use xcb::{xkb, ffi, Connection};
use nix::{self, libc};
use render;
use std::ffi::CStr;
use gdk_pixbuf::Pixbuf;
use glib::translate::ToGlibPtr;
use std::fmt::{self, Display, Formatter};
use std::process::{Command, Stdio};
use std::thread;
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua, Value, LightUserData};
use super::{XCB_CONNECTION_HANDLE, signal};

// TODO this isn't defined in the xcb crate, even though it should be.
// A patch should be open adding this to its generation scheme
extern "C" {
    fn xcb_xkb_get_names_value_list_unpack(buffer: *mut libc::c_void,
                                           nTypes: u8,
                                           indicators: u32,
                                           virtualMods: u16,
                                           groupNames: u8,
                                           nKeys: u8,
                                           nKeyAliases: u8,
                                           nRadioGroups: u8,
                                           which: u32,
                                           aux: *mut ffi::xkb::xcb_xkb_get_names_value_list_t) -> libc::c_int;
}

#[derive(Clone, Debug)]
pub struct AwesomeState {
    preferred_icon_size: u32
}

impl Default for AwesomeState {
    fn default() -> Self {
        AwesomeState {
            preferred_icon_size: 0
        }
    }
}

impl Display for AwesomeState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Awesome: {:p}", self)
    }
}

impl UserData for AwesomeState {}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let awesome_table = lua.create_table();
    state_setup(lua, &awesome_table)?;
    meta_setup(lua, &awesome_table)?;
    method_setup(lua, &awesome_table)?;
    property_setup(lua, &awesome_table)?;
    let globals = lua.globals();
    globals.set("awesome", awesome_table)
}

fn state_setup(lua: &Lua, awesome_table: &Table) -> rlua::Result<()> {
    awesome_table.set("__data", AwesomeState::default().to_lua(lua)?)
}

fn meta_setup(lua: &Lua, awesome_table: &Table) -> rlua::Result<()> {
    let meta_table = awesome_table.get_metatable().unwrap_or_else(|| {
        let table = lua.create_table();
        awesome_table.set_metatable(Some(table.clone()));
        table
    });
    meta_table.set("__tostring", lua.create_function(|_, val: Table| {
        Ok(format!("{}", val.get::<_, AwesomeState>("__data")?))
    }))
}

fn method_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Fill in rest
    awesome_table.set("connect_signal", lua.create_function(signal::global_connect_signal))?;
    awesome_table.set("disconnect_signal", lua.create_function(signal::global_disconnect_signal))?;
    awesome_table.set("emit_signal", lua.create_function(signal::global_emit_signal))?;
    awesome_table.set("xrdb_get_value", lua.create_function(xrdb_get_value))?;
    awesome_table.set("xkb_set_layout_group", lua.create_function(xkb_set_layout_group))?;
    awesome_table.set("xkb_get_layout_group", lua.create_function(xkb_get_layout_group))?;
    awesome_table.set("set_preferred_icon_size", lua.create_function(set_preferred_icon_size))?;
    awesome_table.set("register_xproperty", lua.create_function(register_xproperty))?;
    awesome_table.set("xkb_get_group_names", lua.create_function(xkb_get_group_names))?;
    awesome_table.set("set_xproperty", lua.create_function(set_xproperty))?;
    awesome_table.set("get_xproperty", lua.create_function(get_xproperty))?;
    awesome_table.set("restart", lua.create_function(restart))?;
    awesome_table.set("load_image", lua.create_function(load_image))?;
    awesome_table.set("sync", lua.create_function(sync))?;
    awesome_table.set("exec", lua.create_function(exec))?;
    awesome_table.set("kill", lua.create_function(kill))?;
    awesome_table.set("quit", lua.create_function(quit))
}

fn property_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Do properly
    awesome_table.set("version", "0".to_lua(lua)?)?;
    awesome_table.set("themes_path", "/usr/share/awesome/themes".to_lua(lua)?)?;
    awesome_table.set("conffile", "".to_lua(lua)?)
}

/// Registers a new X property
/// This actually does nothing, since this is Wayland.
fn register_xproperty<'lua>(_: &'lua Lua, _: Value<'lua>) -> rlua::Result<()> {
    warn!("register_xproperty not supported");
    Ok(())
}

/// Get layout short names
fn xkb_get_group_names<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    let xcb_con = lua.globals().get::<_, LightUserData>(XCB_CONNECTION_HANDLE)?.0;
    // FIXME
    // hmm surely we can do this safely
    unsafe {
        let con = Connection::from_raw_conn(xcb_con as _);
        let raw_con = con.get_raw_conn();
        let id = xkb::ID_USE_CORE_KBD as _;
        // The structure here looks weird because we need to ensure Connection
        // isn't cleaned up even in the event of an error...
        let names_r = {
            let names_cookie = xkb::get_names_unchecked(&con, id, xkb::NAME_DETAIL_SYMBOLS);
            names_cookie.get_reply()
        };
        mem::forget(con);
        let names_r = match names_r {
            Ok(names_r) => {
                if names_r.ptr.is_null() {
                    warn!("Failed to get xkb symbols name");
                    return Ok(Value::Nil)
                }
                names_r
            },
            Err(err) => {
                warn!("Failed to get xkb symbols name {:?}", err);
                return Ok(Value::Nil)
            }
        };
        let buffer = ffi::xkb::xcb_xkb_get_names_value_list(names_r.ptr);
        if buffer.is_null() {
            warn!("Returned buffer was NULL");
            return Ok(Value::Nil)
        }
        let names_r_ptr = names_r.ptr;
        if names_r_ptr.is_null() {
            warn!("Name reply pointer was null");
            return Ok(Value::Nil)
        }
        let mut names_list: ffi::xkb::xcb_xkb_get_names_value_list_t = mem::uninitialized();
        xcb_xkb_get_names_value_list_unpack(buffer, (*names_r_ptr).nTypes, (*names_r_ptr).indicators, (*names_r_ptr).virtualMods,
                                            (*names_r_ptr).groupNames, (*names_r_ptr).nKeys, (*names_r_ptr).nKeyAliases,
                                            (*names_r_ptr).nRadioGroups, (*names_r_ptr).which, &mut names_list);
        let atom_name_c = ffi::xproto::xcb_get_atom_name_unchecked(raw_con, names_list.symbolsName);
        let atom_name_r = ffi::xproto::xcb_get_atom_name_reply(raw_con, atom_name_c, ptr::null_mut());
        if atom_name_r.is_null() {
            warn!("Failed to get atom symbols name");
            return Ok(Value::Nil)
        }
        let name_c = ffi::xproto::xcb_get_atom_name_name(atom_name_r);
        CStr::from_ptr(name_c).to_string_lossy().into_owned().to_lua(lua)
    }
}

/// Restart Awesome by restarting the Lua thread
fn restart<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<()> {
    use lua::{self, LuaQuery};
    if let Err(err) = lua::send(LuaQuery::Restart) {
        warn!("Could not restart Lua thread {:#?}", err);
    }
    Ok(())
}

/// Load an image from the given path
/// Returns either a cairo surface as light user data, nil and an error message
fn load_image<'lua>(lua: &'lua Lua, file_path: String) -> rlua::Result<Value<'lua>> {
    let pixbuf = Pixbuf::new_from_file(file_path.as_str())
        .map_err(|err| rlua::Error::RuntimeError(format!("{}", err)))?;
    let surface = render::load_surface_from_pixbuf(pixbuf);
    // UGH, I wanted to do to_glib_full, but that isn't defined apparently
    // So now I have to ignore the lifetime completely and just forget about the surface.
    let surface_ptr = surface.to_glib_none().0;
    ::std::mem::forget(surface);
    rlua::LightUserData(surface_ptr as _).to_lua(lua)
}

fn exec(_: &Lua, command: String) -> rlua::Result<()> {
    trace!("exec: \"{}\"", command);
    thread::Builder::new().name(command.clone()).spawn(|| {
        Command::new(command)
            .stdout(Stdio::null())
            .spawn()
            .expect("Could not spawn command")
    }).expect("Unable to spawn thread");
    Ok(())
}

/// Kills a PID with the given signal
///
/// Returns false if it could not send the signal to that process
fn kill(_: &Lua, (pid, sig): (libc::pid_t, libc::c_int)) -> rlua::Result<bool> {
    Ok(nix::sys::signal::kill(pid, sig).is_ok())
}

fn set_preferred_icon_size(lua: &Lua, val: u32) -> rlua::Result<()> {
    let mut awesome_state: AwesomeState = lua.globals().get::<_, Table>("awesome")?.get("__data")?;
    awesome_state.preferred_icon_size = val;
    lua.globals().get::<_, Table>("awesome")?.set("__data", awesome_state.to_lua(lua)?)

}

fn quit(_: &Lua, _: ()) -> rlua::Result<()> {
    ::rustwlc::terminate();
    Ok(())
}

/// No need to sync in Wayland
fn sync(_: &Lua, _: ()) -> rlua::Result<()> { Ok(()) }

fn set_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("set_xproperty not supported");
    Ok(())
}

fn get_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("get_xproperty not supported");
    Ok(())
}

fn xkb_set_layout_group(_: &Lua, _group: i32) -> rlua::Result<()> {
    warn!("xkb_set_layout_group not supported; Wait until wlroots");
    Ok(())
}

fn xkb_get_layout_group<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    use xcb::ffi::xkb;
    unsafe {
        let xcb_con = lua.globals().get::<_, LightUserData>(XCB_CONNECTION_HANDLE)?.0;
        let con = Connection::from_raw_conn(xcb_con as _);
        mem::forget(xcb_con);
        let raw_con = con.get_raw_conn();
        let state_c = xkb::xcb_xkb_get_state_unchecked(raw_con, xkb::XCB_XKB_ID_USE_CORE_KBD as _);
        let state_r = xkb::xcb_xkb_get_state_reply(raw_con, state_c, ptr::null_mut());
        if state_r.is_null() {
            warn!("State reply was NULL");
            return Ok(Value::Nil);
        }
        (*state_r).group.to_lua(lua)
    }
}

fn xrdb_get_value(lua: &Lua, (_resource_class, resource_name): (String, String))
                  -> rlua::Result<Value> {
    if &resource_name == "Xft.dpi" {
        return Ok("1".to_lua(lua)?)
    }
    warn!("xrdb_get_value not supported");
    Ok(Value::Nil)
}
<|endoftext|>