<commit_before>//! FIXME: write short doc here

// FIXME: this modules relies on strings and AST way too much, and it should be
// rewritten (matklad 2020-05-07)
use std::{
    convert::From,
    fmt::{self, Display},
};

use hir::{Docs, Documentation, HasSource, HirDisplay};
use ra_ide_db::RootDatabase;
use ra_syntax::ast::{self, AstNode, NameOwner, VisibilityOwner};
use stdx::SepBy;

use crate::display::{generic_parameters, where_predicates};

#[derive(Debug)]
pub enum CallableKind {
    Function,
    StructConstructor,
    VariantConstructor,
    Macro,
}

/// Contains information about a function signature
#[derive(Debug)]
pub struct FunctionSignature {
    pub kind: CallableKind,
    /// Optional visibility
    pub visibility: Option<String>,
    /// Qualifiers like `async`, `unsafe`, ...
    pub qualifier: FunctionQualifier,
    /// Name of the function
    pub name: Option<String>,
    /// Documentation for the function
    pub doc: Option<Documentation>,
    /// Generic parameters
    pub generic_parameters: Vec<String>,
    /// Parameters of the function
    pub parameters: Vec<String>,
    /// Parameter names of the function
    pub parameter_names: Vec<String>,
    /// Parameter types of the function
    pub parameter_types: Vec<String>,
    /// Optional return type
    pub ret_type: Option<String>,
    /// Where predicates
    pub where_predicates: Vec<String>,
    /// Self param presence
    pub has_self_param: bool,
}

#[derive(Debug, Default)]
pub struct FunctionQualifier {
    // `async` and `const` are mutually exclusive. Do we need to enforcing it here?
    pub is_async: bool,
    pub is_const: bool,
    pub is_unsafe: bool,
    /// The string `extern ".."`
    pub extern_abi: Option<String>,
}

impl FunctionSignature {
    pub(crate) fn with_doc_opt(mut self, doc: Option<Documentation>) -> Self {
        self.doc = doc;
        self
    }

    pub(crate) fn from_hir(db: &RootDatabase, function: hir::Function) -> Self {
        let doc = function.docs(db);
        let ast_node = function.source(db).value;
        FunctionSignature::from(&ast_node).with_doc_opt(doc)
    }

    pub(crate) fn from_struct(db: &RootDatabase, st: hir::Struct) -> Option<Self> {
        let node: ast::StructDef = st.source(db).value;
        if let ast::StructKind::Record(_) = node.kind() {
            return None;
        };

        let mut params = vec![];
        let mut parameter_types = vec![];
        for field in st.fields(db).into_iter() {
            let ty = field.signature_ty(db);
            let raw_param = format!("{}", ty.display(db));

            if let Some(param_type) = raw_param.split(':').nth(1) {
                parameter_types.push(param_type[1..].to_string());
            } else {
                // useful when you have tuple struct
                parameter_types.push(raw_param.clone());
            }
            params.push(raw_param);
        }

        Some(
            FunctionSignature {
                kind: CallableKind::StructConstructor,
                visibility: node.visibility().map(|n| n.syntax().text().to_string()),
                // Do we need `const`?
                qualifier: Default::default(),
                name: node.name().map(|n| n.text().to_string()),
                ret_type: node.name().map(|n| n.text().to_string()),
                parameters: params,
                parameter_names: vec![],
                parameter_types,
                generic_parameters: generic_parameters(&node),
                where_predicates: where_predicates(&node),
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(st.docs(db)),
        )
    }

    pub(crate) fn from_enum_variant(db: &RootDatabase, variant: hir::EnumVariant) -> Option<Self> {
        let node: ast::EnumVariant = variant.source(db).value;
        match node.kind() {
            ast::StructKind::Record(_) | ast::StructKind::Unit => return None,
            _ => (),
        };

        let parent_name = variant.parent_enum(db).name(db).to_string();

        let name = format!("{}::{}", parent_name, variant.name(db));

        let mut params = vec![];
        let mut parameter_types = vec![];
        for field in variant.fields(db).into_iter() {
            let ty = field.signature_ty(db);
            let raw_param = format!("{}", ty.display(db));
            if let Some(param_type) = raw_param.split(':').nth(1) {
                parameter_types.push(param_type[1..].to_string());
            } else {
                // The unwrap_or_else is useful when you have tuple
                parameter_types.push(raw_param);
            }
            let name = field.name(db);

            params.push(format!("{}: {}", name, ty.display(db)));
        }

        Some(
            FunctionSignature {
                kind: CallableKind::VariantConstructor,
                visibility: None,
                // Do we need `const`?
                qualifier: Default::default(),
                name: Some(name),
                ret_type: None,
                parameters: params,
                parameter_names: vec![],
                parameter_types,
                generic_parameters: vec![],
                where_predicates: vec![],
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(variant.docs(db)),
        )
    }

    pub(crate) fn from_macro(db: &RootDatabase, macro_def: hir::MacroDef) -> Option<Self> {
        let node: ast::MacroCall = macro_def.source(db).value;

        let params = vec![];

        Some(
            FunctionSignature {
                kind: CallableKind::Macro,
                visibility: None,
                qualifier: Default::default(),
                name: node.name().map(|n| n.text().to_string()),
                ret_type: None,
                parameters: params,
                parameter_names: vec![],
                parameter_types: vec![],
                generic_parameters: vec![],
                where_predicates: vec![],
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(macro_def.docs(db)),
        )
    }
}

impl From<&'_ ast::FnDef> for FunctionSignature {
    fn from(node: &ast::FnDef) -> FunctionSignature {
        fn param_list(node: &ast::FnDef) -> (bool, Vec<String>, Vec<String>) {
            let mut res = vec![];
            let mut res_types = vec![];
            let mut has_self_param = false;
            if let Some(param_list) = node.param_list() {
                if let Some(self_param) = param_list.self_param() {
                    has_self_param = true;
                    let raw_param = self_param.syntax().text().to_string();

                    res_types.push(
                        raw_param.split(':').nth(1).unwrap_or_else(|| " Self")[1..].to_string(),
                    );
                    res.push(raw_param);
                }

                res.extend(param_list.params().map(|param| param.syntax().text().to_string()));
                res_types.extend(param_list.params().map(|param| {
                    let param_text = param.syntax().text().to_string();
                    match param_text.split(':').nth(1) {
                        Some(it) => it[1..].to_string(),
                        None => param_text,
                    }
                }));
            }
            (has_self_param, res, res_types)
        }

        fn param_name_list(node: &ast::FnDef) -> Vec<String> {
            let mut res = vec![];
            if let Some(param_list) = node.param_list() {
                if let Some(self_param) = param_list.self_param() {
                    res.push(self_param.syntax().text().to_string())
                }

                res.extend(
                    param_list
                        .params()
                        .map(|param| {
                            Some(
                                param
                                    .pat()?
                                    .syntax()
                                    .descendants()
                                    .find_map(ast::Name::cast)?
                                    .text()
                                    .to_string(),
                            )
                        })
                        .map(|param| param.unwrap_or_default()),
                );
            }
            res
        }

        let (has_self_param, parameters, parameter_types) = param_list(node);

        FunctionSignature {
            kind: CallableKind::Function,
            visibility: node.visibility().map(|n| n.syntax().text().to_string()),
            qualifier: FunctionQualifier {
                is_async: node.async_token().is_some(),
                is_const: node.const_token().is_some(),
                is_unsafe: node.unsafe_token().is_some(),
                extern_abi: node.abi().map(|n| n.to_string()),
            },
            name: node.name().map(|n| n.text().to_string()),
            ret_type: node
                .ret_type()
                .and_then(|r| r.type_ref())
                .map(|n| n.syntax().text().to_string()),
            parameters,
            parameter_names: param_name_list(node),
            parameter_types,
            generic_parameters: generic_parameters(node),
            where_predicates: where_predicates(node),
            // docs are processed separately
            doc: None,
            has_self_param,
        }
    }
}

impl Display for FunctionSignature {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(t) = &self.visibility {
            write!(f, "{} ", t)?;
        }

        if self.qualifier.is_async {
            write!(f, "async ")?;
        }

        if self.qualifier.is_const {
            write!(f, "const ")?;
        }

        if self.qualifier.is_unsafe {
            write!(f, "unsafe ")?;
        }

        if let Some(extern_abi) = &self.qualifier.extern_abi {
            // Keyword `extern` is included in the string.
            write!(f, "{} ", extern_abi)?;
        }

        if let Some(name) = &self.name {
            match self.kind {
                CallableKind::Function => write!(f, "fn {}", name)?,
                CallableKind::StructConstructor => write!(f, "struct {}", name)?,
                CallableKind::VariantConstructor => write!(f, "{}", name)?,
                CallableKind::Macro => write!(f, "{}!", name)?,
            }
        }

        if !self.generic_parameters.is_empty() {
            write!(f, "{}", self.generic_parameters.iter().sep_by(", ").surround_with("<", ">"))?;
        }

        write!(f, "{}", self.parameters.iter().sep_by(", ").surround_with("(", ")"))?;

        if let Some(t) = &self.ret_type {
            write!(f, " -> {}", t)?;
        }

        if !self.where_predicates.is_empty() {
            write!(f, "\nwhere {}", self.where_predicates.iter().sep_by(",\n      "))?;
        }

        Ok(())
    }
}
<commit_msg>Hot fix panic for function_signature<commit_after>//! FIXME: write short doc here

// FIXME: this modules relies on strings and AST way too much, and it should be
// rewritten (matklad 2020-05-07)
use std::{
    convert::From,
    fmt::{self, Display},
};

use hir::{Docs, Documentation, HasSource, HirDisplay};
use ra_ide_db::RootDatabase;
use ra_syntax::ast::{self, AstNode, NameOwner, VisibilityOwner};
use stdx::SepBy;

use crate::display::{generic_parameters, where_predicates};

#[derive(Debug)]
pub enum CallableKind {
    Function,
    StructConstructor,
    VariantConstructor,
    Macro,
}

/// Contains information about a function signature
#[derive(Debug)]
pub struct FunctionSignature {
    pub kind: CallableKind,
    /// Optional visibility
    pub visibility: Option<String>,
    /// Qualifiers like `async`, `unsafe`, ...
    pub qualifier: FunctionQualifier,
    /// Name of the function
    pub name: Option<String>,
    /// Documentation for the function
    pub doc: Option<Documentation>,
    /// Generic parameters
    pub generic_parameters: Vec<String>,
    /// Parameters of the function
    pub parameters: Vec<String>,
    /// Parameter names of the function
    pub parameter_names: Vec<String>,
    /// Parameter types of the function
    pub parameter_types: Vec<String>,
    /// Optional return type
    pub ret_type: Option<String>,
    /// Where predicates
    pub where_predicates: Vec<String>,
    /// Self param presence
    pub has_self_param: bool,
}

#[derive(Debug, Default)]
pub struct FunctionQualifier {
    // `async` and `const` are mutually exclusive. Do we need to enforcing it here?
    pub is_async: bool,
    pub is_const: bool,
    pub is_unsafe: bool,
    /// The string `extern ".."`
    pub extern_abi: Option<String>,
}

impl FunctionSignature {
    pub(crate) fn with_doc_opt(mut self, doc: Option<Documentation>) -> Self {
        self.doc = doc;
        self
    }

