<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::attr::{Attr, AttrHelpers};
use dom::bindings::codegen::Bindings::DOMTokenListBinding;
use dom::bindings::codegen::Bindings::DOMTokenListBinding::DOMTokenListMethods;
use dom::bindings::error::Fallible;
use dom::bindings::error::Error::{InvalidCharacter, Syntax};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalRootable};
use dom::bindings::utils::{Reflector, Reflectable, reflect_dom_object};
use dom::element::{Element, AttributeHandlers};
use dom::node::window_from_node;

use servo_util::str::{DOMString, HTML_SPACE_CHARACTERS};
use string_cache::Atom;

#[dom_struct]
pub struct DOMTokenList {
    reflector_: Reflector,
    element: JS<Element>,
    local_name: Atom,
}

impl DOMTokenList {
    pub fn new_inherited(element: JSRef<Element>, local_name: Atom) -> DOMTokenList {
        DOMTokenList {
            reflector_: Reflector::new(),
            element: JS::from_rooted(element),
            local_name: local_name,
        }
    }

    pub fn new(element: JSRef<Element>, local_name: &Atom) -> Temporary<DOMTokenList> {
        let window = window_from_node(element).root();
        reflect_dom_object(box DOMTokenList::new_inherited(element, local_name.clone()),
                           GlobalRef::Window(*window),
                           DOMTokenListBinding::Wrap)
    }
}

impl Reflectable for DOMTokenList {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}

trait PrivateDOMTokenListHelpers {
    fn attribute(self) -> Option<Temporary<Attr>>;
    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str>;
}

impl<'a> PrivateDOMTokenListHelpers for JSRef<'a, DOMTokenList> {
    fn attribute(self) -> Option<Temporary<Attr>> {
        let element = self.element.root();
        element.get_attribute(ns!(""), &self.local_name)
    }

    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str> {
        match token {
            "" => Err(Syntax),
            token if token.find(HTML_SPACE_CHARACTERS).is_some() => Err(InvalidCharacter),
            token => Ok(token)
        }
    }
}

// http://dom.spec.whatwg.org/#domtokenlist
impl<'a> DOMTokenListMethods for JSRef<'a, DOMTokenList> {
    // http://dom.spec.whatwg.org/#dom-domtokenlist-length
    fn Length(self) -> u32 {
        self.attribute().root().map(|attr| {
            attr.value().tokens().map(|tokens| tokens.len()).unwrap_or(0)
        }).unwrap_or(0) as u32
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-item
    fn Item(self, index: u32) -> Option<DOMString> {
        self.attribute().root().and_then(|attr| attr.value().tokens().and_then(|tokens| {
            tokens.get(index as uint).map(|token| token.as_slice().to_string())
        }))
    }

    fn IndexedGetter(self, index: u32, found: &mut bool) -> Option<DOMString> {
        let item = self.Item(index);
        *found = item.is_some();
        item
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-contains
    fn Contains(self, token: DOMString) -> Fallible<bool> {
        self.check_token_exceptions(token.as_slice()).map(|slice| {
            self.attribute().root().map(|attr| {
                let value = attr.value();
                let tokens = value.tokens()
                                  .expect("Should have parsed this attribute");
                let atom = Atom::from_slice(slice);
                tokens.iter().any(|token| *token == atom)
            }).unwrap_or(false)
        })
    }
}
<commit_msg>DOMTokenList::check_token_exceptions now returns an Atom<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::attr::{Attr, AttrHelpers};
use dom::bindings::codegen::Bindings::DOMTokenListBinding;
use dom::bindings::codegen::Bindings::DOMTokenListBinding::DOMTokenListMethods;
use dom::bindings::error::{ErrorResult, Fallible};
use dom::bindings::error::Error::{InvalidCharacter, Syntax};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalRootable};
use dom::bindings::utils::{Reflector, Reflectable, reflect_dom_object};
use dom::element::{Element, AttributeHandlers};
use dom::node::window_from_node;

use servo_util::str::{DOMString, HTML_SPACE_CHARACTERS};
use string_cache::Atom;

#[dom_struct]
pub struct DOMTokenList {
    reflector_: Reflector,
    element: JS<Element>,
    local_name: Atom,
}

impl DOMTokenList {
    pub fn new_inherited(element: JSRef<Element>, local_name: Atom) -> DOMTokenList {
        DOMTokenList {
            reflector_: Reflector::new(),
            element: JS::from_rooted(element),
            local_name: local_name,
        }
    }

    pub fn new(element: JSRef<Element>, local_name: &Atom) -> Temporary<DOMTokenList> {
        let window = window_from_node(element).root();
        reflect_dom_object(box DOMTokenList::new_inherited(element, local_name.clone()),
                           GlobalRef::Window(*window),
                           DOMTokenListBinding::Wrap)
    }
}

impl Reflectable for DOMTokenList {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}

trait PrivateDOMTokenListHelpers {
    fn attribute(self) -> Option<Temporary<Attr>>;
    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<Atom>;
}

impl<'a> PrivateDOMTokenListHelpers for JSRef<'a, DOMTokenList> {
    fn attribute(self) -> Option<Temporary<Attr>> {
        let element = self.element.root();
        element.get_attribute(ns!(""), &self.local_name)
    }

    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<Atom> {
        match token {
            "" => Err(Syntax),
            slice if slice.find(HTML_SPACE_CHARACTERS).is_some() => Err(InvalidCharacter),
            slice => Ok(Atom::from_slice(slice))
        }
    }
}

// http://dom.spec.whatwg.org/#domtokenlist
impl<'a> DOMTokenListMethods for JSRef<'a, DOMTokenList> {
    // http://dom.spec.whatwg.org/#dom-domtokenlist-length
    fn Length(self) -> u32 {
        self.attribute().root().map(|attr| {
            attr.value().tokens().map(|tokens| tokens.len()).unwrap_or(0)
        }).unwrap_or(0) as u32
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-item
    fn Item(self, index: u32) -> Option<DOMString> {
        self.attribute().root().and_then(|attr| attr.value().tokens().and_then(|tokens| {
            tokens.get(index as uint).map(|token| token.as_slice().to_string())
        }))
    }

    fn IndexedGetter(self, index: u32, found: &mut bool) -> Option<DOMString> {
        let item = self.Item(index);
        *found = item.is_some();
        item
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-contains
    fn Contains(self, token: DOMString) -> Fallible<bool> {
        self.check_token_exceptions(token.as_slice()).map(|token| {
            self.attribute().root().map(|attr| {
                attr.value()
                    .tokens()
                    .expect("Should have parsed this attribute")
                    .iter()
                    .any(|atom| *atom == token)
            }).unwrap_or(false)
        })
    }
}
<|endoftext|>


<commit_before>use error::*;
use anvil_rpc::editor;

use capnp::capability::Promise;
use capnp_rpc::{RpcSystem, twoparty, rpc_twoparty_capnp};
use futures::{Future, Stream};
use tokio_core::io::Io;
use tokio_core::reactor::Core;
use tokio_uds::UnixListener;
// use xi_rope::Rope;

use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::Path;

struct EditorImpl {
    content: Vec<String>,
}

impl EditorImpl {
    pub fn new() -> EditorImpl {
        EditorImpl { content: vec![String::new()] }
    }
}

impl editor::Server for EditorImpl {
    fn open_file(&mut self,
                 params: editor::OpenFileParams,
                 _results: editor::OpenFileResults)
                 -> Promise<(), ::capnp::Error> {
        let file_name = pry!(pry!(params.get()).get_path());
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(file_name)
            .unwrap();

        let mut buf = String::new();
        file.read_to_string(&mut buf).unwrap();

        let mut contents = Vec::new();
        for line in buf.lines() {
            contents.push(line.to_string());
        }

        self.content = contents;

        Promise::ok(())
    }

    fn write_file(&mut self,
                  params: editor::WriteFileParams,
                  _results: editor::WriteFileResults)
                  -> Promise<(), ::capnp::Error> {
        let file_name = pry!(pry!(params.get()).get_path());
        let mut file = File::create(file_name).unwrap();

        for line in &self.content {
            file.write_all(line.as_bytes()).unwrap();
        }

        Promise::ok(())
    }

    fn insert(&mut self,
              params: editor::InsertParams,
              _results: editor::InsertResults)
              -> Promise<(), ::capnp::Error> {
        let mut line_number = pry!(params.get()).get_line() as usize;
        let column = pry!(params.get()).get_column() as usize;
        let text = pry!(pry!(params.get()).get_string());

        {
            let mut line = &mut self.content[line_number];
            for (i, c) in text.char_indices() {
                match c {
                    '\n' => {
                        // TODO: handle new lines in input
                    }
                    _ => {
                        line.insert(column + i, c);
                    }
                }
            }
        }
        println!("{:?}", self.content);
        Promise::ok(())
    }

    fn delete(&mut self,
              params: editor::DeleteParams,
              _results: editor::DeleteResults)
              -> Promise<(), ::capnp::Error> {
        let line = pry!(params.get()).get_line();
        let column = pry!(params.get()).get_column() as usize;
        let length = pry!(params.get()).get_length() as usize;

        {
            let mut line = &mut self.content[line as usize];
            // TODO: handle cases where column + length > line.len().
            // We want to move onto the next line and continue removing elements.
            for i in 0..length {
                line.remove(column + i);
            }
        }
        println!("{:?}", self.content);
        Promise::ok(())
    }

    fn quit(&mut self,
            _params: editor::QuitParams,
            _results: editor::QuitResults)
            -> Promise<(), ::capnp::Error> {
        println!("Quit request");
        Promise::ok(())
    }
}

pub fn server<P>(path: P) -> Result<()>
    where P: AsRef<Path>
{
    let mut core = Core::new().chain_err(|| "unable to create event loop")?;
    let handle = core.handle();

    let socket = UnixListener::bind(path, &handle).chain_err(|| "unable to bind to UDS")?;

    let editor_impl = EditorImpl::new();

    let editor = editor::ToClient::new(editor_impl).from_server::<::capnp_rpc::Server>();

    let handle1 = handle.clone();
    let done = socket.incoming()
        .for_each(move |(stream, _addr)| {
            let (reader, writer) = stream.split();
            let handle = handle1.clone();

            let network = twoparty::VatNetwork::new(reader,
                                                    writer,
                                                    rpc_twoparty_capnp::Side::Server,
                                                    Default::default());

            let rpc_system = RpcSystem::new(Box::new(network), Some(editor.clone().client));

            handle.spawn(rpc_system.map_err(|_| ()));
            Ok(())
        });

    core.run(done).chain_err(|| "unable to run event loop")?;
    Ok(())
}
<commit_msg>Switch to xi-rope for editor contents<commit_after>use error::*;
use anvil_rpc::editor;

use capnp::capability::Promise;
use capnp_rpc::{RpcSystem, twoparty, rpc_twoparty_capnp};
use futures::{Future, Stream};
use tokio_core::io::Io;
use tokio_core::reactor::Core;
use tokio_uds::UnixListener;
use xi_rope::Rope;

use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::Path;

struct EditorImpl {
    content: Rope,
}

impl EditorImpl {
    pub fn new() -> EditorImpl {
        EditorImpl { content: Rope::from("") }
    }
}

impl editor::Server for EditorImpl {
    fn open_file(&mut self,
                 params: editor::OpenFileParams,
                 _results: editor::OpenFileResults)
                 -> Promise<(), ::capnp::Error> {
        let file_name = pry!(pry!(params.get()).get_path());
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(file_name)
            .unwrap();

        let mut buf = String::new();
        file.read_to_string(&mut buf).unwrap();
        let content = Rope::from(buf);

        self.content = content;

        Promise::ok(())
    }

    fn write_file(&mut self,
                  params: editor::WriteFileParams,
                  _results: editor::WriteFileResults)
                  -> Promise<(), ::capnp::Error> {
        let file_name = pry!(pry!(params.get()).get_path());
        let mut file = File::create(file_name).unwrap();

        file.write_all(String::from(self.content.clone()).as_bytes()).unwrap();

        Promise::ok(())
    }

    fn insert(&mut self,
              params: editor::InsertParams,
              _results: editor::InsertResults)
              -> Promise<(), ::capnp::Error> {
        let line = pry!(params.get()).get_line() as usize;
        let column = pry!(params.get()).get_column() as usize;
        let text = pry!(pry!(params.get()).get_string());

        {
            let index = self.content.offset_of_line(line) + column;
            self.content.edit_str(index, index, text)
        }

        println!("{}", String::from(self.content.clone()));
        Promise::ok(())
    }

    fn delete(&mut self,
              params: editor::DeleteParams,
              _results: editor::DeleteResults)
              -> Promise<(), ::capnp::Error> {
        let line = pry!(params.get()).get_line() as usize;
        let column = pry!(params.get()).get_column() as usize;
        let length = pry!(params.get()).get_length() as usize;

        {
            let index = self.content.offset_of_line(line) + column;
            self.content.edit_str(index, index + length, "");
        }

        println!("{}", String::from(self.content.clone()));
        Promise::ok(())
    }

    fn quit(&mut self,
            _params: editor::QuitParams,
            _results: editor::QuitResults)
            -> Promise<(), ::capnp::Error> {
        println!("Quit request");
        Promise::ok(())
    }
}

pub fn server<P>(path: P) -> Result<()>
    where P: AsRef<Path>
{
    let mut core = Core::new().chain_err(|| "unable to create event loop")?;
    let handle = core.handle();

    let socket = UnixListener::bind(path, &handle).chain_err(|| "unable to bind to UDS")?;

    let editor_impl = EditorImpl::new();

    let editor = editor::ToClient::new(editor_impl).from_server::<::capnp_rpc::Server>();

    let handle1 = handle.clone();
    let done = socket.incoming()
        .for_each(move |(stream, _addr)| {
            let (reader, writer) = stream.split();
            let handle = handle1.clone();

            let network = twoparty::VatNetwork::new(reader,
                                                    writer,
                                                    rpc_twoparty_capnp::Side::Server,
                                                    Default::default());

            let rpc_system = RpcSystem::new(Box::new(network), Some(editor.clone().client));

            handle.spawn(rpc_system.map_err(|_| ()));
            Ok(())
        });

    core.run(done).chain_err(|| "unable to run event loop")?;
    Ok(())
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT OR Apache-2.0
//
// Copyright (c) 2019-2020 Andre Richter <andre.o.richter@gmail.com>

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{parse_macro_input, Ident, ItemFn};

#[proc_macro_attribute]
pub fn kernel_test(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);

    let test_name = &format!("{}", f.sig.ident.to_string());
    let test_ident = Ident::new(
        &format!("{}_TEST_CONTAINER", f.sig.ident.to_string().to_uppercase()),
        Span::call_site(),
    );
    let test_code_block = f.block;

    quote!(
        #[test_case]
        const #test_ident: test_types::UnitTest = test_types::UnitTest {
            name: #test_name,
            test_func: || #test_code_block,
        };
    )
    .into()
}

#[proc_macro_attribute]
pub fn unit_test(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);
    let test_ident = f.sig.ident;
    let test_code_block = f.block;

    quote!(
        #[test]
        fn #test_ident() {
            crate::debug::unit_test_logging::setup();
            #test_code_block
        }
    )
    .into()
}
<commit_msg>Copyright for test-macros changes<commit_after>// SPDX-License-Identifier: MIT OR Apache-2.0
//
// Copyright (c) 2019-2020 Andre Richter <andre.o.richter@gmail.com>
// Copyright (c) 2020 Alister Lee <alister@dev.shortepic.com>

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{parse_macro_input, Ident, ItemFn};

#[proc_macro_attribute]
pub fn kernel_test(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);

    let test_name = &format!("{}", f.sig.ident.to_string());
    let test_ident = Ident::new(
        &format!("{}_TEST_CONTAINER", f.sig.ident.to_string().to_uppercase()),
        Span::call_site(),
    );
    let test_code_block = f.block;

    quote!(
        #[test_case]
        const #test_ident: test_types::UnitTest = test_types::UnitTest {
            name: #test_name,
            test_func: || #test_code_block,
        };
    )
    .into()
}

#[proc_macro_attribute]
pub fn unit_test(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);
    let test_ident = f.sig.ident;
    let test_code_block = f.block;

    quote!(
        #[test]
        fn #test_ident() {
            crate::debug::unit_test_logging::setup();
            #test_code_block
        }
    )
    .into()
}
<|endoftext|>


<commit_before>// Copyright (c) 2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use thiserror::Error;

use std::sync::Arc;

use crate::api::{ChromaSampling, Context, ContextInner};
use crate::cpu_features::CpuFeatureLevel;
use crate::rayon::{ThreadPool, ThreadPoolBuilder};
use crate::tiling::TilingInfo;
use crate::util::Pixel;

mod encoder;
pub use encoder::*;

mod rate;
pub use rate::Error as RateControlError;
pub use rate::{RateControlConfig, RateControlSummary};