    pub(crate) fn from_hir(db: &RootDatabase, function: hir::Function) -> Self {
        let doc = function.docs(db);
        let ast_node = function.source(db).value;
        FunctionSignature::from(&ast_node).with_doc_opt(doc)
    }

    pub(crate) fn from_struct(db: &RootDatabase, st: hir::Struct) -> Option<Self> {
        let node: ast::StructDef = st.source(db).value;
        if let ast::StructKind::Record(_) = node.kind() {
            return None;
        };

        let mut params = vec![];
        let mut parameter_types = vec![];
        for field in st.fields(db).into_iter() {
            let ty = field.signature_ty(db);
            let raw_param = format!("{}", ty.display(db));

            if let Some(param_type) = raw_param.split(':').nth(1).and_then(|it| it.get(1..)) {
                parameter_types.push(param_type.to_string());
            } else {
                // useful when you have tuple struct
                parameter_types.push(raw_param.clone());
            }
            params.push(raw_param);
        }

        Some(
            FunctionSignature {
                kind: CallableKind::StructConstructor,
                visibility: node.visibility().map(|n| n.syntax().text().to_string()),
                // Do we need `const`?
                qualifier: Default::default(),
                name: node.name().map(|n| n.text().to_string()),
                ret_type: node.name().map(|n| n.text().to_string()),
                parameters: params,
                parameter_names: vec![],
                parameter_types,
                generic_parameters: generic_parameters(&node),
                where_predicates: where_predicates(&node),
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(st.docs(db)),
        )
    }

    pub(crate) fn from_enum_variant(db: &RootDatabase, variant: hir::EnumVariant) -> Option<Self> {
        let node: ast::EnumVariant = variant.source(db).value;
        match node.kind() {
            ast::StructKind::Record(_) | ast::StructKind::Unit => return None,
            _ => (),
        };

        let parent_name = variant.parent_enum(db).name(db).to_string();

        let name = format!("{}::{}", parent_name, variant.name(db));

        let mut params = vec![];
        let mut parameter_types = vec![];
        for field in variant.fields(db).into_iter() {
            let ty = field.signature_ty(db);
            let raw_param = format!("{}", ty.display(db));
            dbg!(&raw_param);
            if let Some(param_type) = raw_param.split(':').nth(1).and_then(|it| it.get(1..)) {
                parameter_types.push(param_type.to_string());
            } else {
                // The unwrap_or_else is useful when you have tuple
                parameter_types.push(raw_param);
            }
            let name = field.name(db);

            params.push(format!("{}: {}", name, ty.display(db)));
        }

        Some(
            FunctionSignature {
                kind: CallableKind::VariantConstructor,
                visibility: None,
                // Do we need `const`?
                qualifier: Default::default(),
                name: Some(name),
                ret_type: None,
                parameters: params,
                parameter_names: vec![],
                parameter_types,
                generic_parameters: vec![],
                where_predicates: vec![],
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(variant.docs(db)),
        )
    }

    pub(crate) fn from_macro(db: &RootDatabase, macro_def: hir::MacroDef) -> Option<Self> {
        let node: ast::MacroCall = macro_def.source(db).value;

        let params = vec![];

        Some(
            FunctionSignature {
                kind: CallableKind::Macro,
                visibility: None,
                qualifier: Default::default(),
                name: node.name().map(|n| n.text().to_string()),
                ret_type: None,
                parameters: params,
                parameter_names: vec![],
                parameter_types: vec![],
                generic_parameters: vec![],
                where_predicates: vec![],
                doc: None,
                has_self_param: false,
            }
            .with_doc_opt(macro_def.docs(db)),
        )
    }
}

impl From<&'_ ast::FnDef> for FunctionSignature {
    fn from(node: &ast::FnDef) -> FunctionSignature {
        fn param_list(node: &ast::FnDef) -> (bool, Vec<String>, Vec<String>) {
            let mut res = vec![];
            let mut res_types = vec![];
            let mut has_self_param = false;
            if let Some(param_list) = node.param_list() {
                if let Some(self_param) = param_list.self_param() {
                    has_self_param = true;
                    let raw_param = self_param.syntax().text().to_string();

                    res_types.push(
                        raw_param
                            .split(':')
                            .nth(1)
                            .and_then(|it| it.get(1..))
                            .unwrap_or_else(|| "Self")
                            .to_string(),
                    );
                    res.push(raw_param);
                }

                res.extend(param_list.params().map(|param| param.syntax().text().to_string()));
                res_types.extend(param_list.params().map(|param| {
                    let param_text = param.syntax().text().to_string();
                    match param_text.split(':').nth(1).and_then(|it| it.get(1..)) {
                        Some(it) => it.to_string(),
                        None => param_text,
                    }
                }));
            }
            (has_self_param, res, res_types)
        }

        fn param_name_list(node: &ast::FnDef) -> Vec<String> {
            let mut res = vec![];
            if let Some(param_list) = node.param_list() {
                if let Some(self_param) = param_list.self_param() {
                    res.push(self_param.syntax().text().to_string())
                }

                res.extend(
                    param_list
                        .params()
                        .map(|param| {
                            Some(
                                param
                                    .pat()?
                                    .syntax()
                                    .descendants()
                                    .find_map(ast::Name::cast)?
                                    .text()
                                    .to_string(),
                            )
                        })
                        .map(|param| param.unwrap_or_default()),
                );
            }
            res
        }

        let (has_self_param, parameters, parameter_types) = param_list(node);

        FunctionSignature {
            kind: CallableKind::Function,
            visibility: node.visibility().map(|n| n.syntax().text().to_string()),
            qualifier: FunctionQualifier {
                is_async: node.async_token().is_some(),
                is_const: node.const_token().is_some(),
                is_unsafe: node.unsafe_token().is_some(),
                extern_abi: node.abi().map(|n| n.to_string()),
            },
            name: node.name().map(|n| n.text().to_string()),
            ret_type: node
                .ret_type()
                .and_then(|r| r.type_ref())
                .map(|n| n.syntax().text().to_string()),
            parameters,
            parameter_names: param_name_list(node),
            parameter_types,
            generic_parameters: generic_parameters(node),
            where_predicates: where_predicates(node),
            // docs are processed separately
            doc: None,
            has_self_param,
        }
    }
}

impl Display for FunctionSignature {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(t) = &self.visibility {
            write!(f, "{} ", t)?;
        }

        if self.qualifier.is_async {
            write!(f, "async ")?;
        }

        if self.qualifier.is_const {
            write!(f, "const ")?;
        }

        if self.qualifier.is_unsafe {
            write!(f, "unsafe ")?;
        }

        if let Some(extern_abi) = &self.qualifier.extern_abi {
            // Keyword `extern` is included in the string.
            write!(f, "{} ", extern_abi)?;
        }

        if let Some(name) = &self.name {
            match self.kind {
                CallableKind::Function => write!(f, "fn {}", name)?,
                CallableKind::StructConstructor => write!(f, "struct {}", name)?,
                CallableKind::VariantConstructor => write!(f, "{}", name)?,
                CallableKind::Macro => write!(f, "{}!", name)?,
            }
        }

        if !self.generic_parameters.is_empty() {
            write!(f, "{}", self.generic_parameters.iter().sep_by(", ").surround_with("<", ">"))?;
        }

        write!(f, "{}", self.parameters.iter().sep_by(", ").surround_with("(", ")"))?;

        if let Some(t) = &self.ret_type {
            write!(f, " -> {}", t)?;
        }

        if !self.where_predicates.is_empty() {
            write!(f, "\nwhere {}", self.where_predicates.iter().sep_by(",\n      "))?;
        }

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


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::mem;

use basic::Encoding;
use data_type::AsBytes;
use errors::Result;
use util::bit_util::log2;
use util::memory::ByteBufferPtr;
use super::rle_encoding::{RleEncoder, RleDecoder};

/// A encoder for definition/repetition levels. This is a thin
/// wrapper on `RleEncoder`. Currently only support RLE encoding.
pub struct LevelEncoder {
  encoding: Encoding,
  bit_width: usize,
  rle_encoder: RleEncoder
}

impl LevelEncoder {
  pub fn new(encoding: Encoding, max_level: i16) -> Self {
    assert!(encoding == Encoding::RLE, "Currently only support RLE encoding");
    let bit_width = log2(max_level as u64) as usize;
    let max_buffer_size = RleEncoder::min_buffer_size(bit_width);
    let buffer = vec![0; max_buffer_size + mem::size_of::<i32>()];
    Self {
      encoding: encoding,
      bit_width: bit_width,
      rle_encoder: RleEncoder::new_from_buf(bit_width, buffer, mem::size_of::<i32>())
    }
  }

  #[inline]
  pub fn put(&mut self, buffer: &[i16]) -> Result<i16> {
    let mut num_encoded = 0;
    for v in buffer {
      if !self.rle_encoder.put(*v as u64)? { break; }
      num_encoded += 1;
    }
    Ok(num_encoded)
  }