mod speedsettings;
pub use speedsettings::*;

/// Enumeration of possible invalid configuration errors.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Error)]
#[non_exhaustive]
pub enum InvalidConfig {
  /// The width is invalid.
  #[error("invalid width {0} (expected >= 16, <= 32767)")]
  InvalidWidth(usize),
  /// The height is invalid.
  #[error("invalid height {0} (expected >= 16, <= 32767)")]
  InvalidHeight(usize),
  /// RDO lookahead frame count is invalid.
  #[error(
    "invalid rdo lookahead frames {actual} (expected <= {max} and >= {min})"
  )]
  InvalidRdoLookaheadFrames {
    /// The actual value.
    actual: usize,
    /// The maximal supported value.
    max: usize,
    /// The minimal supported value.
    min: usize,
  },
  /// Maximal keyframe interval is invalid.
  #[error("invalid max keyframe interval {actual} (expected <= {max})")]
  InvalidMaxKeyFrameInterval {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Tile columns is invalid.
  #[error("invalid tile cols {0} (expected power of 2)")]
  InvalidTileCols(usize),
  /// Tile rows is invalid.
  #[error("invalid tile rows {0} (expected power of 2)")]
  InvalidTileRows(usize),
  /// Framerate numerator is invalid.
  #[error("invalid framerate numerator {actual} (expected > 0, <= {max})")]
  InvalidFrameRateNum {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Framerate denominator is invalid.
  #[error("invalid framerate denominator {actual} (expected > 0, <= {max})")]
  InvalidFrameRateDen {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Reservoir frame delay is invalid.
  #[error("invalid reservoir frame delay {0} (expected >= 12, <= 131072)")]
  InvalidReservoirFrameDelay(i32),
  /// Reservoir frame delay is invalid.
  #[error(
    "invalid switch frame interval {0} (must only be used with low latency mode)"
  )]
  InvalidSwitchFrameInterval(u64),

  /// The rate control needs a target bitrate in order to produce results
  #[error("The rate control requires a target bitrate")]
  TargetBitrateNeeded,
}

/// Contains the encoder configuration.
#[derive(Clone, Debug, Default)]
pub struct Config {
  /// Settings which impact the produced bitstream.
  pub(crate) enc: EncoderConfig,
  /// Rate control configuration
  pub(crate) rate_control: RateControlConfig,
  /// The number of threads in the threadpool.
  pub(crate) threads: usize,
  /// Shared thread pool
  pub(crate) pool: Option<Arc<ThreadPool>>,
}

impl Config {
  /// Create a default configuration
  ///
  /// same as Default::default()
  pub fn new() -> Self {
    Config::default()
  }

  /// Set the encoder configuration
  ///
  /// EncoderConfig contains the settings impacting the
  /// codec features used in the produced bitstream.
  pub fn with_encoder_config(mut self, enc: EncoderConfig) -> Self {
    self.enc = enc;
    self
  }

  /// Set the number of workers in the threadpool
  ///
  /// The threadpool is shared across all the different parallel
  /// components in the encoder.
  pub fn with_threads(mut self, threads: usize) -> Self {
    self.threads = threads;
    self
  }

  /// Set the rate control configuration
  ///
  /// The default configuration is single pass
  pub fn with_rate_control(mut self, rate_control: RateControlConfig) -> Self {
    self.rate_control = rate_control;
    self
  }

  #[cfg(features = "unstable")]
  /// Use the provided threadpool
  pub fn with_thread_pool(mut self, pool: Arc<ThreadPool>) -> Self {
    self.threadpool = Some(pool);
    self
  }
}

fn check_tile_log2(n: usize) -> bool {
  let tile_log2 = TilingInfo::tile_log2(1, n);
  if tile_log2.is_none() {
    return false;
  }
  let tile_log2 = tile_log2.unwrap();

  ((1 << tile_log2) - n) == 0 || n == 0
}

impl Config {
  /// Creates a [`Context`] with this configuration.
  ///
  /// # Examples
  ///
  /// ```
  /// use rav1e::prelude::*;
  ///
  /// # fn main() -> Result<(), InvalidConfig> {
  /// let cfg = Config::default();
  /// let ctx: Context<u8> = cfg.new_context()?;
  /// # Ok(())
  /// # }
  /// ```
  ///
  /// [`Context`]: struct.Context.html
  pub fn new_context<T: Pixel>(&self) -> Result<Context<T>, InvalidConfig> {
    assert!(
      8 * std::mem::size_of::<T>() >= self.enc.bit_depth,
      "The Pixel u{} does not match the Config bit_depth {}",
      8 * std::mem::size_of::<T>(),
      self.enc.bit_depth
    );

    self.validate()?;

    // Because we don't have a FrameInvariants yet,
    // this is the only way to get the CpuFeatureLevel in use.
    // Since we only call this once, this shouldn't cause
    // performance issues.
    info!("CPU Feature Level: {}", CpuFeatureLevel::default());

    let pool = if let Some(ref p) = self.pool {
      p.clone()
    } else {
      let pool =
        ThreadPoolBuilder::new().num_threads(self.threads).build().unwrap();
      Arc::new(pool)
    };

    let mut config = self.enc;
    config.set_key_frame_interval(
      config.min_key_frame_interval,
      config.max_key_frame_interval,
    );

    // FIXME: inter unsupported with 4:2:2 and 4:4:4 chroma sampling
    let chroma_sampling = config.chroma_sampling;

    // FIXME: tx partition for intra not supported for chroma 422
    if chroma_sampling == ChromaSampling::Cs422 {
      config.speed_settings.rdo_tx_decision = false;
    }

    let mut inner = ContextInner::new(&config);

    if self.rate_control.emit_pass_data {
      let params = inner.rc_state.get_twopass_out_params(&inner, 0);
      inner.rc_state.init_first_pass(params.pass1_log_base_q);
    }

    if let Some(ref s) = self.rate_control.summary {
      inner.rc_state.init_second_pass();
      inner.rc_state.setup_second_pass(s);
    }

    Ok(Context { is_flushing: false, inner, pool, config })
  }

  /// Validates the configuration.
  pub fn validate(&self) -> Result<(), InvalidConfig> {
    use InvalidConfig::*;

    let config = &self.enc;

    if config.width < 16 || config.width > u16::max_value() as usize {
      return Err(InvalidWidth(config.width));
    }
    if config.height < 16 || config.height > u16::max_value() as usize {
      return Err(InvalidHeight(config.height));
    }

    if config.rdo_lookahead_frames > MAX_RDO_LOOKAHEAD_FRAMES
      || config.rdo_lookahead_frames < 1
    {
      return Err(InvalidRdoLookaheadFrames {
        actual: config.rdo_lookahead_frames,
        max: MAX_RDO_LOOKAHEAD_FRAMES,
        min: 1,
      });
    }
    if config.max_key_frame_interval > MAX_MAX_KEY_FRAME_INTERVAL {
      return Err(InvalidMaxKeyFrameInterval {
        actual: config.max_key_frame_interval,
        max: MAX_MAX_KEY_FRAME_INTERVAL,
      });
    }

    if !check_tile_log2(config.tile_cols) {
      return Err(InvalidTileCols(config.tile_cols));
    }
    if !check_tile_log2(config.tile_rows) {
      return Err(InvalidTileRows(config.tile_rows));
    }

    if config.time_base.num == 0
      || config.time_base.num > u32::max_value() as u64
    {
      return Err(InvalidFrameRateNum {
        actual: config.time_base.num,
        max: u32::max_value() as u64,
      });
    }
    if config.time_base.den == 0
      || config.time_base.den > u32::max_value() as u64
    {
      return Err(InvalidFrameRateDen {
        actual: config.time_base.den,
        max: u32::max_value() as u64,
      });
    }

    if let Some(delay) = config.reservoir_frame_delay {
      if delay < 12 || delay > 131_072 {
        return Err(InvalidReservoirFrameDelay(delay));
      }
    }

    if config.switch_frame_interval > 0 && !config.low_latency {
      return Err(InvalidSwitchFrameInterval(config.switch_frame_interval));
    }

    // TODO: add more validation
    let rc = &self.rate_control;

    if (rc.emit_pass_data || rc.summary.is_some()) && config.bitrate == 0 {
      return Err(TargetBitrateNeeded);
    }

    Ok(())
  }
}
<commit_msg>Remove a typo<commit_after>// Copyright (c) 2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use thiserror::Error;

use std::sync::Arc;