  #[inline]
  pub fn consume(&mut self) -> Result<Vec<u8>> {
    let len = (self.rle_encoder.len() as i32).to_le();
    let len_bytes = len.as_bytes();
    let mut encoded_data = self.rle_encoder.consume()?;
    encoded_data[0..len_bytes.len()].copy_from_slice(len_bytes);
    Ok(encoded_data)
  }
}

/// A encoder for definition/repetition levels. This is a thin
/// wrapper on `RleDecoder`. Currently it only support RLE encoding.
pub struct LevelDecoder {
  encoding: Encoding,
  bit_width: usize,
  rle_decoder: RleDecoder
}

impl LevelDecoder {
  pub fn new(encoding: Encoding, max_level: i16) -> Self {
    assert!(encoding == Encoding::RLE, "Currently only support RLE encoding");
    let bit_width = log2(max_level as u64) as usize;
    Self {
      encoding: encoding,
      bit_width: bit_width,
      rle_decoder: RleDecoder::new(bit_width),
    }
  }

  #[inline]
  pub fn set_data(&mut self, data: ByteBufferPtr) -> usize {
    let i32_size = mem::size_of::<i32>();
    let total_bytes = i32_size + read_num_bytes!(i32, i32_size, data.as_ref()) as usize;
    self.rle_decoder.set_data(data.start_from(i32_size));
    total_bytes
  }

  #[inline]
  pub fn get(&mut self, buffer: &mut [i16]) -> Result<usize> {
    self.rle_decoder.get_batch::<i16>(buffer)
  }
}

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

  #[test]
  fn test_level_decoder() {
    let max_level = 10;
    let mut encoder = LevelEncoder::new(Encoding::RLE, max_level);
    let data = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    encoder.put(&data).expect("put() should be OK");

    let buffer = encoder.consume().expect("consume() should be OK");
    let buffer_len = buffer.len();
    let mut decoder = LevelDecoder::new(Encoding::RLE, max_level);
    assert_eq!(decoder.set_data(ByteBufferPtr::new(buffer)), buffer_len);
    let mut result: Vec<i16> = vec![0; 10];
    let num_decoded = decoder.get(&mut result).expect("get() should be OK");
    assert_eq!(num_decoded, 10);
    assert_eq!(result, data);
  }
}
<commit_msg>A few fixes in LevelEncoder<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::mem;

use basic::Encoding;
use data_type::AsBytes;
use errors::Result;
use util::bit_util::log2;
use util::memory::ByteBufferPtr;
use super::rle_encoding::{RleEncoder, RleDecoder};

/// A encoder for definition/repetition levels. This is a thin
/// wrapper on `RleEncoder`. Currently only support RLE encoding.

// TODO: track memory usage.
pub struct LevelEncoder {
  encoding: Encoding,
  bit_width: usize,
  rle_encoder: RleEncoder
}

impl LevelEncoder {
  pub fn new(encoding: Encoding, max_level: i16) -> Self {
    assert!(encoding == Encoding::RLE, "Currently only support RLE encoding");
    let bit_width = log2(max_level as u64 + 1) as usize;
    let max_buffer_size = RleEncoder::min_buffer_size(bit_width);
    let buffer = vec![0; max_buffer_size + mem::size_of::<i32>()];
    Self {
      encoding: encoding,
      bit_width: bit_width,
      rle_encoder: RleEncoder::new_from_buf(bit_width, buffer, mem::size_of::<i32>())
    }
  }

  #[inline]
  pub fn put(&mut self, buffer: &[i16]) -> Result<i16> {
    let mut num_encoded = 0;
    for v in buffer {
      if !self.rle_encoder.put(*v as u64)? { break; }
      num_encoded += 1;
    }
    Ok(num_encoded)
  }

  #[inline]
  pub fn consume(&mut self) -> Result<Vec<u8>> {
    let len = (self.rle_encoder.len() as i32).to_le();
    let len_bytes = len.as_bytes();
    let mut encoded_data = self.rle_encoder.consume()?;
    encoded_data[0..len_bytes.len()].copy_from_slice(len_bytes);
    Ok(encoded_data)
  }
}

/// A encoder for definition/repetition levels. This is a thin
/// wrapper on `RleDecoder`. Currently it only support RLE encoding.
pub struct LevelDecoder {
  encoding: Encoding,
  bit_width: usize,
  rle_decoder: RleDecoder
}

impl LevelDecoder {
  pub fn new(encoding: Encoding, max_level: i16) -> Self {
    assert!(encoding == Encoding::RLE, "Currently only support RLE encoding");
    let bit_width = log2(max_level as u64) as usize;
    Self {
      encoding: encoding,
      bit_width: bit_width,
      rle_decoder: RleDecoder::new(bit_width),
    }
  }

  #[inline]
  pub fn set_data(&mut self, data: ByteBufferPtr) -> usize {
    let i32_size = mem::size_of::<i32>();
    let data_size = read_num_bytes!(i32, i32_size, data.as_ref()) as usize;
    self.rle_decoder.set_data(data.range(i32_size, data_size));
    i32_size + data_size
  }

  #[inline]
  pub fn get(&mut self, buffer: &mut [i16]) -> Result<usize> {
    self.rle_decoder.get_batch::<i16>(buffer)
  }
}

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

  #[test]
  fn test_roundtrip() {
    let max_level = 10;
    let mut encoder = LevelEncoder::new(Encoding::RLE, max_level);
    let data = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    encoder.put(&data).expect("put() should be OK");

    let buffer = encoder.consume().expect("consume() should be OK");
    let buffer_len = buffer.len();
    let mut decoder = LevelDecoder::new(Encoding::RLE, max_level);
    assert_eq!(decoder.set_data(ByteBufferPtr::new(buffer)), buffer_len);
    let mut result: Vec<i16> = vec![0; 10];
    let num_decoded = decoder.get(&mut result).expect("get() should be OK");
    assert_eq!(num_decoded, 10);
    assert_eq!(result, data);
  }
}
<|endoftext|>


<commit_before>use lazy_static::lazy_static;
use reqwest::header::{HeaderMap, ACCEPT};
use reqwest::StatusCode;

use config::LinkChecker;
use errors::Result;

use std::collections::HashMap;
use std::error::Error;
use std::sync::{Arc, RwLock};

#[derive(Clone, Debug, PartialEq)]
pub struct LinkResult {
    pub code: Option<StatusCode>,
    /// Whether the HTTP request didn't make it to getting a HTTP code
    pub error: Option<String>,
}

impl LinkResult {
    pub fn is_valid(&self) -> bool {
        if self.error.is_some() {
            return false;
        }

        if let Some(c) = self.code {
            return c.is_success();
        }

        true
    }

    pub fn message(&self) -> String {
        if let Some(ref e) = self.error {
            return e.clone();
        }

        if let Some(c) = self.code {
            return format!("{}", c);
        }

        "Unknown error".to_string()
    }
}

lazy_static! {
    // Keep history of link checks so a rebuild doesn't have to check again
    static ref LINKS: Arc<RwLock<HashMap<String, LinkResult>>> = Arc::new(RwLock::new(HashMap::new()));
}

pub fn check_url(url: &str, config: &LinkChecker) -> LinkResult {
    {
        let guard = LINKS.read().unwrap();
        if let Some(res) = guard.get(url) {
            return res.clone();
        }
    }

    let mut headers = HeaderMap::new();
    headers.insert(ACCEPT, "text/html".parse().unwrap());
    headers.append(ACCEPT, "*/*".parse().unwrap());

    let client = reqwest::Client::new();

    let check_anchor = !config.skip_anchor_prefixes.iter().any(|prefix| url.starts_with(prefix));

    // Need to actually do the link checking
    let res = match client.get(url).headers(headers).send() {
        Ok(ref mut response) if check_anchor && has_anchor(url) => {
            match check_page_for_anchor(url, response.text()) {
                Ok(_) => LinkResult { code: Some(response.status()), error: None },
                Err(e) => LinkResult { code: None, error: Some(e.to_string()) },
            }
        }
        Ok(response) => LinkResult { code: Some(response.status()), error: None },
        Err(e) => LinkResult { code: None, error: Some(e.description().to_string()) },
    };

    LINKS.write().unwrap().insert(url.to_string(), res.clone());
    res
}

fn has_anchor(url: &str) -> bool {
    match url.find('#') {
        Some(index) => match url.get(index..=index + 1) {
            Some("#/") | Some("#!") | None => false,
            Some(_) => true,
        },
        None => false,
    }
}

fn check_page_for_anchor(url: &str, body: reqwest::Result<String>) -> Result<()> {
    let body = body.unwrap();
    let index = url.find('#').unwrap();
    let anchor = url.get(index + 1..).unwrap();
    let checks: [String; 4] = [
        format!(" id='{}'", anchor),
        format!(r#" id="{}""#, anchor),
        format!(" name='{}'", anchor),
        format!(r#" name="{}""#, anchor),
    ];

    if checks.iter().any(|check| body[..].contains(&check[..])) {
        Ok(())
    } else {
        Err(errors::Error::from(format!("Anchor `#{}` not found on page", anchor)))
    }
}

#[cfg(test)]
mod tests {
    use super::{check_page_for_anchor, check_url, has_anchor, LinkChecker, LINKS};

    #[test]
    fn can_validate_ok_links() {
        let url = "https://google.com";
        let res = check_url(url, &LinkChecker::default());
        assert!(res.is_valid());
        assert!(LINKS.read().unwrap().get(url).is_some());
        let res = check_url(url, &LinkChecker::default());
        assert!(res.is_valid());
    }

    #[test]
    fn can_fail_404_links() {
        let res = check_url("https://google.comys", &LinkChecker::default());
        assert_eq!(res.is_valid(), false);
        assert!(res.code.is_none());
        assert!(res.error.is_some());
    }

    #[test]
    fn can_validate_anchors() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = "
collect
".to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_validate_anchors_with_other_quotes() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = r#"
collect
"#.to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_validate_anchors_with_name_attr() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = r#"
collect
"#.to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_fail_when_anchor_not_found() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#me";
        let body = "
collect
".to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_err());
    }

    #[test]
    fn can_check_url_for_anchor() {
        let url = "https://doc.rust-lang.org/std/index.html#the-rust-standard-library";
        let res = has_anchor(url);
        assert_eq!(res, true);
    }

    #[test]
    fn will_return_false_when_no_anchor() {
        let url = "https://doc.rust-lang.org/std/index.html";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn will_return_false_when_has_router_url() {
        let url = "https://doc.rust-lang.org/#/std";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn will_return_false_when_has_router_url_alt() {
        let url = "https://doc.rust-lang.org/#!/std";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn skip_anchor_prefixes() {
        let config = LinkChecker {
            skip_prefixes: vec![],
            skip_anchor_prefixes: vec!["https://github.com/rust-lang/rust/blob/".to_owned()],
        };

        // anchor check is ignored because the url matches the prefix
        let permalink = "https://github.com/rust-lang/rust/blob/c772948b687488a087356cb91432425662e034b9/src/librustc_back/target/mod.rs#L194-L214";
        assert!(check_url(&permalink, &config).is_valid());

        // other anchors are checked
        let glossary = "https://help.github.com/en/articles/github-glossary#blame";
        assert!(check_url(&glossary, &config).is_valid());

        let glossary_invalid =
            "https://help.github.com/en/articles/github-glossary#anchor-does-not-exist";
        assert_eq!(check_url(&glossary_invalid, &config).is_valid(), false);
    }
}
<commit_msg>link_checker: Handle non-success status codes (#897)<commit_after>use lazy_static::lazy_static;
use reqwest::header::{HeaderMap, ACCEPT};
use reqwest::StatusCode;

use config::LinkChecker;
use errors::Result;

use std::collections::HashMap;
use std::error::Error;
use std::sync::{Arc, RwLock};

#[derive(Clone, Debug, PartialEq)]
pub struct LinkResult {
    pub code: Option<StatusCode>,
    /// Whether the HTTP request didn't make it to getting a HTTP code
    pub error: Option<String>,
}

impl LinkResult {
    pub fn is_valid(&self) -> bool {
        if self.error.is_some() {
            return false;
        }

        if let Some(c) = self.code {
            return c.is_success();
        }

        true
    }

    pub fn message(&self) -> String {
        if let Some(ref e) = self.error {
            return e.clone();
        }

        if let Some(c) = self.code {
            return format!("{}", c);
        }

        "Unknown error".to_string()
    }
}

lazy_static! {
    // Keep history of link checks so a rebuild doesn't have to check again
    static ref LINKS: Arc<RwLock<HashMap<String, LinkResult>>> = Arc::new(RwLock::new(HashMap::new()));
}

pub fn check_url(url: &str, config: &LinkChecker) -> LinkResult {
    {
        let guard = LINKS.read().unwrap();
        if let Some(res) = guard.get(url) {
            return res.clone();
        }
    }

    let mut headers = HeaderMap::new();
    headers.insert(ACCEPT, "text/html".parse().unwrap());
    headers.append(ACCEPT, "*/*".parse().unwrap());

    let client = reqwest::Client::new();

    let check_anchor = !config.skip_anchor_prefixes.iter().any(|prefix| url.starts_with(prefix));

    // Need to actually do the link checking
    let res = match client.get(url).headers(headers).send() {
        Ok(ref mut response) if check_anchor && has_anchor(url) => {
            match check_page_for_anchor(url, response.text()) {
                Ok(_) => LinkResult { code: Some(response.status()), error: None },
                Err(e) => LinkResult { code: None, error: Some(e.to_string()) },
            }
        }
        Ok(response) => {
            if response.status().is_success() {
                LinkResult { code: Some(response.status()), error: None }
            } else {
                let error_string = if response.status().is_informational() {
                    String::from(format!(
                        "Informational status code ({}) received",
                        response.status()
                    ))
                } else if response.status().is_redirection() {
                    String::from(format!(
                        "Redirection status code ({}) received",
                        response.status()
                    ))
                } else if response.status().is_client_error() {
                    String::from(format!(
                        "Client error status code ({}) received",
                        response.status()
                    ))
                } else if response.status().is_server_error() {
                    String::from(format!(
                        "Server error status code ({}) received",
                        response.status()
                    ))
                } else {
                    String::from("Non-success status code received")
                };

                LinkResult { code: None, error: Some(error_string) }
            }
        }
        Err(e) => LinkResult { code: None, error: Some(e.description().to_string()) },
    };

    LINKS.write().unwrap().insert(url.to_string(), res.clone());
    res
}

fn has_anchor(url: &str) -> bool {
    match url.find('#') {
        Some(index) => match url.get(index..=index + 1) {
            Some("#/") | Some("#!") | None => false,
            Some(_) => true,
        },
        None => false,
    }
}

fn check_page_for_anchor(url: &str, body: reqwest::Result<String>) -> Result<()> {
    let body = body.unwrap();
    let index = url.find('#').unwrap();
    let anchor = url.get(index + 1..).unwrap();
    let checks: [String; 4] = [
        format!(" id='{}'", anchor),
        format!(r#" id="{}""#, anchor),
        format!(" name='{}'", anchor),
        format!(r#" name="{}""#, anchor),
    ];

    if checks.iter().any(|check| body[..].contains(&check[..])) {
        Ok(())
    } else {
        Err(errors::Error::from(format!("Anchor `#{}` not found on page", anchor)))
    }
}

#[cfg(test)]
mod tests {
    use super::{check_page_for_anchor, check_url, has_anchor, LinkChecker, LINKS};

    #[test]
    fn can_validate_ok_links() {
        let url = "https://google.com";
        let res = check_url(url, &LinkChecker::default());
        assert!(res.is_valid());
        assert!(LINKS.read().unwrap().get(url).is_some());
        let res = check_url(url, &LinkChecker::default());
        assert!(res.is_valid());
    }

    #[test]
    fn can_fail_404_links() {
        let res = check_url("https://google.comys", &LinkChecker::default());
        assert_eq!(res.is_valid(), false);
        assert!(res.code.is_none());
        assert!(res.error.is_some());
    }

    #[test]
    fn can_validate_anchors() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = "
collect
".to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_validate_anchors_with_other_quotes() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = r#"
collect
"#.to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_validate_anchors_with_name_attr() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.collect";
        let body = r#"
collect
"#.to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_ok());
    }

    #[test]
    fn can_fail_when_anchor_not_found() {
        let url = "https://doc.rust-lang.org/std/iter/trait.Iterator.html#me";
        let body = "
collect
".to_string();
        let res = check_page_for_anchor(url, Ok(body));
        assert!(res.is_err());
    }

    #[test]
    fn can_check_url_for_anchor() {
        let url = "https://doc.rust-lang.org/std/index.html#the-rust-standard-library";
        let res = has_anchor(url);
        assert_eq!(res, true);
    }

    #[test]
    fn will_return_false_when_no_anchor() {
        let url = "https://doc.rust-lang.org/std/index.html";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn will_return_false_when_has_router_url() {
        let url = "https://doc.rust-lang.org/#/std";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn will_return_false_when_has_router_url_alt() {
        let url = "https://doc.rust-lang.org/#!/std";
        let res = has_anchor(url);
        assert_eq!(res, false);
    }

    #[test]
    fn skip_anchor_prefixes() {
        let config = LinkChecker {
            skip_prefixes: vec![],
            skip_anchor_prefixes: vec!["https://github.com/rust-lang/rust/blob/".to_owned()],
        };

        // anchor check is ignored because the url matches the prefix
        let permalink = "https://github.com/rust-lang/rust/blob/c772948b687488a087356cb91432425662e034b9/src/librustc_back/target/mod.rs#L194-L214";
        assert!(check_url(&permalink, &config).is_valid());

        // other anchors are checked
        let glossary = "https://help.github.com/en/articles/github-glossary#blame";
        assert!(check_url(&glossary, &config).is_valid());

        let glossary_invalid =
            "https://help.github.com/en/articles/github-glossary#anchor-does-not-exist";
        assert_eq!(check_url(&glossary_invalid, &config).is_valid(), false);
    }
}
<|endoftext|>


<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::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketConstants;
use dom::bindings::error::Fallible;
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{Temporary, JSRef};
use dom::bindings::trace::JSTraceable;
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::eventtarget::{EventTarget, EventTargetHelpers, WebSocketTypeId};
use servo_net::resource_task::{Load, LoadData, LoadResponse};
use servo_util::str::DOMString;
use std::ascii::IntoBytes;
use std::comm::channel;
use url::Url;

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: DOMString,
    response: LoadResponse,
    state: u16
}

impl WebSocket {
    pub fn new_inherited(url: DOMString, resp: LoadResponse, state_value: u16) -> WebSocket {
        WebSocket {
            eventtarget: EventTarget::new_inherited(WebSocketTypeId),
            url: url,
            response: resp,
            state: state_value
        }
    }

    pub fn new(global: &GlobalRef, url: DOMString) -> Temporary<WebSocket> {
        let resource_task = global.resource_task();
        let ws_url = Url::parse(url.as_slice()).unwrap();
        let(start_chan, start_port) = channel();
        resource_task.send(Load(LoadData::new(ws_url), start_chan));
        let resp = start_port.recv();
        reflect_dom_object(box WebSocket::new_inherited(url, resp, WebSocketConstants::OPEN),
                           global,
                           WebSocketBinding::Wrap)
    }

    pub fn Constructor(global: &GlobalRef, url: DOMString) -> Fallible<Temporary<WebSocket>> {
        Ok(WebSocket::new(global, url))
    }

}

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

impl<'a> WebSocketMethods for JSRef<'a, WebSocket> {
    event_handler!(open, GetOnopen, SetOnopen)
    event_handler!(error, GetOnerror, SetOnerror)
    event_handler!(close, GetOnclose, SetOnclose)
    event_handler!(message, GetOnmessage, SetOnmessage)

    fn Send(self, message: DOMString){
        let message_u8: Vec<u8>=message.to_string().into_bytes();
        assert!(message.len() <= 125);
        let mut payload: Vec<u8> = Vec::with_capacity(2 + message_u8.len());
//We are sending a single framed unmasked text message. Referring to http://tools.ietf.org/html/rfc6455#section-5.7.
// 10000001 this indicates FIN bit=1 and the opcode is 0001 which means it is a text frame and the decimal equivalent is 129

        payload.push(129u8);
        payload.push(message.len() as u8);
        payload.push_all(message_u8.as_slice());
    }

    fn Url(self) -> DOMString {
        self.url.clone()
    }
}
<commit_msg> added const value 0x81<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::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketConstants;
use dom::bindings::error::Fallible;
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{Temporary, JSRef};
use dom::bindings::trace::JSTraceable;
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::eventtarget::{EventTarget, EventTargetHelpers, WebSocketTypeId};
use servo_net::resource_task::{Load, LoadData, LoadResponse};
use servo_util::str::DOMString;
use std::ascii::IntoBytes;
use std::comm::channel;
use url::Url;

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: DOMString,
    response: LoadResponse,
    state: u16
}

impl WebSocket {
    pub fn new_inherited(url: DOMString, resp: LoadResponse, state_value: u16) -> WebSocket {
        WebSocket {
            eventtarget: EventTarget::new_inherited(WebSocketTypeId),
            url: url,
            response: resp,
            state: state_value
        }
    }

    pub fn new(global: &GlobalRef, url: DOMString) -> Temporary<WebSocket> {
        let resource_task = global.resource_task();
        let ws_url = Url::parse(url.as_slice()).unwrap();
        let(start_chan, start_port) = channel();
        resource_task.send(Load(LoadData::new(ws_url), start_chan));
        let resp = start_port.recv();
        reflect_dom_object(box WebSocket::new_inherited(url, resp, WebSocketConstants::OPEN),
                           global,
                           WebSocketBinding::Wrap)
    }

    pub fn Constructor(global: &GlobalRef, url: DOMString) -> Fallible<Temporary<WebSocket>> {
        Ok(WebSocket::new(global, url))
    }

}

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

impl<'a> WebSocketMethods for JSRef<'a, WebSocket> {
    event_handler!(open, GetOnopen, SetOnopen)
    event_handler!(error, GetOnerror, SetOnerror)
    event_handler!(close, GetOnclose, SetOnclose)
    event_handler!(message, GetOnmessage, SetOnmessage)

    fn Send(self, message: DOMString){
        let message_u8: Vec<u8>=message.to_string().into_bytes();
        assert!(message.len() <= 125);
        let mut payload: Vec<u8> = Vec::with_capacity(2 + message_u8.len());
//We are sending a single framed unmasked text message. Referring to http://tools.ietf.org/html/rfc6455#section-5.7.
// 10000001 this indicates FIN bit=1 and the opcode is 0001 which means it is a text frame and the decimal equivalent is 129
        const ENTRY1: u8 = 0x81;
        payload.push(ENTRY1);
        payload.push(message.len() as u8);
        payload.push_all(message_u8.as_slice());
    }

    fn Url(self) -> DOMString {
        self.url.clone()
    }
}
<|endoftext|>


<commit_before>use rand::{Rng, SeedableRng, StdRng};
use std::cmp::max;
use std::usize;
use std::collections::{BTreeSet, HashSet};
use petgraph::Graph;
use petgraph::prelude::NodeIndex;
use generators::Gen;
use resources::AstronomicalNamesResource as anr;

/// Basic non deterministic name generator for generating random strings which
/// are similar to the trained data provided.
pub struct NameGen {
    rng: StdRng,
    generated: HashSet<String>,
    graph: Graph<char, f64>,
    start: NodeIndex,
    end: NodeIndex,
}

impl Gen for NameGen {
    type GenItem = String;
    type TrainData = anr;

    /// Creates a new NameGen with the given seed
    fn new(seed: u32) -> NameGen {

        // Create and initialize random generator using seed
        let new_seed: &[_] = &[seed.clone() as usize];
        let rng: StdRng = SeedableRng::from_seed(new_seed);
        let mut graph = Graph::<char, f64>::new();
        let start = graph.add_node('<');
        let end = graph.add_node('>');
        let generated = HashSet::<String>::new();

        NameGen {
            rng,
            generated,
            graph,
            start,
            end,
        }
    }

    /// Trains the underlying model using the given AstronomicalNamesResource
    fn train(&mut self, data: &anr) {

        let depth = data.names.iter().fold(0, |acc, ref s| max(acc, s.len()));

        // Instansiate layers, number of layers is equal to the longest training
        // string
        let mut layers = Vec::<BTreeSet<NodeIndex>>::new();
        for _ in 0..depth {
            layers.push(BTreeSet::<NodeIndex>::new());
        }

        // Add edges between all characters following each other at every
        // position producing a forward connected graph
        for name in &data.names {
            let chars = name.chars();
            let mut prev = self.start;

            for (index, chr) in chars.enumerate() {
                let node = match layers[index].iter().find(|&&node| {
                    *self.graph.node_weight(node).unwrap() == chr
                }) {
                    Some(&node) => node,
                    _ => self.graph.add_node(chr),
                };
                layers[index].insert(node.clone());

                self.graph.update_edge(prev, node, 0.0);
                prev = node;
            }
            // Add connection to end from last character
            self.graph.update_edge(prev, self.end, 0.0);
        }
        info!(
            "Model has been trained, is using {} layers, {} nodes and {} edges",
            depth,
            self.graph.node_count(),
            self.graph.edge_count()
        );
    }

    /// Attempts to generate a new name from the model.
    /// This name is guaranteed to exist in the training set or to have been
    /// previously generated.
    /// Attempts N number of tries, if no unique name could be found it will
    /// return None
    fn generate(&mut self) -> Option<String> {

        // Non deterministicly generate a new string from the model,
        // Note: This may produce an exisiting string in the training set or
        // previously generated set
        fn generate_attempt(
            graph: &Graph<char, f64>,
            start: &NodeIndex,
            end: &NodeIndex,
            rng: &mut StdRng,
        ) -> String {
            let mut final_string = String::new();
            let mut current_node = start.clone();

            // Traverse until we hit end
            while current_node != end.clone() {

                // Step to random neighbor in next layer for which it exists
                // an edge
                let neighbors = graph.neighbors(current_node).collect::<Vec<NodeIndex>>();
                let next_node = rng.choose(neighbors.as_slice()).unwrap();
                final_string.push(*graph.node_weight(current_node).unwrap());
                current_node = *next_node;
            }
            // Remove start node character
            final_string.remove(0);
            final_string
        };

        // Check if a name is valid according to constraints
        // TODO: Extract this to a seperate method, should also probably be
        // based on configuration entries
        let is_valid_name = |name: &String| name.contains(" ") || name.len() < 9;

        // Attempt N number of attempts retuning none if no unique string was
        // generated which fullfils the criteria
        let gen_num_attempts = 27;
        for _ in 0..gen_num_attempts {
            let name = generate_attempt(&self.graph, &self.start, &self.end, &mut self.rng);
            if is_valid_name(&name) && !self.generated.contains(&name) {
                self.generated.insert(name.clone());
                return Some(name);
            }
        }
        info!("Unsuccessful generation used {} tries", gen_num_attempts);
        None

    }
}

#[cfg(test)]
mod names_test {
    use super::*;
    extern crate env_logger;
    use resources::{AstronomicalNamesResource, ResourceHandler};

    #[test]
    // All genrated names must be unique
    fn test_generate_unique() {
        let _ = env_logger::init();

        let mut gen = NameGen::new(0);
        let factory = ResourceHandler::new();
        let res = factory
            .fetch_resource::<AstronomicalNamesResource>()
            .unwrap();
        gen.train(&res);

        let mut names = HashSet::<String>::new();
        for _ in 0..30 {
            names.insert(gen.generate().unwrap());
        }

        assert_eq!(names.len(), 30);
    }
}
<commit_msg>Add suffix generation to NameGen<commit_after>use rand::{Rng, SeedableRng, StdRng};
use std::cmp::max;
use std::usize;
use std::collections::{BTreeSet, HashSet};
use petgraph::Graph;
use petgraph::prelude::NodeIndex;
use generators::Gen;
use resources::AstronomicalNamesResource as anr;

/// Basic non deterministic name generator for generating random strings which
/// are similar to the trained data provided.
pub struct NameGen {
    rng: StdRng,
    generated: HashSet<String>,
    graph: Graph<char, f64>,
    start: NodeIndex,
    end: NodeIndex,
    suffixes: Vec<String>
}

impl Gen for NameGen {
    type GenItem = String;
    type TrainData = anr;

    /// Creates a new NameGen with the given seed
    fn new(seed: u32) -> NameGen {

        // Create and initialize random generator using seed
        let new_seed: &[_] = &[seed.clone() as usize];
        let rng: StdRng = SeedableRng::from_seed(new_seed);
        let mut graph = Graph::<char, f64>::new();
        let start = graph.add_node('<');
        let end = graph.add_node('>');
        let generated = HashSet::<String>::new();
        let suffixes = Vec::<String>::new();

        NameGen {
            rng,
            generated,
            graph,
            start,
            end,
            suffixes
        }
    }

    /// Trains the underlying model using the given AstronomicalNamesResource
    fn train(&mut self, data: &anr) {

        let depth = data.names.iter().fold(0, |acc, ref s| max(acc, s.len()));

        // Instansiate layers, number of layers is equal to the longest training
        // string
        let mut layers = Vec::<BTreeSet<NodeIndex>>::new();
        for _ in 0..depth {
            layers.push(BTreeSet::<NodeIndex>::new());
        }

        // Add edges between all characters following each other at every
        // position producing a forward connected graph
        for name in &data.names {
            let chars = name.chars();
            let mut prev = self.start;

            for (index, chr) in chars.enumerate() {
                let node = match layers[index].iter().find(|&&node| {
                    *self.graph.node_weight(node).unwrap() == chr
                }) {
                    Some(&node) => node,
                    _ => self.graph.add_node(chr),
                };
                layers[index].insert(node.clone());

                self.graph.update_edge(prev, node, 0.0);
                prev = node;
            }
            // Add connection to end from last character
            self.graph.update_edge(prev, self.end, 0.0);
        }
        info!(
            "Model has been trained, is using {} layers, {} nodes and {} edges",
            depth,
            self.graph.node_count(),
            self.graph.edge_count()
        );

        // Load suffixes
        self.suffixes.extend_from_slice(&data.greek[..]);
        self.suffixes.extend_from_slice(&data.decorators[..]);
    }

    /// Attempts to generate a new name from the model.
    /// This name is guaranteed to exist in the training set or to have been
    /// previously generated.
    /// Attempts N number of tries, if no unique name could be found it will
    /// return None
    fn generate(&mut self) -> Option<String> {

        // Non deterministicly generate a new string from the model,
        // Note: This may produce an exisiting string in the training set or
        // previously generated set
        fn generate_attempt(
            graph: &Graph<char, f64>,
            start: &NodeIndex,
            end: &NodeIndex,
            rng: &mut StdRng,
        ) -> String {
            let mut final_string = String::new();
            let mut current_node = start.clone();

            // Traverse until we hit end
            while current_node != end.clone() {

                // Step to random neighbor in next layer for which it exists
                // an edge
                let neighbors = graph.neighbors(current_node).collect::<Vec<NodeIndex>>();
                let next_node = rng.choose(neighbors.as_slice()).unwrap();
                final_string.push(*graph.node_weight(current_node).unwrap());
                current_node = *next_node;
            }
            // Remove start node character
            final_string.remove(0);
            final_string
        };

        // Randomly generate a suffix uniformlly from training data, has a high
        // chance of returning None
        // TODO: Load threshold from config
        fn get_suffix(rng: &mut StdRng, suffixes: &Vec<String>) -> Option<String> {
            let suffix_chance = 0.1;
            match rng.next_f64() {
                 x if x < suffix_chance => Some(rng.choose(&suffixes[..]).unwrap().clone()),
                 _ => None
            }
        }

        // Check if a name is valid according to constraints
        // TODO: Extract this to a seperate method, should also probably be
        // based on configuration entries
        let is_valid_name = |name: &String| name.contains(" ") || name.len() < 9;

        // Attempt N number of attempts retuning none if no unique string was
        // generated which fullfils the criteria
        let gen_num_attempts = 27;
        for _ in 0..gen_num_attempts {
            let name = generate_attempt(&self.graph, &self.start, &self.end, &mut self.rng);
            if is_valid_name(&name) && !self.generated.contains(&name) {
                self.generated.insert(name.clone());
                return match get_suffix(&mut self.rng, &self.suffixes) {
                     Some(suffix) => Some(format!("{} {}", name, &suffix)),
                     _ => Some(name)
                }
            }
        }
        info!("Unsuccessful generation used {} tries", gen_num_attempts);
        None

    }
}

#[cfg(test)]
mod names_test {
    use super::*;
    extern crate env_logger;
    use resources::{AstronomicalNamesResource, ResourceHandler};

    #[test]
    // All genrated names must be unique
    fn test_generate_unique() {
        let _ = env_logger::init();

        let mut gen = NameGen::new(0);
        let factory = ResourceHandler::new();
        let res = factory
            .fetch_resource::<AstronomicalNamesResource>()
            .unwrap();
        gen.train(&res);

        let mut names = HashSet::<String>::new();
        for _ in 0..30 {
            names.insert(gen.generate().unwrap());
        }

        assert_eq!(names.len(), 30);
    }
}
<|endoftext|>


<commit_before>use std;

use std::cmp::PartialOrd;

extern crate num;
use self::num::traits::{Zero, One, FromPrimitive, ToPrimitive, Float};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate image;
use self::image::{ImageBuffer, Luma, DynamicImage};

quick_error! {
    #[derive(Debug)]
    pub enum ImageSaveError {
        /// an error has occured when opening the file
        Io(err: std::io::Error) {}
        /// an error has occured when writing the image into the file
        Image(err: image::ImageError) {}
    }
}

pub trait ToImage {
    /// # Panics
    /// panics unless all the values in `self` are between
    /// `0` (inclusive) and `1` (inclusive)
    fn to_image(&self) -> DynamicImage;
    /// # Panics
    /// panics unless all the values in `self` are between
    /// `0` (inclusive) and `1` (inclusive)
    fn save_to_png(&self, filename: &str) -> Result<(), ImageSaveError>;
}

impl<T: Zero + One + FromPrimitive + ToPrimitive + PartialOrd + Float> ToImage for DMat<T> {
    fn to_image(&self) -> DynamicImage {
        let mut image_buffer = ImageBuffer::new(self.ncols() as u32, self.nrows() as u32);
        for (x, y, pixel) in image_buffer.enumerate_pixels_mut() {
            let value = self[(y as usize, x as usize)];
            assert!(T::zero() <= value);
            assert!(value <= T::one());
            let max = T::from_u8(std::u8::MAX).unwrap();
            let byte = (value * max).round().to_u8().unwrap();
            *pixel = Luma([byte]);
        }
        DynamicImage::ImageLuma8(image_buffer)
    }

    fn save_to_png(&self, filename: &str) -> Result<(), ImageSaveError> {
        let image = self.to_image();
        let path = std::path::Path::new(filename);
        let mut file = try!(std::fs::File::create(&path).map_err(ImageSaveError::Io));
        image.save(&mut file, image::PNG).map_err(ImageSaveError::Image)
    }
}
<commit_msg>docstring<commit_after>use std;

use std::cmp::PartialOrd;

extern crate num;
use self::num::traits::{Zero, One, FromPrimitive, ToPrimitive, Float};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate image;
use self::image::{ImageBuffer, Luma, DynamicImage};

quick_error! {
    #[derive(Debug)]
    pub enum ImageSaveError {
        /// an error has occured when opening the file
        Io(err: std::io::Error) {}
        /// an error has occured when writing the image into the file
        Image(err: image::ImageError) {}
    }
}

/// make images from things containing brightness information
pub trait ToImage {
    /// # Panics
    /// panics unless all the values in `self` are between
    /// `0` (inclusive) and `1` (inclusive)
    fn to_image(&self) -> DynamicImage;
    /// # Panics
    /// panics unless all the values in `self` are between
    /// `0` (inclusive) and `1` (inclusive)
    fn save_to_png(&self, filename: &str) -> Result<(), ImageSaveError>;
}

impl<T: Zero + One + FromPrimitive + ToPrimitive + PartialOrd + Float> ToImage for DMat<T> {
    fn to_image(&self) -> DynamicImage {
        let mut image_buffer = ImageBuffer::new(self.ncols() as u32, self.nrows() as u32);
        for (x, y, pixel) in image_buffer.enumerate_pixels_mut() {
            let value = self[(y as usize, x as usize)];
            assert!(T::zero() <= value);
            assert!(value <= T::one());
            let max = T::from_u8(std::u8::MAX).unwrap();
            let byte = (value * max).round().to_u8().unwrap();
            *pixel = Luma([byte]);
        }
        DynamicImage::ImageLuma8(image_buffer)
    }

    fn save_to_png(&self, filename: &str) -> Result<(), ImageSaveError> {
        let image = self.to_image();
        let path = std::path::Path::new(filename);
        let mut file = try!(std::fs::File::create(&path).map_err(ImageSaveError::Io));
        image.save(&mut file, image::PNG).map_err(ImageSaveError::Image)
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

macro_rules! want_macro {
    ($t:expr, $n:ident, $isf:ident, $asf:ident) => {
        /// Helper that returns an Option<Value::$t>.
        /// You can optionally pass a JSON pointer to retrieve a
        /// nested key.
        #[macro_export]
        macro_rules! $n {
            ($v:expr) => (if $v.$isf() {
                Some($v.$asf().unwrap())
            } else {
                None
            });

            ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
                $n!(v)
            } else {
                None
            });
        }
    }
}

want_macro!("Null", wantnull, is_null, as_null);
want_macro!("Bool", wantbool, is_boolean, as_bool);
want_macro!("I64", wanti64, is_i64, as_i64);
want_macro!("U64", wantu64, is_u64, as_u64);
want_macro!("F64", wantf64, is_f64, as_f64);
want_macro!("String", wantstr, is_string, as_str);
want_macro!("Array", wantarray, is_array, as_array);
want_macro!("Object", wantobj, is_object, as_object);

// XXX We can't auto-generate macros due to the eager expansion of
// `$crate`, which causes any external libs to look for ::error::...
// instead of inapi::error::...
//
// macro_rules! need_macro {
//     ($t:expr, $n:ident, $isf:ident, $asf:ident) => {
//         /// Helper that returns a Result<Value::$t>.
//         /// You can optionally pass a JSON pointer to retrieve a
//         /// nested key.
//         #[macro_export]
//         macro_rules! $n {
//             ($v:expr) => (if $v.$isf() {
//                 Ok($v.$asf().unwrap())
//             } else {
//                 Err($crate::Error::Generic(format!("Value is not $t")))
//             });
//
//             ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
//                 $n!(v)
//             } else {
//                 Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
//             });
//         }
//     }
// }
//
// need_macro!("Null", neednull, is_null, as_null);
// need_macro!("Bool", needbool, is_boolean, as_bool);
// need_macro!("I64", needi64, is_i64, as_i64);
// need_macro!("U64", needu64, is_u64, as_u64);
// need_macro!("F64", needf64, is_f64, as_f64);
// need_macro!("String", needstr, is_string, as_str);
// need_macro!("Array", needarray, is_array, as_array);
// need_macro!("Object", needobj, is_object, as_object);

/// Helper that returns a Result<Value::Null>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! neednull {
    ($v:expr) => (if $v.is_null() {
        Ok($v.as_null().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not null")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        neednull!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::Bool>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needbool {
    ($v:expr) => (if $v.is_boolean() {
        Ok($v.as_bool().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not boolean")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needbool!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::I64>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needi64 {
    ($v:expr) => (if $v.is_i64() {
        Ok($v.as_i64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not i64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needi64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::U64>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needu64 {
    ($v:expr) => (if $v.is_u64() {
        Ok($v.as_u64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not u64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needu64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::F64>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needf64 {
    ($v:expr) => (if $v.is_f64() {
        Ok($v.as_f64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not f64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needf64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::String>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needstr {
    ($v:expr) => (if $v.is_string() {
        Ok($v.as_str().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not string")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needstr!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::Array>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needarray {
    ($v:expr) => (if $v.is_array() {
        Ok($v.as_array().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not array")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needarray!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a Result<Value::Object>.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needobj {
    ($v:expr) => (if $v.is_object() {
        Ok($v.as_object().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not object")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needobj!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}
<commit_msg>Add backticks to Rust types in macro docs<commit_after>// Copyright 2015-2017 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

macro_rules! want_macro {
    ($t:expr, $n:ident, $isf:ident, $asf:ident) => {
        /// Helper that returns an `Option<Value>`.
        /// You can optionally pass a JSON pointer to retrieve a
        /// nested key.
        #[macro_export]
        macro_rules! $n {
            ($v:expr) => (if $v.$isf() {
                Some($v.$asf().unwrap())
            } else {
                None
            });

            ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
                $n!(v)
            } else {
                None
            });
        }
    }
}

want_macro!("Null", wantnull, is_null, as_null);
want_macro!("Bool", wantbool, is_boolean, as_bool);
want_macro!("I64", wanti64, is_i64, as_i64);
want_macro!("U64", wantu64, is_u64, as_u64);
want_macro!("F64", wantf64, is_f64, as_f64);
want_macro!("String", wantstr, is_string, as_str);
want_macro!("Array", wantarray, is_array, as_array);
want_macro!("Object", wantobj, is_object, as_object);

// XXX We can't auto-generate macros due to the eager expansion of
// `$crate`, which causes any external libs to look for ::error::...
// instead of inapi::error::...
//
// macro_rules! need_macro {
//     ($t:expr, $n:ident, $isf:ident, $asf:ident) => {
//         /// Helper that returns a Result<Value::$t>.
//         /// You can optionally pass a JSON pointer to retrieve a
//         /// nested key.
//         #[macro_export]
//         macro_rules! $n {
//             ($v:expr) => (if $v.$isf() {
//                 Ok($v.$asf().unwrap())
//             } else {
//                 Err($crate::Error::Generic(format!("Value is not $t")))
//             });
//
//             ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
//                 $n!(v)
//             } else {
//                 Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
//             });
//         }
//     }
// }
//
// need_macro!("Null", neednull, is_null, as_null);
// need_macro!("Bool", needbool, is_boolean, as_bool);
// need_macro!("I64", needi64, is_i64, as_i64);
// need_macro!("U64", needu64, is_u64, as_u64);
// need_macro!("F64", needf64, is_f64, as_f64);
// need_macro!("String", needstr, is_string, as_str);
// need_macro!("Array", needarray, is_array, as_array);
// need_macro!("Object", needobj, is_object, as_object);

/// Helper that returns a `Result<Value::Null>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! neednull {
    ($v:expr) => (if $v.is_null() {
        Ok($v.as_null().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not null")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        neednull!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::Bool>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needbool {
    ($v:expr) => (if $v.is_boolean() {
        Ok($v.as_bool().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not boolean")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needbool!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::I64>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needi64 {
    ($v:expr) => (if $v.is_i64() {
        Ok($v.as_i64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not i64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needi64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::U64>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needu64 {
    ($v:expr) => (if $v.is_u64() {
        Ok($v.as_u64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not u64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needu64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::F64>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needf64 {
    ($v:expr) => (if $v.is_f64() {
        Ok($v.as_f64().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not f64")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needf64!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::String>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needstr {
    ($v:expr) => (if $v.is_string() {
        Ok($v.as_str().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not string")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needstr!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::Array>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needarray {
    ($v:expr) => (if $v.is_array() {
        Ok($v.as_array().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not array")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needarray!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}

/// Helper that returns a `Result<Value::Object>`.
/// You can optionally pass a JSON pointer to retrieve a nested key.
#[macro_export]
macro_rules! needobj {
    ($v:expr) => (if $v.is_object() {
        Ok($v.as_object().unwrap())
    } else {
        Err($crate::Error::Generic(format!("Value is not object")))
    });

    ($v:expr => $p:expr) => (if let Some(v) = $v.pointer($p) {
        needobj!(v)
    } else {
        Err($crate::Error::Generic(format!("Could not find {} in data", $p)))
    });
}
<|endoftext|>


<commit_before>use rustc_data_structures::fx::FxHashSet;
use std::fs;
use std::hash::{Hash, Hasher};
use std::path::Path;

use errors::Handler;

macro_rules! try_something {
    ($e:expr, $diag:expr, $out:expr) => ({
        match $e {
            Ok(c) => c,
            Err(e) => {
                $diag.struct_err(&e.to_string()).emit();
                return $out;
            }
        }
    })
}

#[derive(Debug, Clone, Eq)]
pub struct CssPath {
    pub name: String,
    pub children: FxHashSet<CssPath>,
}

// This PartialEq implementation IS NOT COMMUTATIVE!!!
//
// The order is very important: the second object must have all first's rules.
// However, the first doesn't require to have all second's rules.
impl PartialEq for CssPath {
    fn eq(&self, other: &CssPath) -> bool {
        if self.name != other.name {
            false
        } else {
            for child in &self.children {
                if !other.children.iter().any(|c| child == c) {
                    return false;
                }
            }
            true
        }
    }
}

impl Hash for CssPath {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for x in &self.children {
            x.hash(state);
        }
    }
}

impl CssPath {
    fn new(name: String) -> CssPath {
        CssPath {
            name,
            children: FxHashSet::default(),
        }
    }
}

/// All variants contain the position they occur.
#[derive(Debug, Clone, Copy)]
enum Events {
    StartLineComment(usize),
    StartComment(usize),
    EndComment(usize),
    InBlock(usize),
    OutBlock(usize),
}

impl Events {
    fn get_pos(&self) -> usize {
        match *self {
            Events::StartLineComment(p) |
            Events::StartComment(p) |
            Events::EndComment(p) |
            Events::InBlock(p) |
            Events::OutBlock(p) => p,
        }
    }

    fn is_comment(&self) -> bool {
        match *self {
            Events::StartLineComment(_) |
            Events::StartComment(_) |
            Events::EndComment(_) => true,
            _ => false,
        }
    }
}

fn previous_is_line_comment(events: &[Events]) -> bool {
    if let Some(&Events::StartLineComment(_)) = events.last() {
        true
    } else {
        false
    }
}

fn is_line_comment(pos: usize, v: &[u8], events: &[Events]) -> bool {
    if let Some(&Events::StartComment(_)) = events.last() {
        return false;
    }
    pos + 1 < v.len() && v[pos + 1] == b'/'
}

fn load_css_events(v: &[u8]) -> Vec<Events> {
    let mut pos = 0;
    let mut events = Vec::with_capacity(100);

    while pos + 1 < v.len() {
        match v[pos] {
            b'/' if pos + 1 < v.len() && v[pos + 1] == b'*' => {
                events.push(Events::StartComment(pos));
                pos += 1;
            }
            b'/' if is_line_comment(pos, v, &events) => {
                events.push(Events::StartLineComment(pos));
                pos += 1;
            }
            b'\n' if previous_is_line_comment(&events) => {
                events.push(Events::EndComment(pos));
            }
            b'*' if pos + 1 < v.len() && v[pos + 1] == b'/' => {
                events.push(Events::EndComment(pos + 2));
                pos += 1;
            }
            b'{' if !previous_is_line_comment(&events) => {
                if let Some(&Events::StartComment(_)) = events.last() {
                    pos += 1;
                    continue
                }
                events.push(Events::InBlock(pos + 1));
            }
            b'}' if !previous_is_line_comment(&events) => {
                if let Some(&Events::StartComment(_)) = events.last() {
                    pos += 1;
                    continue
                }
                events.push(Events::OutBlock(pos + 1));
            }
            _ => {}
        }
        pos += 1;
    }
    events
}

fn get_useful_next(events: &[Events], pos: &mut usize) -> Option<Events> {
    while *pos < events.len() {
        if !events[*pos].is_comment() {
            return Some(events[*pos]);
        }
        *pos += 1;
    }
    None
}

fn get_previous_positions(events: &[Events], mut pos: usize) -> Vec<usize> {
    let mut ret = Vec::with_capacity(3);

    ret.push(events[pos].get_pos());
    if pos > 0 {
        pos -= 1;
    }
    loop {
        if pos < 1 || !events[pos].is_comment() {
            let x = events[pos].get_pos();
            if *ret.last().unwrap() != x {
                ret.push(x);
            } else {
                ret.push(0);
            }
            break
        }
        ret.push(events[pos].get_pos());
        pos -= 1;
    }
    if ret.len() & 1 != 0 && events[pos].is_comment() {
        ret.push(0);
    }
    ret.iter().rev().cloned().collect()
}

fn build_rule(v: &[u8], positions: &[usize]) -> String {
    positions.chunks(2)
             .map(|x| ::std::str::from_utf8(&v[x[0]..x[1]]).unwrap_or(""))
             .collect::<String>()
             .trim()
             .replace("\n", " ")
             .replace("/", "")
             .replace("\t", " ")
             .replace("{", "")
             .replace("}", "")
             .split(' ')
             .filter(|s| s.len() > 0)
             .collect::<Vec<&str>>()
             .join(" ")
}

fn inner(v: &[u8], events: &[Events], pos: &mut usize) -> FxHashSet<CssPath> {
    let mut paths = Vec::with_capacity(50);

    while *pos < events.len() {
        if let Some(Events::OutBlock(_)) = get_useful_next(events, pos) {
            *pos += 1;
            break
        }
        if let Some(Events::InBlock(_)) = get_useful_next(events, pos) {
            paths.push(CssPath::new(build_rule(v, &get_previous_positions(events, *pos))));
            *pos += 1;
        }
        while let Some(Events::InBlock(_)) = get_useful_next(events, pos) {
            if let Some(ref mut path) = paths.last_mut() {
                for entry in inner(v, events, pos).iter() {
                    path.children.insert(entry.clone());
                }
            }
        }
        if let Some(Events::OutBlock(_)) = get_useful_next(events, pos) {
            *pos += 1;
        }
    }
    paths.iter().cloned().collect()
}

pub fn load_css_paths(v: &[u8]) -> CssPath {
    let events = load_css_events(v);
    let mut pos = 0;

    let mut parent = CssPath::new("parent".to_owned());
    parent.children = inner(v, &events, &mut pos);
    parent
}

pub fn get_differences(against: &CssPath, other: &CssPath, v: &mut Vec<String>) {
    if against.name != other.name {
        return
    } else {
        for child in &against.children {
            let mut found = false;
            let mut found_working = false;
            let mut tmp = Vec::new();

            for other_child in &other.children {
                if child.name == other_child.name {
                    if child != other_child {
                        get_differences(child, other_child, &mut tmp);
                    } else {
                        found_working = true;
                    }
                    found = true;
                    break
                }
            }
            if found == false {
                v.push(format!("  Missing \"{}\" rule", child.name));
            } else if found_working == false {
                v.extend(tmp.iter().cloned());
            }
        }
    }
}

pub fn test_theme_against<P: AsRef<Path>>(f: &P, against: &CssPath, diag: &Handler)
    -> (bool, Vec<String>)
{
    let data = try_something!(fs::read(f), diag, (false, vec![]));
    let paths = load_css_paths(&data);
    let mut ret = vec![];
    get_differences(against, &paths, &mut ret);
    (true, ret)
}

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

    #[test]
    fn test_comments_in_rules() {
        let text = r#"
rule a {}

rule b, c
// a line comment
{}

rule d
// another line comment
e {}

rule f/* a multine

comment*/{}

rule g/* another multine

comment*/h

i {}

rule j/*commeeeeent

you like things like "{}" in there? :)
*/
end {}"#;

        let against = r#"
rule a {}

rule b, c {}

rule d e {}

rule f {}

rule gh i {}

rule j end {}
"#;

        let mut ret = Vec::new();
        get_differences(&load_css_paths(against.as_bytes()),
                        &load_css_paths(text.as_bytes()),
                        &mut ret);
        assert!(ret.is_empty());
    }

    #[test]
    fn test_text() {
        let text = r#"
a
/* sdfs
*/ b
c // sdf
d {}
"#;
        let paths = load_css_paths(text.as_bytes());
        assert!(paths.children.contains(&CssPath::new("a b c d".to_owned())));
    }

    #[test]
    fn test_comparison() {
        let x = r#"
a {
    b {
        c {}
    }
}
"#;

        let y = r#"
a {
    b {}
}
"#;

        let against = load_css_paths(y.as_bytes());
        let other = load_css_paths(x.as_bytes());

        let mut ret = Vec::new();
        get_differences(&against, &other, &mut ret);
        assert!(ret.is_empty());
        get_differences(&other, &against, &mut ret);
        assert_eq!(ret, vec!["  Missing \"c\" rule".to_owned()]);
    }

    #[test]
    fn check_empty_css() {
        let events = load_css_events(&[]);
        assert_eq!(events.len(), 0);
    }
}
<commit_msg>Improve theme checker by removing unneeded conditions<commit_after>use rustc_data_structures::fx::FxHashSet;
use std::fs;
use std::hash::{Hash, Hasher};
use std::path::Path;

use errors::Handler;

macro_rules! try_something {
    ($e:expr, $diag:expr, $out:expr) => ({
        match $e {
            Ok(c) => c,
            Err(e) => {
                $diag.struct_err(&e.to_string()).emit();
                return $out;
            }
        }
    })
}

#[derive(Debug, Clone, Eq)]
pub struct CssPath {
    pub name: String,
    pub children: FxHashSet<CssPath>,
}

// This PartialEq implementation IS NOT COMMUTATIVE!!!
//
// The order is very important: the second object must have all first's rules.
// However, the first doesn't require to have all second's rules.
impl PartialEq for CssPath {
    fn eq(&self, other: &CssPath) -> bool {
        if self.name != other.name {
            false
        } else {
            for child in &self.children {
                if !other.children.iter().any(|c| child == c) {
                    return false;
                }
            }
            true
        }
    }
}

impl Hash for CssPath {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for x in &self.children {
            x.hash(state);
        }
    }
}

impl CssPath {
    fn new(name: String) -> CssPath {
        CssPath {
            name,
            children: FxHashSet::default(),
        }
    }
}

/// All variants contain the position they occur.
#[derive(Debug, Clone, Copy)]
enum Events {
    StartLineComment(usize),
    StartComment(usize),
    EndComment(usize),
    InBlock(usize),
    OutBlock(usize),
}

impl Events {
    fn get_pos(&self) -> usize {
        match *self {
            Events::StartLineComment(p) |
            Events::StartComment(p) |
            Events::EndComment(p) |
            Events::InBlock(p) |
            Events::OutBlock(p) => p,
        }
    }

    fn is_comment(&self) -> bool {
        match *self {
            Events::StartLineComment(_) |
            Events::StartComment(_) |
            Events::EndComment(_) => true,
            _ => false,
        }
    }
}

fn previous_is_line_comment(events: &[Events]) -> bool {
    if let Some(&Events::StartLineComment(_)) = events.last() {
        true
    } else {
        false
    }
}

fn is_line_comment(pos: usize, v: &[u8], events: &[Events]) -> bool {
    if let Some(&Events::StartComment(_)) = events.last() {
        return false;
    }
    v[pos + 1] == b'/'
}

fn load_css_events(v: &[u8]) -> Vec<Events> {
    let mut pos = 0;
    let mut events = Vec::with_capacity(100);

    while pos + 1 < v.len() {
        match v[pos] {
            b'/' if v[pos + 1] == b'*' => {
                events.push(Events::StartComment(pos));
                pos += 1;
            }
            b'/' if is_line_comment(pos, v, &events) => {
                events.push(Events::StartLineComment(pos));
                pos += 1;
            }
            b'\n' if previous_is_line_comment(&events) => {
                events.push(Events::EndComment(pos));
            }
            b'*' if v[pos + 1] == b'/' => {
                events.push(Events::EndComment(pos + 2));
                pos += 1;
            }
            b'{' if !previous_is_line_comment(&events) => {
                if let Some(&Events::StartComment(_)) = events.last() {
                    pos += 1;
                    continue
                }
                events.push(Events::InBlock(pos + 1));
            }
            b'}' if !previous_is_line_comment(&events) => {
                if let Some(&Events::StartComment(_)) = events.last() {
                    pos += 1;
                    continue
                }
                events.push(Events::OutBlock(pos + 1));
            }
            _ => {}
        }
        pos += 1;
    }
    events
}

fn get_useful_next(events: &[Events], pos: &mut usize) -> Option<Events> {
    while *pos < events.len() {
        if !events[*pos].is_comment() {
            return Some(events[*pos]);
        }
        *pos += 1;
    }
    None
}

fn get_previous_positions(events: &[Events], mut pos: usize) -> Vec<usize> {
    let mut ret = Vec::with_capacity(3);

    ret.push(events[pos].get_pos());
    if pos > 0 {
        pos -= 1;
    }
    loop {
        if pos < 1 || !events[pos].is_comment() {
            let x = events[pos].get_pos();
            if *ret.last().unwrap() != x {
                ret.push(x);
            } else {
                ret.push(0);
            }
            break
        }
        ret.push(events[pos].get_pos());
        pos -= 1;
    }
    if ret.len() & 1 != 0 && events[pos].is_comment() {
        ret.push(0);
    }
    ret.iter().rev().cloned().collect()
}

fn build_rule(v: &[u8], positions: &[usize]) -> String {
    positions.chunks(2)
             .map(|x| ::std::str::from_utf8(&v[x[0]..x[1]]).unwrap_or(""))
             .collect::<String>()
             .trim()
             .replace("\n", " ")
             .replace("/", "")
             .replace("\t", " ")
             .replace("{", "")
             .replace("}", "")
             .split(' ')
             .filter(|s| s.len() > 0)
             .collect::<Vec<&str>>()
             .join(" ")
}

fn inner(v: &[u8], events: &[Events], pos: &mut usize) -> FxHashSet<CssPath> {
    let mut paths = Vec::with_capacity(50);

    while *pos < events.len() {
        if let Some(Events::OutBlock(_)) = get_useful_next(events, pos) {
            *pos += 1;
            break
        }
        if let Some(Events::InBlock(_)) = get_useful_next(events, pos) {
            paths.push(CssPath::new(build_rule(v, &get_previous_positions(events, *pos))));
            *pos += 1;
        }
        while let Some(Events::InBlock(_)) = get_useful_next(events, pos) {
            if let Some(ref mut path) = paths.last_mut() {
                for entry in inner(v, events, pos).iter() {
                    path.children.insert(entry.clone());
                }
            }
        }
        if let Some(Events::OutBlock(_)) = get_useful_next(events, pos) {
            *pos += 1;
        }
    }
    paths.iter().cloned().collect()
}

pub fn load_css_paths(v: &[u8]) -> CssPath {
    let events = load_css_events(v);
    let mut pos = 0;

    let mut parent = CssPath::new("parent".to_owned());
    parent.children = inner(v, &events, &mut pos);
    parent
}

pub fn get_differences(against: &CssPath, other: &CssPath, v: &mut Vec<String>) {
    if against.name != other.name {
        return
    } else {
        for child in &against.children {
            let mut found = false;
            let mut found_working = false;
            let mut tmp = Vec::new();

            for other_child in &other.children {
                if child.name == other_child.name {
                    if child != other_child {
                        get_differences(child, other_child, &mut tmp);
                    } else {
                        found_working = true;
                    }
                    found = true;
                    break
                }
            }
            if found == false {
                v.push(format!("  Missing \"{}\" rule", child.name));
            } else if found_working == false {
                v.extend(tmp.iter().cloned());
            }
        }
    }
}

pub fn test_theme_against<P: AsRef<Path>>(
    f: &P,
    against: &CssPath,
    diag: &Handler,
) -> (bool, Vec<String>) {
    let data = try_something!(fs::read(f), diag, (false, vec![]));
    let paths = load_css_paths(&data);
    let mut ret = vec![];
    get_differences(against, &paths, &mut ret);
    (true, ret)
}

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

    #[test]
    fn test_comments_in_rules() {
        let text = r#"
rule a {}

rule b, c
// a line comment
{}

rule d
// another line comment
e {}

rule f/* a multine

comment*/{}

rule g/* another multine

comment*/h

i {}

rule j/*commeeeeent

you like things like "{}" in there? :)
*/
end {}"#;

        let against = r#"
rule a {}

rule b, c {}

rule d e {}

rule f {}

rule gh i {}

rule j end {}
"#;

        let mut ret = Vec::new();
        get_differences(&load_css_paths(against.as_bytes()),
                        &load_css_paths(text.as_bytes()),
                        &mut ret);
        assert!(ret.is_empty());
    }

    #[test]
    fn test_text() {
        let text = r#"
a
/* sdfs
*/ b
c // sdf
d {}
"#;
        let paths = load_css_paths(text.as_bytes());
        assert!(paths.children.contains(&CssPath::new("a b c d".to_owned())));
    }

    #[test]
    fn test_comparison() {
        let x = r#"
a {
    b {
        c {}
    }
}
"#;

        let y = r#"
a {
    b {}
}
"#;

        let against = load_css_paths(y.as_bytes());
        let other = load_css_paths(x.as_bytes());

        let mut ret = Vec::new();
        get_differences(&against, &other, &mut ret);
        assert!(ret.is_empty());
        get_differences(&other, &against, &mut ret);
        assert_eq!(ret, vec!["  Missing \"c\" rule".to_owned()]);
    }

    #[test]
    fn check_empty_css() {
        let events = load_css_events(&[]);
        assert_eq!(events.len(), 0);
    }

    #[test]
    fn check_invalid_css() {
        let events = load_css_events(b"*");
        assert_eq!(events.len(), 0);
    }
}
<|endoftext|>