use crate::api::{ChromaSampling, Context, ContextInner};
use crate::cpu_features::CpuFeatureLevel;
use crate::rayon::{ThreadPool, ThreadPoolBuilder};
use crate::tiling::TilingInfo;
use crate::util::Pixel;

mod encoder;
pub use encoder::*;

mod rate;
pub use rate::Error as RateControlError;
pub use rate::{RateControlConfig, RateControlSummary};

mod speedsettings;
pub use speedsettings::*;

/// Enumeration of possible invalid configuration errors.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Error)]
#[non_exhaustive]
pub enum InvalidConfig {
  /// The width is invalid.
  #[error("invalid width {0} (expected >= 16, <= 32767)")]
  InvalidWidth(usize),
  /// The height is invalid.
  #[error("invalid height {0} (expected >= 16, <= 32767)")]
  InvalidHeight(usize),
  /// RDO lookahead frame count is invalid.
  #[error(
    "invalid rdo lookahead frames {actual} (expected <= {max} and >= {min})"
  )]
  InvalidRdoLookaheadFrames {
    /// The actual value.
    actual: usize,
    /// The maximal supported value.
    max: usize,
    /// The minimal supported value.
    min: usize,
  },
  /// Maximal keyframe interval is invalid.
  #[error("invalid max keyframe interval {actual} (expected <= {max})")]
  InvalidMaxKeyFrameInterval {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Tile columns is invalid.
  #[error("invalid tile cols {0} (expected power of 2)")]
  InvalidTileCols(usize),
  /// Tile rows is invalid.
  #[error("invalid tile rows {0} (expected power of 2)")]
  InvalidTileRows(usize),
  /// Framerate numerator is invalid.
  #[error("invalid framerate numerator {actual} (expected > 0, <= {max})")]
  InvalidFrameRateNum {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Framerate denominator is invalid.
  #[error("invalid framerate denominator {actual} (expected > 0, <= {max})")]
  InvalidFrameRateDen {
    /// The actual value.
    actual: u64,
    /// The maximal supported value.
    max: u64,
  },
  /// Reservoir frame delay is invalid.
  #[error("invalid reservoir frame delay {0} (expected >= 12, <= 131072)")]
  InvalidReservoirFrameDelay(i32),
  /// Reservoir frame delay is invalid.
  #[error(
    "invalid switch frame interval {0} (must only be used with low latency mode)"
  )]
  InvalidSwitchFrameInterval(u64),

  /// The rate control needs a target bitrate in order to produce results
  #[error("The rate control requires a target bitrate")]
  TargetBitrateNeeded,
}

/// Contains the encoder configuration.
#[derive(Clone, Debug, Default)]
pub struct Config {
  /// Settings which impact the produced bitstream.
  pub(crate) enc: EncoderConfig,
  /// Rate control configuration
  pub(crate) rate_control: RateControlConfig,
  /// The number of threads in the threadpool.
  pub(crate) threads: usize,
  /// Shared thread pool
  pub(crate) pool: Option<Arc<ThreadPool>>,
}

impl Config {
  /// Create a default configuration
  ///
  /// same as Default::default()
  pub fn new() -> Self {
    Config::default()
  }

  /// Set the encoder configuration
  ///
  /// EncoderConfig contains the settings impacting the
  /// codec features used in the produced bitstream.
  pub fn with_encoder_config(mut self, enc: EncoderConfig) -> Self {
    self.enc = enc;
    self
  }

  /// Set the number of workers in the threadpool
  ///
  /// The threadpool is shared across all the different parallel
  /// components in the encoder.
  pub fn with_threads(mut self, threads: usize) -> Self {
    self.threads = threads;
    self
  }

  /// Set the rate control configuration
  ///
  /// The default configuration is single pass
  pub fn with_rate_control(mut self, rate_control: RateControlConfig) -> Self {
    self.rate_control = rate_control;
    self
  }

  #[cfg(feature = "unstable")]
  /// Use the provided threadpool
  pub fn with_thread_pool(mut self, pool: Arc<ThreadPool>) -> Self {
    self.pool = Some(pool);
    self
  }
}

fn check_tile_log2(n: usize) -> bool {
  let tile_log2 = TilingInfo::tile_log2(1, n);
  if tile_log2.is_none() {
    return false;
  }
  let tile_log2 = tile_log2.unwrap();

  ((1 << tile_log2) - n) == 0 || n == 0
}

impl Config {
  /// Creates a [`Context`] with this configuration.
  ///
  /// # Examples
  ///
  /// ```
  /// use rav1e::prelude::*;
  ///
  /// # fn main() -> Result<(), InvalidConfig> {
  /// let cfg = Config::default();
  /// let ctx: Context<u8> = cfg.new_context()?;
  /// # Ok(())
  /// # }
  /// ```
  ///
  /// [`Context`]: struct.Context.html
  pub fn new_context<T: Pixel>(&self) -> Result<Context<T>, InvalidConfig> {
    assert!(
      8 * std::mem::size_of::<T>() >= self.enc.bit_depth,
      "The Pixel u{} does not match the Config bit_depth {}",
      8 * std::mem::size_of::<T>(),
      self.enc.bit_depth
    );

    self.validate()?;

    // Because we don't have a FrameInvariants yet,
    // this is the only way to get the CpuFeatureLevel in use.
    // Since we only call this once, this shouldn't cause
    // performance issues.
    info!("CPU Feature Level: {}", CpuFeatureLevel::default());

    let pool = if let Some(ref p) = self.pool {
      p.clone()
    } else {
      let pool =
        ThreadPoolBuilder::new().num_threads(self.threads).build().unwrap();
      Arc::new(pool)
    };

    let mut config = self.enc;
    config.set_key_frame_interval(
      config.min_key_frame_interval,
      config.max_key_frame_interval,
    );

    // FIXME: inter unsupported with 4:2:2 and 4:4:4 chroma sampling
    let chroma_sampling = config.chroma_sampling;

    // FIXME: tx partition for intra not supported for chroma 422
    if chroma_sampling == ChromaSampling::Cs422 {
      config.speed_settings.rdo_tx_decision = false;
    }

    let mut inner = ContextInner::new(&config);

    if self.rate_control.emit_pass_data {
      let params = inner.rc_state.get_twopass_out_params(&inner, 0);
      inner.rc_state.init_first_pass(params.pass1_log_base_q);
    }

    if let Some(ref s) = self.rate_control.summary {
      inner.rc_state.init_second_pass();
      inner.rc_state.setup_second_pass(s);
    }

    Ok(Context { is_flushing: false, inner, pool, config })
  }

  /// Validates the configuration.
  pub fn validate(&self) -> Result<(), InvalidConfig> {
    use InvalidConfig::*;

    let config = &self.enc;

    if config.width < 16 || config.width > u16::max_value() as usize {
      return Err(InvalidWidth(config.width));
    }
    if config.height < 16 || config.height > u16::max_value() as usize {
      return Err(InvalidHeight(config.height));
    }

    if config.rdo_lookahead_frames > MAX_RDO_LOOKAHEAD_FRAMES
      || config.rdo_lookahead_frames < 1
    {
      return Err(InvalidRdoLookaheadFrames {
        actual: config.rdo_lookahead_frames,
        max: MAX_RDO_LOOKAHEAD_FRAMES,
        min: 1,
      });
    }
    if config.max_key_frame_interval > MAX_MAX_KEY_FRAME_INTERVAL {
      return Err(InvalidMaxKeyFrameInterval {
        actual: config.max_key_frame_interval,
        max: MAX_MAX_KEY_FRAME_INTERVAL,
      });
    }

    if !check_tile_log2(config.tile_cols) {
      return Err(InvalidTileCols(config.tile_cols));
    }
    if !check_tile_log2(config.tile_rows) {
      return Err(InvalidTileRows(config.tile_rows));
    }

    if config.time_base.num == 0
      || config.time_base.num > u32::max_value() as u64
    {
      return Err(InvalidFrameRateNum {
        actual: config.time_base.num,
        max: u32::max_value() as u64,
      });
    }
    if config.time_base.den == 0
      || config.time_base.den > u32::max_value() as u64
    {
      return Err(InvalidFrameRateDen {
        actual: config.time_base.den,
        max: u32::max_value() as u64,
      });
    }

    if let Some(delay) = config.reservoir_frame_delay {
      if delay < 12 || delay > 131_072 {
        return Err(InvalidReservoirFrameDelay(delay));
      }
    }

    if config.switch_frame_interval > 0 && !config.low_latency {
      return Err(InvalidSwitchFrameInterval(config.switch_frame_interval));
    }

    // TODO: add more validation
    let rc = &self.rate_control;

    if (rc.emit_pass_data || rc.summary.is_some()) && config.bitrate == 0 {
      return Err(TargetBitrateNeeded);
    }

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


<commit_before>extern crate flate2;
extern crate std;
extern crate tar;
extern crate tempdir;

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Custom(&'static str),
}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Self {
        Error::Io(e)
    }
}

#[derive(Debug, Clone)]
pub struct Abs<'a> {
    repo_name: &'a str,
    abs_path: std::path::PathBuf,
}

impl<'a> Abs<'a> {
    pub fn new<P>(repo_name: &'a str, abs_path: P) -> Self
        where P: AsRef<std::path::Path>
    {
        Abs {
            repo_name: repo_name,
            abs_path: abs_path.as_ref().to_path_buf(),
        }
    }

    pub fn add<P, Q>(&self, package_dir: P, srcdest: Q) -> Result<(), Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let root = try!(tempdir::TempDir::new("guzuta-abs-root"));
        try!(self.unarchive(root.as_ref(), self.abs_path.as_path()));
        try!(self.add_srcpkg(root.as_ref(), package_dir, srcdest));
        try!(self.archive(root, self.abs_path.as_path()));
        Ok(())
    }

    pub fn remove(&self, package_name: &str) -> Result<(), Error> {
        let root = try!(tempdir::TempDir::new("guzuta-abs-root"));
        try!(self.unarchive(root.as_ref(), self.abs_path.as_path()));
        try!(std::fs::remove_dir_all(root.path().join(self.repo_name).join(package_name)));
        try!(self.archive(root, self.abs_path.as_path()));
        Ok(())
    }

    fn unarchive<P, Q>(&self, root_dir: P, abs_path: Q) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        match std::fs::File::open(abs_path) {
            Ok(file) => self.unarchive_file(root_dir, file),
            Err(e) => {
                if e.kind() != std::io::ErrorKind::NotFound {
                    Err(e)
                } else {
                    Ok(())
                }
            }
        }
    }

    fn unarchive_file<P, R>(&self, root_dir: P, abs_file: R) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              R: std::io::Read
    {
        let gz_reader = try!(flate2::read::GzDecoder::new(abs_file));
        let mut tar_reader = tar::Archive::new(gz_reader);
        try!(tar_reader.unpack(root_dir));
        Ok(())
    }

    fn add_srcpkg<P, Q, R>(&self, root_dir: P, package_dir: Q, srcdest: R) -> Result<(), Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>,
              R: AsRef<std::path::Path>
    {
        let current_dir_buf = try!(std::env::current_dir());
        let current_dir = current_dir_buf.as_path();
        let srcdest = current_dir.join(srcdest);
        let srcpkgdest = try!(tempdir::TempDir::new("guzuta-abs-srcpkgdest"));
        let builddir = try!(tempdir::TempDir::new("guzuta-abs-builddir"));
        let mut cmd = std::process::Command::new("makepkg");
        cmd.env("SRCDEST", srcdest)
            .env("SRCPKGDEST", srcpkgdest.path())
            .env("BUILDDIR", builddir.path())
            .current_dir(package_dir.as_ref())
            .arg("--source");
        info!("{:?}", cmd);
        let status = try!(cmd.status());
        if !status.success() {
            return Err(Error::Custom("makepkg --source failed"));
        }

        for entry in try!(std::fs::read_dir(srcpkgdest.path())) {
            let path = try!(entry).path();
            info!("Unarchive source package {} into {}",
                  path.display(),
                  root_dir.as_ref().display());
            try!(self.unarchive(root_dir.as_ref().join(self.repo_name), path));
            return Ok(());
        }
        return Err(Error::Custom("No source pakcage is generated"));
    }

    fn archive<P, Q>(&self, root_dir: P, abs_path: Q) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let file = try!(std::fs::File::create(abs_path.as_ref()));
        let gz_writer = flate2::write::GzEncoder::new(file, flate2::Compression::Default);
        let mut builder = tar::Builder::new(gz_writer);
        try!(self.archive_path(&mut builder, root_dir.as_ref(), root_dir.as_ref()));
        let gz_writer = try!(builder.into_inner());
        try!(gz_writer.finish());
        Ok(())
    }

    fn archive_path<W, P, Q>(&self,
                             mut builder: &mut tar::Builder<W>,
                             root_dir: P,
                             path: Q)
                             -> Result<(), std::io::Error>
        where W: std::io::Write,
              P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let path_in_archive =
            path.as_ref().strip_prefix(root_dir.as_ref()).expect("Failed to strip prefix");
        if path.as_ref().is_dir() {
            if !path_in_archive.as_os_str().is_empty() {
                let mut path_in_archive = path_in_archive.to_path_buf().into_os_string();
                path_in_archive.push("/");
                try!(builder.append_dir(path_in_archive, path.as_ref()));
            }
            for entry in try!(std::fs::read_dir(path.as_ref())) {
                try!(self.archive_path(&mut builder, root_dir.as_ref(), try!(entry).path()));
            }
            Ok(())
        } else if path.as_ref().is_file() {
            let mut file = try!(std::fs::File::open(path.as_ref()));
            try!(builder.append_file(path_in_archive, &mut file));
            Ok(())
        } else {
            // Ignore unknown file type
            Ok(())
        }
    }
}
<commit_msg>Unlink symlink to invalid source package<commit_after>extern crate flate2;
extern crate std;
extern crate tar;
extern crate tempdir;

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Custom(&'static str),
}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Self {
        Error::Io(e)
    }
}

#[derive(Debug, Clone)]
pub struct Abs<'a> {
    repo_name: &'a str,
    abs_path: std::path::PathBuf,
}

impl<'a> Abs<'a> {
    pub fn new<P>(repo_name: &'a str, abs_path: P) -> Self
        where P: AsRef<std::path::Path>
    {
        Abs {
            repo_name: repo_name,
            abs_path: abs_path.as_ref().to_path_buf(),
        }
    }

    pub fn add<P, Q>(&self, package_dir: P, srcdest: Q) -> Result<(), Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let root = try!(tempdir::TempDir::new("guzuta-abs-root"));
        try!(self.unarchive(root.as_ref(), self.abs_path.as_path()));
        try!(self.add_srcpkg(root.as_ref(), package_dir, srcdest));
        try!(self.archive(root, self.abs_path.as_path()));
        Ok(())
    }

    pub fn remove(&self, package_name: &str) -> Result<(), Error> {
        let root = try!(tempdir::TempDir::new("guzuta-abs-root"));
        try!(self.unarchive(root.as_ref(), self.abs_path.as_path()));
        try!(std::fs::remove_dir_all(root.path().join(self.repo_name).join(package_name)));
        try!(self.archive(root, self.abs_path.as_path()));
        Ok(())
    }

    fn unarchive<P, Q>(&self, root_dir: P, abs_path: Q) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        match std::fs::File::open(abs_path) {
            Ok(file) => self.unarchive_file(root_dir, file),
            Err(e) => {
                if e.kind() != std::io::ErrorKind::NotFound {
                    Err(e)
                } else {
                    Ok(())
                }
            }
        }
    }

    fn unarchive_file<P, R>(&self, root_dir: P, abs_file: R) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              R: std::io::Read
    {
        let gz_reader = try!(flate2::read::GzDecoder::new(abs_file));
        let mut tar_reader = tar::Archive::new(gz_reader);
        try!(tar_reader.unpack(root_dir));
        Ok(())
    }

    fn add_srcpkg<P, Q, R>(&self, root_dir: P, package_dir: Q, srcdest: R) -> Result<(), Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>,
              R: AsRef<std::path::Path>
    {
        let current_dir_buf = try!(std::env::current_dir());
        let current_dir = current_dir_buf.as_path();
        let srcdest = current_dir.join(srcdest);
        let srcpkgdest = try!(tempdir::TempDir::new("guzuta-abs-srcpkgdest"));
        let builddir = try!(tempdir::TempDir::new("guzuta-abs-builddir"));
        let mut cmd = std::process::Command::new("makepkg");
        cmd.env("SRCDEST", srcdest)
            .env("SRCPKGDEST", srcpkgdest.path())
            .env("BUILDDIR", builddir.path())
            .current_dir(package_dir.as_ref())
            .arg("--source");
        info!("{:?}", cmd);
        let status = try!(cmd.status());
        if !status.success() {
            return Err(Error::Custom("makepkg --source failed"));
        }

        for entry in try!(std::fs::read_dir(srcpkgdest.path())) {
            let entry = try!(entry);
            let symlink_source_package_path = package_dir.as_ref().join(entry.file_name());
            if symlink_source_package_path.read_link().is_ok() {
                info!("Unlink symlink {}", symlink_source_package_path.display());
                try!(std::fs::remove_file(symlink_source_package_path));
            }
            let path = entry.path();
            info!("Unarchive source package {} into {}",
                  path.display(),
                  root_dir.as_ref().display());
            try!(self.unarchive(root_dir.as_ref().join(self.repo_name), path));
            return Ok(());
        }
        return Err(Error::Custom("No source pakcage is generated"));
    }

    fn archive<P, Q>(&self, root_dir: P, abs_path: Q) -> Result<(), std::io::Error>
        where P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let file = try!(std::fs::File::create(abs_path.as_ref()));
        let gz_writer = flate2::write::GzEncoder::new(file, flate2::Compression::Default);
        let mut builder = tar::Builder::new(gz_writer);
        try!(self.archive_path(&mut builder, root_dir.as_ref(), root_dir.as_ref()));
        let gz_writer = try!(builder.into_inner());
        try!(gz_writer.finish());
        Ok(())
    }

    fn archive_path<W, P, Q>(&self,
                             mut builder: &mut tar::Builder<W>,
                             root_dir: P,
                             path: Q)
                             -> Result<(), std::io::Error>
        where W: std::io::Write,
              P: AsRef<std::path::Path>,
              Q: AsRef<std::path::Path>
    {
        let path_in_archive =
            path.as_ref().strip_prefix(root_dir.as_ref()).expect("Failed to strip prefix");
        if path.as_ref().is_dir() {
            if !path_in_archive.as_os_str().is_empty() {
                let mut path_in_archive = path_in_archive.to_path_buf().into_os_string();
                path_in_archive.push("/");
                try!(builder.append_dir(path_in_archive, path.as_ref()));
            }
            for entry in try!(std::fs::read_dir(path.as_ref())) {
                try!(self.archive_path(&mut builder, root_dir.as_ref(), try!(entry).path()));
            }
            Ok(())
        } else if path.as_ref().is_file() {
            let mut file = try!(std::fs::File::open(path.as_ref()));
            try!(builder.append_file(path_in_archive, &mut file));
            Ok(())
        } else {
            // Ignore unknown file type
            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>//! Check Disk usage

extern crate nix;
extern crate regex;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::cmp::max;

use structopt::StructOpt;
use regex::Regex;
use nix::sys::statvfs::vfs;
use tabin_plugins::Status;
use tabin_plugins::procfs::Mount;
use tabin_plugins::linux::bytes_to_human_size;

/// Check all mounted file systems for disk usage.
///
/// For some reason this check generally generates values that are between 1% and
/// 3% higher than `df`, even though AFAICT we're both just calling statvfs a bunch
/// of times.
#[derive(StructOpt, Deserialize, Debug)]
#[structopt(name = "check-disk (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "80")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "90")]
    crit: f64,
    #[structopt(short = "W", long = "warn-inodes", help = "Percent of inode usage to warn at",
                default_value = "80")]
    warn_inodes: f64,
    #[structopt(short = "C", long = "crit-inodes",
                help = "Percent of inode usage to go critical at", default_value = "90")]
    crit_inodes: f64,

    #[structopt(long = "pattern", name = "regex",
                help = "Only check filesystems that match this regex")]
    pattern: Option<String>,
    #[structopt(long = "exclude-pattern", name = "exclude-regex",
                help = "Only check filesystems that match this regex")]
    exclude_pattern: Option<String>,
    #[structopt(long = "type", name = "fs-type",
                help = "Only check filesystems that are of this type, e.g. \
                        ext4 or tmpfs. See 'man 8 mount' for more examples.")]
    fs_type: Option<String>,
    #[structopt(long = "exclude-type", name = "exclude-fs-type",
                help = "Do not check filesystems that are of this type.")]
    exclude_type: Option<String>,
    #[structopt(long = "info",
                help = "Print information of all known filesystems. \
                        Similar to df.")]
    info: bool,
}

fn main() {
    let args = Args::from_args();

    let mut mounts = Mount::load_all().unwrap();
    mounts.sort_by(|l, r| l.file.len().cmp(&r.file.len()));

    let status = match filter(mounts, &args) {
        Ok(ms) => do_check(&ms, &args),
        Err(e) => {
            println!("{}", e.msg);
            Status::Critical
        }
    };

    status.exit();
}

#[derive(Debug, PartialEq, Eq)]
struct ErrorMsg {
    msg: String,
}

type DiskResult<T> = Result<T, ErrorMsg>;

#[derive(Debug)]
struct MountStat {
    mount: Mount,
    stat: vfs::Statvfs,
}

fn percent(part: u64, whole: u64) -> f64 {
    100.0 - (part as f64 / whole as f64) * 100.0
}

fn maybe_regex(pattern: &Option<String>) -> DiskResult<Option<Regex>> {
    if let Some(ref pattern) = *pattern {
        let re = match Regex::new(&pattern) {
            Ok(re) => re,
            Err(e) => {
                return Err(ErrorMsg {
                    msg: format!("Unable to filter disks like {:?}: {}", pattern, e),
                })
            }
        };
        Ok(Some(re))
    } else {
        Ok(None)
    }
}

/// Convert Mounts into MountStats, applying filters from args
///
/// This:
///
/// * calls `statvfs` on every filesystem in `/proc/mounts`
/// * filters out dummy (0-block) and /proc filesystems
/// * Only shows one of any given `/dev/` filesystem's mount points (the
///   shortest, same as df)
/// * Applies the `pattern` and `type` filters
fn filter(mounts: Vec<Mount>, args: &Args) -> DiskResult<Vec<MountStat>> {
    let mut devices = HashSet::new();
    let include_regex = try!(maybe_regex(&args.pattern));
    let exclude_regex = try!(maybe_regex(&args.exclude_pattern));
    let ms = mounts.into_iter()
        .filter_map(|mount| {
            let stat = vfs::Statvfs::for_path(mount.file.as_bytes()).unwrap();
            if stat.f_blocks > 0 && !mount.file.starts_with("/proc") {
                Some(MountStat {
                    mount: mount,
                    stat: stat,
                })
            } else {
                None
            }
        })
        // If the same device is mounted multiple times, we only want one. In
        // To match df we expect these to come in in sorted order and just keep
        // the first one.
        .filter_map(|ms| {
            if ms.mount.spec.starts_with("/dev") {
                if devices.contains(&ms.mount.spec) {
                    None
                } else {
                    devices.insert(ms.mount.spec.clone());
                    Some(ms)
                }
            } else {
                Some(ms)
            }
        })
        .filter(|ms|
                if let Some(ref re) = include_regex {
                    re.is_match(&ms.mount.file)
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref re) = exclude_regex {
                    !re.is_match(&ms.mount.file)
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref vfstype) = args.fs_type {
                    ms.mount.vfstype == *vfstype
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref vfstype) = args.exclude_type {
                    ms.mount.vfstype != *vfstype
                } else {
                    true
                })
        .collect::<Vec<_>>();
    Ok(ms)
}

fn do_check(mountstats: &[MountStat], args: &Args) -> Status {
    let mut status = Status::Ok;
    for ms in mountstats {
        let pcnt = percent(ms.stat.f_bavail, ms.stat.f_blocks);
        if pcnt > args.crit {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.crit
            )
        } else if pcnt > args.warn {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.warn
            )
        }

        let ipcnt = percent(ms.stat.f_favail, ms.stat.f_files);
        if ipcnt > args.crit_inodes {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.crit_inodes
            );
        } else if ipcnt > args.warn_inodes {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.warn_inodes
            );
        }
    }
    if status == Status::Ok {
        println!(
            "OKAY: {} filesystems checked, none are above {}% disk or {}% inode usage",
            mountstats.len(),
            args.warn,
            args.warn_inodes
        );
    }

    if args.info {
        println!(
            "{:<15} {:>7} {:>5}% {:>7} {:>5}% {:<20}",
            "Filesystem", "Size", "Use", "INodes", "IUse", "Mounted on"
        );
        for ms in mountstats {
            println!(
                "{:<15} {:>7} {:>5.1}% {:>7} {:>5.1}% {:<20}",
                ms.mount.spec,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                percent(ms.stat.f_bavail, ms.stat.f_blocks),
                bytes_to_human_size(ms.stat.f_files),
                percent(ms.stat.f_favail, ms.stat.f_files),
                ms.mount.file
            );
        }
    }

    status
}

#[cfg(test)]
mod unit {
    use super::{maybe_regex, Args, ErrorMsg};
    use structopt::StructOpt;

    #[test]
    fn validate_docstring() {
        let args: Args = Args::from_iter(["arg0", "--crit", "5"].into_iter());
        assert_eq!(args.crit, 5.0);
        let args: Args = Args::from_iter(["arg0", "--pattern", "hello"].into_iter());
        assert_eq!(args.pattern.unwrap(), "hello");
    }

    #[test]
    fn check_maybe_regex() {
        if let Err(emsg) = maybe_regex(&Some("[hello".to_owned())) {
            assert_eq!(
                emsg,
                ErrorMsg {
                    msg: "Unable to filter disks like \"[hello\": \
                          Error parsing regex near \'hello\' at character offset 6: \
                          Character class was not closed before the end of the regex \
                          (missing a \']\')."
                        .to_owned(),
                }
            )
        } else {
            panic!("Should have gotten an error");
        }
    }
}
<commit_msg>Fix test that was relying on the exact error message from Regex::new<commit_after>//! Check Disk usage

extern crate nix;
extern crate regex;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::cmp::max;

use structopt::StructOpt;
use regex::Regex;
use nix::sys::statvfs::vfs;
use tabin_plugins::Status;
use tabin_plugins::procfs::Mount;
use tabin_plugins::linux::bytes_to_human_size;

/// Check all mounted file systems for disk usage.
///
/// For some reason this check generally generates values that are between 1% and
/// 3% higher than `df`, even though AFAICT we're both just calling statvfs a bunch
/// of times.
#[derive(StructOpt, Deserialize, Debug)]
#[structopt(name = "check-disk (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "80")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "90")]
    crit: f64,
    #[structopt(short = "W", long = "warn-inodes", help = "Percent of inode usage to warn at",
                default_value = "80")]
    warn_inodes: f64,
    #[structopt(short = "C", long = "crit-inodes",
                help = "Percent of inode usage to go critical at", default_value = "90")]
    crit_inodes: f64,

    #[structopt(long = "pattern", name = "regex",
                help = "Only check filesystems that match this regex")]
    pattern: Option<String>,
    #[structopt(long = "exclude-pattern", name = "exclude-regex",
                help = "Only check filesystems that match this regex")]
    exclude_pattern: Option<String>,
    #[structopt(long = "type", name = "fs-type",
                help = "Only check filesystems that are of this type, e.g. \
                        ext4 or tmpfs. See 'man 8 mount' for more examples.")]
    fs_type: Option<String>,
    #[structopt(long = "exclude-type", name = "exclude-fs-type",
                help = "Do not check filesystems that are of this type.")]
    exclude_type: Option<String>,
    #[structopt(long = "info",
                help = "Print information of all known filesystems. \
                        Similar to df.")]
    info: bool,
}

fn main() {
    let args = Args::from_args();

    let mut mounts = Mount::load_all().unwrap();
    mounts.sort_by(|l, r| l.file.len().cmp(&r.file.len()));

    let status = match filter(mounts, &args) {
        Ok(ms) => do_check(&ms, &args),
        Err(e) => {
            println!("{}", e.msg);
            Status::Critical
        }
    };

    status.exit();
}

#[derive(Debug, PartialEq, Eq)]
struct ErrorMsg {
    msg: String,
}

type DiskResult<T> = Result<T, ErrorMsg>;

#[derive(Debug)]
struct MountStat {
    mount: Mount,
    stat: vfs::Statvfs,
}

fn percent(part: u64, whole: u64) -> f64 {
    100.0 - (part as f64 / whole as f64) * 100.0
}

fn maybe_regex(pattern: &Option<String>) -> DiskResult<Option<Regex>> {
    if let Some(ref pattern) = *pattern {
        let re = match Regex::new(&pattern) {
            Ok(re) => re,
            Err(e) => {
                return Err(ErrorMsg {
                    msg: format!("Unable to filter disks like {:?}: {}", pattern, e),
                })
            }
        };
        Ok(Some(re))
    } else {
        Ok(None)
    }
}

/// Convert Mounts into MountStats, applying filters from args
///
/// This:
///
/// * calls `statvfs` on every filesystem in `/proc/mounts`
/// * filters out dummy (0-block) and /proc filesystems
/// * Only shows one of any given `/dev/` filesystem's mount points (the
///   shortest, same as df)
/// * Applies the `pattern` and `type` filters
fn filter(mounts: Vec<Mount>, args: &Args) -> DiskResult<Vec<MountStat>> {
    let mut devices = HashSet::new();
    let include_regex = try!(maybe_regex(&args.pattern));
    let exclude_regex = try!(maybe_regex(&args.exclude_pattern));
    let ms = mounts.into_iter()
        .filter_map(|mount| {
            let stat = vfs::Statvfs::for_path(mount.file.as_bytes()).unwrap();
            if stat.f_blocks > 0 && !mount.file.starts_with("/proc") {
                Some(MountStat {
                    mount: mount,
                    stat: stat,
                })
            } else {
                None
            }
        })
        // If the same device is mounted multiple times, we only want one. In
        // To match df we expect these to come in in sorted order and just keep
        // the first one.
        .filter_map(|ms| {
            if ms.mount.spec.starts_with("/dev") {
                if devices.contains(&ms.mount.spec) {
                    None
                } else {
                    devices.insert(ms.mount.spec.clone());
                    Some(ms)
                }
            } else {
                Some(ms)
            }
        })
        .filter(|ms|
                if let Some(ref re) = include_regex {
                    re.is_match(&ms.mount.file)
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref re) = exclude_regex {
                    !re.is_match(&ms.mount.file)
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref vfstype) = args.fs_type {
                    ms.mount.vfstype == *vfstype
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref vfstype) = args.exclude_type {
                    ms.mount.vfstype != *vfstype
                } else {
                    true
                })
        .collect::<Vec<_>>();
    Ok(ms)
}

fn do_check(mountstats: &[MountStat], args: &Args) -> Status {
    let mut status = Status::Ok;
    for ms in mountstats {
        let pcnt = percent(ms.stat.f_bavail, ms.stat.f_blocks);
        if pcnt > args.crit {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.crit
            )
        } else if pcnt > args.warn {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.warn
            )
        }

        let ipcnt = percent(ms.stat.f_favail, ms.stat.f_files);
        if ipcnt > args.crit_inodes {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.crit_inodes
            );
        } else if ipcnt > args.warn_inodes {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.warn_inodes
            );
        }
    }
    if status == Status::Ok {
        println!(
            "OKAY: {} filesystems checked, none are above {}% disk or {}% inode usage",
            mountstats.len(),
            args.warn,
            args.warn_inodes
        );
    }

    if args.info {
        println!(
            "{:<15} {:>7} {:>5}% {:>7} {:>5}% {:<20}",
            "Filesystem", "Size", "Use", "INodes", "IUse", "Mounted on"
        );
        for ms in mountstats {
            println!(
                "{:<15} {:>7} {:>5.1}% {:>7} {:>5.1}% {:<20}",
                ms.mount.spec,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                percent(ms.stat.f_bavail, ms.stat.f_blocks),
                bytes_to_human_size(ms.stat.f_files),
                percent(ms.stat.f_favail, ms.stat.f_files),
                ms.mount.file
            );
        }
    }

    status
}

#[cfg(test)]
mod unit {
    use super::{maybe_regex, Args};
    use structopt::StructOpt;

    #[test]
    fn validate_docstring() {
        let args: Args = Args::from_iter(["arg0", "--crit", "5"].into_iter());
        assert_eq!(args.crit, 5.0);
        let args: Args = Args::from_iter(["arg0", "--pattern", "hello"].into_iter());
        assert_eq!(args.pattern.unwrap(), "hello");
    }

    #[test]
    fn check_maybe_regex() {
        if let Err(emsg) = maybe_regex(&Some("[hello".to_owned())) {
            let expected = r#"Unable to filter disks like "[hello":"#;
            assert!(
                emsg.msg.contains(expected),
                "\nExpected something containing: {}\n\
                 But instead received         : {}",
                expected,
                emsg.msg
            );
        } else {
            panic!("Should have gotten an error");
        }
    }
}
<|endoftext|>


<commit_before>extern crate getopts;
extern crate markovpass;

use markovpass::PassphraseMarkovChain;
use std::io::Write;

fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt("n", "", "Number of passphrases to generate (default 1)", "NUM");
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3)", "LENGTH");
    opts.optopt("w", "", "Minimum word length for corpus (default 5)", "LENGTH");
    opts.optflag("h", "help", "display this help and exit");

    opts
}

fn parse_args(opts: &getopts::Options, args: &Vec<String>)
        -> Result<(String, usize, f64, usize, usize), &'static str> {
    let matches = match opts.parse(&args[1..]) {
        Ok(x) => { x },
        Err(_) => { return Err("Failed to parse arguments."); },
    };

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    macro_rules! get_num_flag {
        ($num_type:ty, $flag:expr, $default:expr) => {
            match matches.opt_str($flag) {
                Some(s) => match s.parse::<$num_type>() {
                    Ok(n) => { n },
                    Err(_) => { return Err("Failed to parse arguments."); }
                },
                None => $default,
            }
        }
    }

    let number = get_num_flag!(usize, "n", 1);
    let min_entropy = get_num_flag!(f64, "e", 60.0);
    let ngram_length = get_num_flag!(usize, "l", 3);
    let min_word_length = get_num_flag!(usize, "w", 5);

    let filename = if matches.free.is_empty() {
        "-".to_string()
    } else {
        matches.free[0].clone()
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

fn get_corpus(filename: &str) -> Result<String, std::io::Error> {
    let mut input: Box<std::io::Read> = if filename == "-" {
        Box::new(std::io::stdin())
    } else {
        let file = try!(std::fs::File::open(&filename));
        Box::new(file)
    };
    let mut corpus = String::new();
    try!(input.read_to_string(&mut corpus));

    Ok(corpus)
}

fn clean_word(word: &str, min_length: usize) -> Option<&str> {
    let word = word.trim_matches(|c: char| !c.is_alphabetic());

    if word.chars().all(|c: char| c.is_alphabetic()) && word.len() >= min_length {
        Some(word)
    } else {
        None
    }
}

fn get_ngrams(corpus: &str, ngram_length: usize, min_word_length: usize) -> Vec<String> {
    let words = corpus.split_whitespace()
      .filter_map(|word| clean_word(word, min_word_length));
    let cleaned_corpus = Some("").into_iter().chain(words).collect::<Vec<&str>>().join(" ");
    let char_count = cleaned_corpus.chars().count();
    if char_count < ngram_length { return vec![]; };
    let count = char_count - ngram_length + 1;

    let mut chars = cleaned_corpus.chars();
    let mut ngrams = Vec::with_capacity(count);
    for _ in 0..count {
        let ngram: String = chars.clone().take(ngram_length).collect();
        ngrams.push(ngram.to_lowercase());
        chars.next();
    };

    ngrams
}

fn gen_passphrases(ngrams: Vec<String>, number: usize, min_entropy: f64)
        -> Result<Vec<(String, f64)>, &'static str> {
    if ngrams.is_empty() { return Err("No NGrams found."); };
    let chain = try!(PassphraseMarkovChain::new(ngrams.iter().cloned()));
    let mut passphrases = Vec::with_capacity(number);
    for _ in 0..number {
        passphrases.push(chain.passphrase(min_entropy));
    };

    Ok(passphrases)
}

fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn write_error(program: &str, filename: &str, message: &str) {
    writeln!(&mut std::io::stderr(), "{}: {}: {}", &program, &filename, &message)
        .expect("Failed writing to stderr.");
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let program = args[0].clone();
    let opts = build_opts();

    let (filename, number, min_entropy, ngram_length, min_word_length) =
            match parse_args(&opts, &args) {
        Ok(parsed_args) => { parsed_args },
        Err(_) => {
            print_usage(&program, &opts);
            return;
        }
    };
    let corpus = match get_corpus(&filename) {
        Ok(corpus) => { corpus },
        Err(_) => {
            write_error(&program, &filename, "Failed to read input");
            return;
        },
    };
    let ngrams = get_ngrams(&corpus, ngram_length, min_word_length);
    let passphrases = match gen_passphrases(ngrams, number, min_entropy) {
        Ok(passphrases) => { passphrases },
        Err(e) => {
            write_error(&program, &filename, e);
            return;
        },
    };

    for (passphrase, entropy) in passphrases {
        println!("{} <{}>", passphrase, entropy);
    };
}
<commit_msg>Fix bug in l = 1 case<commit_after>extern crate getopts;
extern crate markovpass;

use markovpass::PassphraseMarkovChain;
use std::io::Write;

fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt("n", "", "Number of passphrases to generate (default 1)", "NUM");
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3)", "LENGTH");
    opts.optopt("w", "", "Minimum word length for corpus (default 5)", "LENGTH");
    opts.optflag("h", "help", "display this help and exit");

    opts
}

fn parse_args(opts: &getopts::Options, args: &Vec<String>)
        -> Result<(String, usize, f64, usize, usize), &'static str> {
    let matches = match opts.parse(&args[1..]) {
        Ok(x) => { x },
        Err(_) => { return Err("Failed to parse arguments."); },
    };

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    macro_rules! get_num_flag {
        ($num_type:ty, $flag:expr, $default:expr) => {
            match matches.opt_str($flag) {
                Some(s) => match s.parse::<$num_type>() {
                    Ok(n) => { n },
                    Err(_) => { return Err("Failed to parse arguments."); }
                },
                None => $default,
            }
        }
    }

    let number = get_num_flag!(usize, "n", 1);
    let min_entropy = get_num_flag!(f64, "e", 60.0);
    let ngram_length = get_num_flag!(usize, "l", 3);
    let min_word_length = get_num_flag!(usize, "w", 5);

    let filename = if matches.free.is_empty() {
        "-".to_string()
    } else {
        matches.free[0].clone()
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

fn get_corpus(filename: &str) -> Result<String, std::io::Error> {
    let mut input: Box<std::io::Read> = if filename == "-" {
        Box::new(std::io::stdin())
    } else {
        let file = try!(std::fs::File::open(&filename));
        Box::new(file)
    };
    let mut corpus = String::new();
    try!(input.read_to_string(&mut corpus));

    Ok(corpus)
}

fn clean_word(word: &str, min_length: usize) -> Option<&str> {
    let word = word.trim_matches(|c: char| !c.is_alphabetic());

    if word.chars().all(|c: char| c.is_alphabetic()) && word.len() >= min_length {
        Some(word)
    } else {
        None
    }
}

fn get_ngrams(corpus: &str, ngram_length: usize, min_word_length: usize) -> Vec<String> {
    let words = corpus.split_whitespace()
      .filter_map(|word| clean_word(word, min_word_length));
    let cleaned_corpus = Some("").into_iter().chain(words).collect::<Vec<&str>>().join(" ");
    let char_count = cleaned_corpus.chars().count();
    if char_count < ngram_length { return vec![]; };
    let count = char_count - ngram_length + 1;

    let mut chars = cleaned_corpus.chars();
    let mut ngrams = Vec::with_capacity(count);
    for _ in 0..count {
        let ngram: String = chars.clone().take(ngram_length).collect();
        ngrams.push(ngram.to_lowercase());
        chars.next();
    };

    ngrams
}

fn gen_passphrases(ngrams: Vec<String>, number: usize, min_entropy: f64)
        -> Result<Vec<(String, f64)>, &'static str> {
    if ngrams.is_empty() { return Err("No NGrams found."); };
    let chain = try!(PassphraseMarkovChain::new(ngrams.iter().cloned()));
    let mut passphrases = Vec::with_capacity(number);
    for _ in 0..number {
        passphrases.push(chain.passphrase(min_entropy));
    };

    Ok(passphrases)
}

fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn write_error(message: &str) {
    writeln!(&mut std::io::stderr(), "{}", &message).expect("Failed to write to stderr.");
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let program = args[0].clone();
    let opts = build_opts();

    let (filename, number, min_entropy, ngram_length, min_word_length) =
            match parse_args(&opts, &args) {
        Ok(parsed_args) => { parsed_args },
        Err(_) => {
            print_usage(&program, &opts);
            return;
        }
    };
    if ngram_length < 2 {
        write_error("Ngram length must be greater than one.");
    };
    let corpus = match get_corpus(&filename) {
        Ok(corpus) => { corpus },
        Err(_) => {
            write_error(&format!("{}: {}: {}", &program, &filename, "Failed to read input."));
            return;
        },
    };
    let ngrams = get_ngrams(&corpus, ngram_length, min_word_length);
    let passphrases = match gen_passphrases(ngrams, number, min_entropy) {
        Ok(passphrases) => { passphrases },
        Err(e) => {
            write_error(&format!("{}: {}: {}", &program, &filename, e));
            return;
        },
    };

    for (passphrase, entropy) in passphrases {
        println!("{} <{}>", passphrase, entropy);
    };
}
<|endoftext|>