<commit_before>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::Options;
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::type_case;

pub struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    created_interfaces: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn typescript_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        created_interfaces: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if ident != name {
        code = format!("export type {} = {};\n\n", name, ident) + &code;
    }

    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => ("any".into(), None),
        Bool => ("boolean".into(), None),
        StringT => ("string".into(), None),
        Integer => ("number".into(), None),
        Floating => ("number".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_interface_type(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                (format!("{} | undefined", inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("{}[]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("{{ [key: string]: {} }}", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (format!("[{}]", types.join(", ")), Some(defs.join("\n\n")))
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii_digit() {
            let temp = String::from("n") + name;
            type_case(&temp)
        } else {
            type_case(name)
        }
    } else {
        type_case(name)
    };
    if output_name.is_empty() {
        output_name.push_str("GeneratedType");
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

pub fn collapse_option(typ: &Shape) -> (bool, &Shape) {
    if let Shape::Optional(inner) = typ {
        return (true, &**inner);
    }
    (false, typ)
}

#[rustfmt::skip]
const RESERVED_WORDS: &[&str] = &["break", "case", "catch", "class", "const",
    "continue", "debugger", "default", "delete", "do", "else", "enum", "export", "extends", "false",
    "finally", "for", "function", "if", "import", "in", "instanceof", "new", "null", "return",
    "super", "switch", "this", "throw", "true", "try", "typeof", "var", "void", "while", "with",
    "implements", "interface", "let", "package", "private", "protected", "public", "static",
    "yield"];

pub fn is_ts_identifier(s: &str) -> bool {
    if RESERVED_WORDS.contains(&s) {
        return false;
    }

    if let Some((first, rest)) = s.as_bytes().split_first() {
        let first_valid = (b'a'..b'z').contains(first)
            || (b'A'..b'Z').contains(first)
            || *first == b'_'
            || *first == b'$';
        return first_valid
            && rest.iter().all(|b| {
                (b'a'..b'z').contains(b)
                    || (b'A'..b'Z').contains(b)
                    || *b == b'_'
                    || *b == b'$'
                    || (b'0'..b'9').contains(b)
            });
    }
    false
}

fn generate_interface_type(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_interfaces.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_interfaces
        .push((containing_shape.clone(), type_name.clone()));

    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let (was_optional, collapsed) = collapse_option(typ);

            let (field_type, child_defs) = type_from_shape(ctxt, name, collapsed);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            let escape_name = !is_ts_identifier(name);

            format!(
                "    {}{}{}{}: {};",
                if escape_name { "\"" } else { "" },
                name,
                if escape_name { "\"" } else { "" },
                if was_optional { "?" } else { "" },
                field_type
            )
        })
        .collect();

    let mut code = format!("export interface {} {{\n", type_name);

    if !fields.is_empty() {
        code += &fields.join("\n");
        code += "\n";
    }
    code += "}";

    if !defs.is_empty() {
        code += "\n\n";
        code += &defs.join("\n\n");
    }

    (type_name, Some(code))
}

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

    #[test]
    fn test_is_ts_identifier() {
        // Valid:
        assert!(is_ts_identifier("foobar"));
        assert!(is_ts_identifier("FOOBAR"));
        assert!(is_ts_identifier("foo_bar"));
        assert!(is_ts_identifier("$"));
        assert!(is_ts_identifier("foobar1"));

        // Invalid:
        assert!(!is_ts_identifier("1foobar"));
        assert!(!is_ts_identifier(""));
        assert!(!is_ts_identifier(" "));
        assert!(!is_ts_identifier(" foobar"));
        assert!(!is_ts_identifier("foobar "));
        assert!(!is_ts_identifier("foo bar"));
        assert!(!is_ts_identifier("foo.bar"));
        assert!(!is_ts_identifier("true"));
    }
}
<commit_msg>Use inclusive ranges for checking letter/number ranges<commit_after>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::Options;
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::type_case;

pub struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    created_interfaces: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn typescript_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        created_interfaces: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if ident != name {
        code = format!("export type {} = {};\n\n", name, ident) + &code;
    }

    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => ("any".into(), None),
        Bool => ("boolean".into(), None),
        StringT => ("string".into(), None),
        Integer => ("number".into(), None),
        Floating => ("number".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_interface_type(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                (format!("{} | undefined", inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("{}[]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("{{ [key: string]: {} }}", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (format!("[{}]", types.join(", ")), Some(defs.join("\n\n")))
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii_digit() {
            let temp = String::from("n") + name;
            type_case(&temp)
        } else {
            type_case(name)
        }
    } else {
        type_case(name)
    };
    if output_name.is_empty() {
        output_name.push_str("GeneratedType");
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

pub fn collapse_option(typ: &Shape) -> (bool, &Shape) {
    if let Shape::Optional(inner) = typ {
        return (true, &**inner);
    }
    (false, typ)
}

#[rustfmt::skip]
const RESERVED_WORDS: &[&str] = &["break", "case", "catch", "class", "const",
    "continue", "debugger", "default", "delete", "do", "else", "enum", "export", "extends", "false",
    "finally", "for", "function", "if", "import", "in", "instanceof", "new", "null", "return",
    "super", "switch", "this", "throw", "true", "try", "typeof", "var", "void", "while", "with",
    "implements", "interface", "let", "package", "private", "protected", "public", "static",
    "yield"];

pub fn is_ts_identifier(s: &str) -> bool {
    if RESERVED_WORDS.contains(&s) {
        return false;
    }

    if let Some((first, rest)) = s.as_bytes().split_first() {
        let first_valid = (b'a'..=b'z').contains(first)
            || (b'A'..=b'Z').contains(first)
            || *first == b'_'
            || *first == b'$';
        return first_valid
            && rest.iter().all(|b| {
                (b'a'..=b'z').contains(b)
                    || (b'A'..=b'Z').contains(b)
                    || *b == b'_'
                    || *b == b'$'
                    || (b'0'..=b'9').contains(b)
            });
    }
    false
}

fn generate_interface_type(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_interfaces.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_interfaces
        .push((containing_shape.clone(), type_name.clone()));

    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let (was_optional, collapsed) = collapse_option(typ);

            let (field_type, child_defs) = type_from_shape(ctxt, name, collapsed);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            let escape_name = !is_ts_identifier(name);

            format!(
                "    {}{}{}{}: {};",
                if escape_name { "\"" } else { "" },
                name,
                if escape_name { "\"" } else { "" },
                if was_optional { "?" } else { "" },
                field_type
            )
        })
        .collect();

    let mut code = format!("export interface {} {{\n", type_name);

    if !fields.is_empty() {
        code += &fields.join("\n");
        code += "\n";
    }
    code += "}";

    if !defs.is_empty() {
        code += "\n\n";
        code += &defs.join("\n\n");
    }

    (type_name, Some(code))
}

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

    #[test]
    fn test_is_ts_identifier() {
        // Valid:
        assert!(is_ts_identifier("foobar"));
        assert!(is_ts_identifier("FOOBAR"));
        assert!(is_ts_identifier("foo_bar"));
        assert!(is_ts_identifier("$"));
        assert!(is_ts_identifier("foobar1"));
        assert!(is_ts_identifier("zoneId"));

        // Invalid:
        assert!(!is_ts_identifier("1foobar"));
        assert!(!is_ts_identifier(""));
        assert!(!is_ts_identifier(" "));
        assert!(!is_ts_identifier(" foobar"));
        assert!(!is_ts_identifier("foobar "));
        assert!(!is_ts_identifier("foo bar"));
        assert!(!is_ts_identifier("foo.bar"));
        assert!(!is_ts_identifier("true"));
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019 Sanguk Park.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://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.
 */

// TODO: Refactor type names and remove this.
#![allow(non_camel_case_types)]

use core::mem;

use crate::cpu::*;
use crate::spinlock::*;
use crate::types::*;

const FLOAT_REG_BYTES: usize = 16;

/// The ID of a physical or virtual CPU.
pub type cpu_id_t = u64;

/// The integer type corresponding to the native register size.
pub type uintreg_t = u64;

/// Integer type large enough to hold a physical address.
pub type uintpaddr_t = uintptr_t;

/// Integer type large enough to hold a virtual address.
pub type uintvaddr_t = uintptr_t;

/// The type of a page table entry (PTE).
pub type pte_t = u64;

/// The struct for storing a floating point register.
///
/// 2 64-bit integers used to avoid need for FP support at this level.
#[repr(C, align(16))]
struct float_reg {
    // TODO: alignas(FLOAT_REG_BYTES)
    low: u64,
    high: u64,
}

const_assert_eq!(float_reg_size; mem::size_of::<float_reg>(), FLOAT_REG_BYTES);

/// Arch-specific information about a VM.
#[repr(C)]
pub struct ArchVm {
    /// The index of the last vCPU of this VM which ran on each pCPU. Each
    /// element of this array should only be read or written by code running
    /// on that CPU, which avoids contention and so no lock is needed to
    /// access this field.
    last_vcpu_on_cpu: [spci_vcpu_index_t; MAX_CPUS],
}

/// Type to represent the register state of a vCPU.
#[repr(C)]
pub struct ArchRegs {
    /// General purpose registers.
    r: [uintreg_t; 31],
    pc: uintreg_t,
    spsr: uintreg_t,

    /// System registers.
    lazy: ArchSysRegs,

    // Floating point registers.
    fp: [float_reg; 32],
    fpsr: uintreg_t,
    fpcr: uintreg_t,

    // TODO: 'hikey' environment has GIC version 2.
    //#[cfg(any(feature = "GIC_VERSION=3", feature = "GIC_VERSION=4"))]
    gic_ich_hcr_el2: uintreg_t,

    /// Peripheral registers, handled separately from other system registers.
    peripherals: ArchPeriRegs,
}

// from src/arch/aarch64/hypervisor/offset.h
// Note: always keep this constants same as ones in offset.h
const CPU_ID: usize = 0;
const CPU_STACK_BOTTOM: usize = 8;
const VCPU_REGS: usize = 32;
const REGS_LAZY: usize = 264;
const REGS_FREGS: usize = REGS_LAZY + 232;
//#[cfg(any(feature = "GIC_VERSION=3", feature = "GIC_VERSION=4"))]
const REGS_GIC: usize = REGS_FREGS + 528;

/// Checks above constants are correct.
/// HfO2: This checking was originally done in compile time in C. But it was
/// impossible because Rust compiler rejects construction of variables with
/// interior mutability (`VCpu` has `SpinLock`) in constant expressions. Hence
/// we check those constants in runtime.
pub fn arch_cpu_module_init() {
    assert_eq!(offset_of!(Cpu, id), CPU_ID);
    assert_eq!(offset_of!(Cpu, stack_bottom), CPU_STACK_BOTTOM);
    assert_eq!(
        offset_of!(VCpu, inner)
        + 8 // expected value of offset_of!(SpinLock<VCpuState>, data), but it
            // is not working. see Gilnaa/memoffset#21.
        + offset_of!(VCpuInner, regs),
        VCPU_REGS
    );
    assert_eq!(offset_of!(ArchRegs, lazy), REGS_LAZY);
    assert_eq!(offset_of!(ArchRegs, fp), REGS_FREGS);
    assert_eq!(offset_of!(ArchRegs, gic_ich_hcr_el2), REGS_GIC);
}

#[repr(C)]
pub struct ArchSysRegs {
    vmpidr_el2: uintreg_t,
    csselr_el1: uintreg_t,
    sctlr_el1: uintreg_t,
    actlr_el1: uintreg_t,
    cpacr_el1: uintreg_t,
    ttbr0_el1: uintreg_t,
    ttbr1_el1: uintreg_t,
    tcr_el1: uintreg_t,
    esr_el1: uintreg_t,
    afsr0_el1: uintreg_t,
    afsr1_el1: uintreg_t,
    far_el1: uintreg_t,
    mair_el1: uintreg_t,
    vbar_el1: uintreg_t,
    contextidr_el1: uintreg_t,
    tpidr_el0: uintreg_t,
    tpidrro_el0: uintreg_t,
    tpidr_el1: uintreg_t,
    amair_el1: uintreg_t,
    cntkctl_el1: uintreg_t,
    sp_el0: uintreg_t,
    sp_el1: uintreg_t,
    elr_el1: uintreg_t,
    spsr_el1: uintreg_t,
    par_el1: uintreg_t,
    hcr_el2: uintreg_t,
    cptr_el2: uintreg_t,
    cnthctl_el2: uintreg_t,
    vttbr_el2: uintreg_t,
}

#[repr(C)]
struct ArchPeriRegs {
    cntv_cval_el0: uintreg_t,
    cntv_ctl_el0: uintreg_t,
}
<commit_msg>Disable checking VCPU_REGS temporarily<commit_after>/*
 * Copyright 2019 Sanguk Park.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://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.
 */

// TODO: Refactor type names and remove this.
#![allow(non_camel_case_types)]

use core::mem;

use crate::cpu::*;
use crate::spinlock::*;
use crate::types::*;

const FLOAT_REG_BYTES: usize = 16;

/// The ID of a physical or virtual CPU.
pub type cpu_id_t = u64;

/// The integer type corresponding to the native register size.
pub type uintreg_t = u64;

/// Integer type large enough to hold a physical address.
pub type uintpaddr_t = uintptr_t;

/// Integer type large enough to hold a virtual address.
pub type uintvaddr_t = uintptr_t;

/// The type of a page table entry (PTE).
pub type pte_t = u64;

/// The struct for storing a floating point register.
///
/// 2 64-bit integers used to avoid need for FP support at this level.
#[repr(C, align(16))]
struct float_reg {
    // TODO: alignas(FLOAT_REG_BYTES)
    low: u64,
    high: u64,
}

const_assert_eq!(float_reg_size; mem::size_of::<float_reg>(), FLOAT_REG_BYTES);

/// Arch-specific information about a VM.
#[repr(C)]
pub struct ArchVm {
    /// The index of the last vCPU of this VM which ran on each pCPU. Each
    /// element of this array should only be read or written by code running
    /// on that CPU, which avoids contention and so no lock is needed to
    /// access this field.
    last_vcpu_on_cpu: [spci_vcpu_index_t; MAX_CPUS],
}

/// Type to represent the register state of a vCPU.
#[repr(C)]
pub struct ArchRegs {
    /// General purpose registers.
    r: [uintreg_t; 31],
    pc: uintreg_t,
    spsr: uintreg_t,

    /// System registers.
    lazy: ArchSysRegs,

    // Floating point registers.
    fp: [float_reg; 32],
    fpsr: uintreg_t,
    fpcr: uintreg_t,

    // TODO: 'hikey' environment has GIC version 2.
    //#[cfg(any(feature = "GIC_VERSION=3", feature = "GIC_VERSION=4"))]
    gic_ich_hcr_el2: uintreg_t,

    /// Peripheral registers, handled separately from other system registers.
    peripherals: ArchPeriRegs,
}

// from src/arch/aarch64/hypervisor/offset.h
// Note: always keep this constants same as ones in offset.h
const CPU_ID: usize = 0;
const CPU_STACK_BOTTOM: usize = 8;
const VCPU_REGS: usize = 32;
const REGS_LAZY: usize = 264;
const REGS_FREGS: usize = REGS_LAZY + 232;
//#[cfg(any(feature = "GIC_VERSION=3", feature = "GIC_VERSION=4"))]
const REGS_GIC: usize = REGS_FREGS + 528;

/// Checks above constants are correct.
/// HfO2: This checking was originally done in compile time in C. But it was
/// impossible because Rust compiler rejects construction of variables with
/// interior mutability (`VCpu` has `SpinLock`) in constant expressions. Hence
/// we check those constants in runtime.
pub fn arch_cpu_module_init() {
    assert_eq!(offset_of!(Cpu, id), CPU_ID);
    assert_eq!(offset_of!(Cpu, stack_bottom), CPU_STACK_BOTTOM);
    // assert_eq!(
    //     offset_of!(VCpu, inner)
    //     + 8 // expected value of offset_of!(SpinLock<VCpuState>, data), but it
    //         // is not working. see Gilnaa/memoffset#21.
    //     + offset_of!(VCpuInner, regs),
    //     VCPU_REGS
    // );
    assert_eq!(offset_of!(ArchRegs, lazy), REGS_LAZY);
    assert_eq!(offset_of!(ArchRegs, fp), REGS_FREGS);
    assert_eq!(offset_of!(ArchRegs, gic_ich_hcr_el2), REGS_GIC);
}

#[repr(C)]
pub struct ArchSysRegs {
    vmpidr_el2: uintreg_t,
    csselr_el1: uintreg_t,
    sctlr_el1: uintreg_t,
    actlr_el1: uintreg_t,
    cpacr_el1: uintreg_t,
    ttbr0_el1: uintreg_t,
    ttbr1_el1: uintreg_t,
    tcr_el1: uintreg_t,
    esr_el1: uintreg_t,
    afsr0_el1: uintreg_t,
    afsr1_el1: uintreg_t,
    far_el1: uintreg_t,
    mair_el1: uintreg_t,
    vbar_el1: uintreg_t,
    contextidr_el1: uintreg_t,
    tpidr_el0: uintreg_t,
    tpidrro_el0: uintreg_t,
    tpidr_el1: uintreg_t,
    amair_el1: uintreg_t,
    cntkctl_el1: uintreg_t,
    sp_el0: uintreg_t,
    sp_el1: uintreg_t,
    elr_el1: uintreg_t,
    spsr_el1: uintreg_t,
    par_el1: uintreg_t,
    hcr_el2: uintreg_t,
    cptr_el2: uintreg_t,
    cnthctl_el2: uintreg_t,
    vttbr_el2: uintreg_t,
}

#[repr(C)]
struct ArchPeriRegs {
    cntv_cval_el0: uintreg_t,
    cntv_ctl_el0: uintreg_t,
}
<|endoftext|>


<commit_before>#![crate_type = "lib"]
#![crate_name = "coinaddress"]
#![license = "MIT"]

//! Functions for validating the base58 hash checksums, including specifically
//! the bitcoin and litecoin addresses.

extern crate num;
extern crate rustc;
extern crate test;

use num::bigint::BigUint;
use std::num::{from_uint, from_int};
use std::num::Zero;
use rustc::util::sha2::{Sha256, Digest};

#[deriving(PartialEq, Show)]
pub enum ValidationError {
    TooShort,
    InvalidEncoding,
    HashMismatch,

    // currency specific
    NotBitcoin,
    NotLitecoin,
}

static BASE58_CHARS: &'static str = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

fn decode_base58(bc: &str) -> Option<BigUint> {
    let mut res: BigUint = Zero::zero();
    let b58: BigUint = from_int(58).unwrap();

    for c in bc.chars() {
        match BASE58_CHARS.find(c) {
            None => {
                return None
            },
            Some(x) => {
                res = res * b58 + from_uint(x).unwrap();
            }
        }
    }
    Some(res)
}

fn bigint_to_bytes(n: BigUint) -> Vec<u8> {
    let mut res: Vec<u8> = Vec::new();
    let mut tmp = n;
    let b256: BigUint = from_int(256).unwrap();

    while !tmp.is_zero() {
        res.insert(0, (tmp % b256).to_u8().unwrap());
        tmp = tmp / b256;
    }

    if res.len() == 0 {
        return vec!(0);
    }

    res
}

fn pad_to(v: Vec<u8>, len: uint) -> Vec<u8> {
    let mut tmp: Vec<u8> = v;
    while tmp.len() < len {
        tmp.insert(0, 0)
    }
    tmp
}

fn double_sha256(chunk: &[u8]) -> Vec<u8> {
    let mut hash = Sha256::new();
    hash.input(chunk);
    let mut hash2 = Sha256::new();
    hash2.input(hash.result_bytes().as_slice());
    hash2.result_bytes()
}

/// Validate provided generic base58 hash
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_base58_hash(addr: &str) -> Result<uint, ValidationError> {
    if addr.len() == 0 {
        return Err(TooShort);
    }

    let big = match decode_base58(addr) {
        None => return Err(InvalidEncoding),
        Some(x) => x
    };
    let bytes = bigint_to_bytes(big);
    let padded = pad_to(bytes, 25);

    let hash = double_sha256(padded.slice(0, padded.len() - 4));
    let short_hash = hash.slice(0, 4);
    let known = padded.slice(padded.len()-4, padded.len());
    if short_hash.as_slice() == known {
        Ok(padded[0].to_uint().unwrap())
    } else {
        Err(HashMismatch)
    }
}

/// Validate bitcoin address checksum
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_btc_address(addr: &str) -> Result<uint, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(0) => Ok(0),      // real address
        Ok(5) => Ok(5),      // script hash
        Ok(111) => Ok(111),  // testnet address
        Ok(_) => Err(NotBitcoin),
        Err(x) => Err(x)
    }
}

/// Validate litecoin address checksum
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_ltc_address(addr: &str) -> Result<uint, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(48) => Ok(48),      // real address
        Ok(111) => Ok(111),    // testnet address
        Ok(_) => Err(NotLitecoin),
        Err(x) => Err(x)
    }
}

#[test]
fn test_decoding() {
    use std::num::from_str_radix;

    assert_eq!(decode_base58("1"), from_int(0));
    assert_eq!(decode_base58("1AaZz"), from_int(0x1C8485));
    assert_eq!(decode_base58("AO"), None);
    assert_eq!(decode_base58("1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i"), from_str_radix("65A16059864A2FDBC7C99A4723A8395BC6F188EBC046B2FF", 16));
}

#[test]
fn test_to_bytes() {
    assert_eq!(bigint_to_bytes(Zero::zero()), vec!(0));
    assert_eq!(bigint_to_bytes(from_int(1).unwrap()), vec!(1));
    assert_eq!(bigint_to_bytes(from_int(256).unwrap()), vec!(1, 0));
}

#[test]
fn test_validate_hash() {
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_base58_hash("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYheX"), Err(HashMismatch));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYh  "), Err(InvalidEncoding));
    assert_eq!(validate_base58_hash("1"), Err(HashMismatch));
    assert_eq!(validate_base58_hash(""), Err(TooShort));
}

#[test]
fn test_validate_btc_address() {
    assert_eq!(validate_btc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_btc_address("3EktnHQD7RiAE6uzMj2ZifT9YgRrkSgzQX"), Ok(5));
    assert_eq!(validate_btc_address("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_btc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Err(NotBitcoin));
}

#[test]
fn test_validate_ltc_address() {
    assert_eq!(validate_ltc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Ok(48));
    assert_eq!(validate_ltc_address("muen9zszN6rVwXaFw48xh6YkdUSjJcfzek"), Ok(111));
    assert_eq!(validate_ltc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Err(NotLitecoin));
}

#[bench]
fn bench_validate_hash(b: &mut test::Bencher) {
    b.iter(|| {
        validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem")
    })
}
<commit_msg>Update imports<commit_after>#![crate_type = "lib"]
#![crate_name = "coinaddress"]
#![license = "MIT"]

//! Functions for validating the base58 hash checksums, including specifically
//! the bitcoin and litecoin addresses.

extern crate num;
extern crate rustc;
extern crate test;

use num::bigint::BigUint;
use num::Zero;
use std::num::{from_uint, from_int};
use rustc::util::sha2::{Sha256, Digest};

#[deriving(PartialEq, Show)]
pub enum ValidationError {
    TooShort,
    InvalidEncoding,
    HashMismatch,

    // currency specific
    NotBitcoin,
    NotLitecoin,
}

static BASE58_CHARS: &'static str = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

fn decode_base58(bc: &str) -> Option<BigUint> {
    let mut res: BigUint = Zero::zero();
    let b58: BigUint = from_int(58).unwrap();

    for c in bc.chars() {
        match BASE58_CHARS.find(c) {
            None => {
                return None
            },
            Some(x) => {
                res = res * b58 + from_uint(x).unwrap();
            }
        }
    }
    Some(res)
}

fn bigint_to_bytes(n: BigUint) -> Vec<u8> {
    let mut res: Vec<u8> = Vec::new();
    let mut tmp = n;
    let b256: BigUint = from_int(256).unwrap();

    while !tmp.is_zero() {
        res.insert(0, (tmp % b256).to_u8().unwrap());
        tmp = tmp / b256;
    }

    if res.len() == 0 {
        return vec!(0);
    }

    res
}

fn pad_to(v: Vec<u8>, len: uint) -> Vec<u8> {
    let mut tmp: Vec<u8> = v;
    while tmp.len() < len {
        tmp.insert(0, 0)
    }
    tmp
}

fn double_sha256(chunk: &[u8]) -> Vec<u8> {
    let mut hash = Sha256::new();
    hash.input(chunk);
    let mut hash2 = Sha256::new();
    hash2.input(hash.result_bytes().as_slice());
    hash2.result_bytes()
}

/// Validate provided generic base58 hash
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_base58_hash(addr: &str) -> Result<uint, ValidationError> {
    if addr.len() == 0 {
        return Err(TooShort);
    }

    let big = match decode_base58(addr) {
        None => return Err(InvalidEncoding),
        Some(x) => x
    };
    let bytes = bigint_to_bytes(big);
    let padded = pad_to(bytes, 25);

    let hash = double_sha256(padded.slice(0, padded.len() - 4));
    let short_hash = hash.slice(0, 4);
    let known = padded.slice(padded.len()-4, padded.len());
    if short_hash.as_slice() == known {
        Ok(padded[0].to_uint().unwrap())
    } else {
        Err(HashMismatch)
    }
}

/// Validate bitcoin address checksum
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_btc_address(addr: &str) -> Result<uint, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(0) => Ok(0),      // real address
        Ok(5) => Ok(5),      // script hash
        Ok(111) => Ok(111),  // testnet address
        Ok(_) => Err(NotBitcoin),
        Err(x) => Err(x)
    }
}

/// Validate litecoin address checksum
/// Returning the hash version/type if correct and an error otherwise
pub fn validate_ltc_address(addr: &str) -> Result<uint, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(48) => Ok(48),      // real address
        Ok(111) => Ok(111),    // testnet address
        Ok(_) => Err(NotLitecoin),
        Err(x) => Err(x)
    }
}

#[test]
fn test_decoding() {
    use std::num::from_str_radix;

    assert_eq!(decode_base58("1"), from_int(0));
    assert_eq!(decode_base58("1AaZz"), from_int(0x1C8485));
    assert_eq!(decode_base58("AO"), None);
    assert_eq!(decode_base58("1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i"), from_str_radix("65A16059864A2FDBC7C99A4723A8395BC6F188EBC046B2FF", 16));
}

#[test]
fn test_to_bytes() {
    assert_eq!(bigint_to_bytes(Zero::zero()), vec!(0));
    assert_eq!(bigint_to_bytes(from_int(1).unwrap()), vec!(1));
    assert_eq!(bigint_to_bytes(from_int(256).unwrap()), vec!(1, 0));
}

#[test]
fn test_validate_hash() {
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_base58_hash("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYheX"), Err(HashMismatch));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYh  "), Err(InvalidEncoding));
    assert_eq!(validate_base58_hash("1"), Err(HashMismatch));
    assert_eq!(validate_base58_hash(""), Err(TooShort));
}

#[test]
fn test_validate_btc_address() {
    assert_eq!(validate_btc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_btc_address("3EktnHQD7RiAE6uzMj2ZifT9YgRrkSgzQX"), Ok(5));
    assert_eq!(validate_btc_address("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_btc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Err(NotBitcoin));
}

#[test]
fn test_validate_ltc_address() {
    assert_eq!(validate_ltc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Ok(48));
    assert_eq!(validate_ltc_address("muen9zszN6rVwXaFw48xh6YkdUSjJcfzek"), Ok(111));
    assert_eq!(validate_ltc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Err(NotLitecoin));
}

#[bench]
fn bench_validate_hash(b: &mut test::Bencher) {
    b.iter(|| {
        validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem")
    })
}
<|endoftext|>


<commit_before>use crate::prelude::CanGetFile;
use crate::types::{requests::get_file::GetFile, ChannelPost, Message, MessageKind, MessageOrChannelPost};

/// A trait to obtain text from a message.
///
/// For example, this will return the text from text messages, or the caption of a photo.
pub trait MessageText {
    /// Obtain text from a message if available.
    fn text<'a>(&'a self) -> Option<String>;
}

impl MessageText for MessageOrChannelPost {
    fn text<'a>(&'a self) -> Option<String> {
        match self {
            MessageOrChannelPost::Message(msg) => msg.text(),
            MessageOrChannelPost::ChannelPost(post) => post.text(),
        }
    }
}

impl MessageText for Message {
    fn text<'a>(&'a self) -> Option<String> {
        self.kind.text()
    }
}

impl MessageText for MessageKind {
    fn text<'a>(&'a self) -> Option<String> {
        match self {
            MessageKind::Text { data, .. } => Some(data.to_owned()),
            MessageKind::Audio { data } => data.title.to_owned(),
            MessageKind::Document { data, caption } => {
                caption.clone().or_else(|| data.file_name.clone())
            }
            MessageKind::Photo { caption, .. } => caption.to_owned(),
            MessageKind::Sticker { .. } => None,
            MessageKind::Video { caption, .. } => caption.to_owned(),
            MessageKind::Voice { .. } => None,
            MessageKind::VideoNote { .. } => None,
            MessageKind::Contact { data } => Some(data.first_name.to_owned()),
            MessageKind::Location { .. } => None,
            MessageKind::Venue { data } => Some(data.title.to_owned()),
            MessageKind::NewChatMembers { .. } => None,
            MessageKind::LeftChatMember { .. } => None,
            MessageKind::NewChatTitle { data } => Some(data.to_owned()),
            MessageKind::NewChatPhoto { .. } => None,
            MessageKind::DeleteChatPhoto => None,
            MessageKind::GroupChatCreated => None,
            MessageKind::SupergroupChatCreated => None,
            MessageKind::ChannelChatCreated => None,
            MessageKind::MigrateToChatId { .. } => None,
            MessageKind::MigrateFromChatId { .. } => None,
            MessageKind::PinnedMessage { data } => data.text(),
            MessageKind::Unknown { .. } => None,
        }
    }
}

impl MessageText for ChannelPost {
    fn text<'a>(&'a self) -> Option<String> {
        self.kind.text()
    }
}

/// A trait to obtain `GetFile` requests from a message.
///
/// Many message kinds such as `Sticker` return a single `GetFile`.
/// Message kinds like `Photo` might return more if an album is posted.
pub trait MessageGetFiles {
    /// Obtain files from a message if available.
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>>;
}

impl MessageGetFiles for MessageOrChannelPost {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        match self {
            MessageOrChannelPost::Message(msg) => msg.get_files(),
            MessageOrChannelPost::ChannelPost(post) => post.get_files(),
        }
    }
}

impl MessageGetFiles for Message {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        self.kind.get_files()
    }
}

impl MessageGetFiles for MessageKind {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        match self {
            MessageKind::Text { .. } => None,
            MessageKind::Audio { data } => Some(vec![data.get_file()]),
            MessageKind::Document { data, .. } => Some(vec![data.get_file()]),
            MessageKind::Photo { data, .. } => {
                Some(data.into_iter().map(|f| f.get_file()).collect())
            }
            MessageKind::Sticker { data } => Some(vec![data.get_file()]),
            MessageKind::Video { data, .. } => Some(vec![data.get_file()]),
            MessageKind::Voice { data } => Some(vec![data.get_file()]),
            MessageKind::VideoNote { data } => Some(vec![data.get_file()]),
            MessageKind::Contact { .. } => None,
            MessageKind::Location { .. } => None,
            MessageKind::Venue { .. } => None,
            MessageKind::NewChatMembers { .. } => None,
            MessageKind::LeftChatMember { .. } => None,
            MessageKind::NewChatTitle { .. } => None,
            MessageKind::NewChatPhoto { data } => {
                Some(data.into_iter().map(|f| f.get_file()).collect())
            }
            MessageKind::DeleteChatPhoto => None,
            MessageKind::GroupChatCreated => None,
            MessageKind::SupergroupChatCreated => None,
            MessageKind::ChannelChatCreated => None,
            MessageKind::MigrateToChatId { .. } => None,
            MessageKind::MigrateFromChatId { .. } => None,
            MessageKind::PinnedMessage { .. } => None,
            MessageKind::Unknown { .. } => None,
        }
    }
}

impl MessageGetFiles for ChannelPost {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        self.kind.get_files()
    }
}
<commit_msg>Add thumbnails to message file list as well<commit_after>use crate::prelude::CanGetFile;
use crate::types::{requests::get_file::GetFile, ChannelPost, Message, MessageKind, MessageOrChannelPost};

/// A trait to obtain text from a message.
///
/// For example, this will return the text from text messages, or the caption of a photo.
pub trait MessageText {
    /// Obtain text from a message if available.
    fn text<'a>(&'a self) -> Option<String>;
}

impl MessageText for MessageOrChannelPost {
    fn text<'a>(&'a self) -> Option<String> {
        match self {
            MessageOrChannelPost::Message(msg) => msg.text(),
            MessageOrChannelPost::ChannelPost(post) => post.text(),
        }
    }
}

impl MessageText for Message {
    fn text<'a>(&'a self) -> Option<String> {
        self.kind.text()
    }
}

impl MessageText for MessageKind {
    fn text<'a>(&'a self) -> Option<String> {
        match self {
            MessageKind::Text { data, .. } => Some(data.to_owned()),
            MessageKind::Audio { data } => data.title.to_owned(),
            MessageKind::Document { data, caption } => {
                caption.clone().or_else(|| data.file_name.clone())
            }
            MessageKind::Photo { caption, .. } => caption.to_owned(),
            MessageKind::Sticker { .. } => None,
            MessageKind::Video { caption, .. } => caption.to_owned(),
            MessageKind::Voice { .. } => None,
            MessageKind::VideoNote { .. } => None,
            MessageKind::Contact { data } => Some(data.first_name.to_owned()),
            MessageKind::Location { .. } => None,
            MessageKind::Venue { data } => Some(data.title.to_owned()),
            MessageKind::NewChatMembers { .. } => None,
            MessageKind::LeftChatMember { .. } => None,
            MessageKind::NewChatTitle { data } => Some(data.to_owned()),
            MessageKind::NewChatPhoto { .. } => None,
            MessageKind::DeleteChatPhoto => None,
            MessageKind::GroupChatCreated => None,
            MessageKind::SupergroupChatCreated => None,
            MessageKind::ChannelChatCreated => None,
            MessageKind::MigrateToChatId { .. } => None,
            MessageKind::MigrateFromChatId { .. } => None,
            MessageKind::PinnedMessage { data } => data.text(),
            MessageKind::Unknown { .. } => None,
        }
    }
}

impl MessageText for ChannelPost {
    fn text<'a>(&'a self) -> Option<String> {
        self.kind.text()
    }
}

/// A trait to obtain `GetFile` requests from a message.
///
/// Many message kinds such as `Sticker` return a single `GetFile`.
/// Message kinds like `Photo` might return more if an album is posted.
/// A video, video note or document returns any thumbnail as well.
pub trait MessageGetFiles {
    /// Obtain files from a message if available.
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>>;
}

impl MessageGetFiles for MessageOrChannelPost {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        match self {
            MessageOrChannelPost::Message(msg) => msg.get_files(),
            MessageOrChannelPost::ChannelPost(post) => post.get_files(),
        }
    }
}

impl MessageGetFiles for Message {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        self.kind.get_files()
    }
}

impl MessageGetFiles for MessageKind {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        match self {
            MessageKind::Text { .. } => None,
            MessageKind::Audio { data } => Some(vec![data.get_file()]),
            MessageKind::Document { data, .. } => {
                let mut files = vec![data.get_file()];
                if let Some(thumb) = &data.thumb {
                    files.push(thumb.get_file());
                }
                Some(files)
            },
            MessageKind::Photo { data, .. } => {
                Some(data.into_iter().map(|f| f.get_file()).collect())
            }
            MessageKind::Sticker { data } => Some(vec![data.get_file()]),
            MessageKind::Video { data, .. } => {
                let mut files = vec![data.get_file()];
                if let Some(thumb) = &data.thumb {
                    files.push(thumb.get_file());
                }
                Some(files)
            },
            MessageKind::Voice { data } => Some(vec![data.get_file()]),
            MessageKind::VideoNote { data, .. } => {
                let mut files = vec![data.get_file()];
                if let Some(thumb) = &data.thumb {
                    files.push(thumb.get_file());
                }
                Some(files)
            },
            MessageKind::Contact { .. } => None,
            MessageKind::Location { .. } => None,
            MessageKind::Venue { .. } => None,
            MessageKind::NewChatMembers { .. } => None,
            MessageKind::LeftChatMember { .. } => None,
            MessageKind::NewChatTitle { .. } => None,
            MessageKind::NewChatPhoto { data } => {
                Some(data.into_iter().map(|f| f.get_file()).collect())
            }
            MessageKind::DeleteChatPhoto => None,
            MessageKind::GroupChatCreated => None,
            MessageKind::SupergroupChatCreated => None,
            MessageKind::ChannelChatCreated => None,
            MessageKind::MigrateToChatId { .. } => None,
            MessageKind::MigrateFromChatId { .. } => None,
            MessageKind::PinnedMessage { .. } => None,
            MessageKind::Unknown { .. } => None,
        }
    }
}

impl MessageGetFiles for ChannelPost {
    fn get_files<'a>(&'a self) -> Option<Vec<GetFile>> {
        self.kind.get_files()
    }
}
<|endoftext|>


<commit_before>use bot::BotEnv;
use database::models::EphemeralMessage;
use error::*;
use tasks::{RunsTask, Wait};

use chrono::{Utc, Duration};

use diesel;
use diesel::prelude::*;

use serenity::model::id::ChannelId;

use std::sync::Arc;
use std::thread;

#[derive(Debug, Default)]
pub struct EphemeralMessageTask {
  next_sleep: i64
}

impl RunsTask for EphemeralMessageTask {
  fn start(mut self, _: Arc<BotEnv>) {
    loop {
      thread::sleep(Duration::seconds(self.next_sleep).to_std().unwrap());
      if self.next_sleep == 0 {
        self.next_sleep = 1800;
      }
      let next_half_hour = (Utc::now() + Duration::minutes(30)).timestamp();
      info!("Checking for ephemeral messages");
      let res: Result<Vec<EphemeralMessage>> = ::bot::with_connection(|c| {
        use database::schema::ephemeral_messages::dsl;
        dsl::ephemeral_messages
          .filter(dsl::expires_on.le(next_half_hour))
          .load(c)
      }).chain_err(|| "could not load ephemeral messages");
      let mut msgs = match res {
        Ok(m) => m,
        Err(e) => {
          warn!("error loading ephemeral messages: {}", e);
          continue;
        }
      };

      if msgs.is_empty() {
        continue;
      }

      msgs.sort_by_key(|m| m.expires_on);

      ::std::thread::spawn(move || {
        for (wait, eph) in Wait::new(msgs.into_iter().map(|m| (m.expires_on, m))) {
          ::std::thread::sleep(Duration::seconds(wait).to_std().unwrap());

          let channel = ChannelId(*eph.channel_id);
          match channel.delete_message(*eph.message_id) {
            Ok(_) => {
              if let Err(e) = ::bot::with_connection(|c| diesel::delete(&eph).execute(c)) {
                warn!("could not delete ephemeral message (id: {}) from database: {}", eph.id, e);
              }
            }
            Err(e) => warn!("could not delete ephemeral message (id: {}): {}", eph.id, e)
          }
        }
      });

      info!("Done checking for ephemeral messages");
    }
  }
}
<commit_msg>chore: add comma<commit_after>use bot::BotEnv;
use database::models::EphemeralMessage;
use error::*;
use tasks::{RunsTask, Wait};

use chrono::{Utc, Duration};

use diesel;
use diesel::prelude::*;

use serenity::model::id::ChannelId;

use std::sync::Arc;
use std::thread;

#[derive(Debug, Default)]
pub struct EphemeralMessageTask {
  next_sleep: i64
}

impl RunsTask for EphemeralMessageTask {
  fn start(mut self, _: Arc<BotEnv>) {
    loop {
      thread::sleep(Duration::seconds(self.next_sleep).to_std().unwrap());
      if self.next_sleep == 0 {
        self.next_sleep = 1800;
      }
      let next_half_hour = (Utc::now() + Duration::minutes(30)).timestamp();
      info!("Checking for ephemeral messages");
      let res: Result<Vec<EphemeralMessage>> = ::bot::with_connection(|c| {
        use database::schema::ephemeral_messages::dsl;
        dsl::ephemeral_messages
          .filter(dsl::expires_on.le(next_half_hour))
          .load(c)
      }).chain_err(|| "could not load ephemeral messages");
      let mut msgs = match res {
        Ok(m) => m,
        Err(e) => {
          warn!("error loading ephemeral messages: {}", e);
          continue;
        }
      };

      if msgs.is_empty() {
        continue;
      }

      msgs.sort_by_key(|m| m.expires_on);

      ::std::thread::spawn(move || {
        for (wait, eph) in Wait::new(msgs.into_iter().map(|m| (m.expires_on, m))) {
          ::std::thread::sleep(Duration::seconds(wait).to_std().unwrap());

          let channel = ChannelId(*eph.channel_id);
          match channel.delete_message(*eph.message_id) {
            Ok(_) => {
              if let Err(e) = ::bot::with_connection(|c| diesel::delete(&eph).execute(c)) {
                warn!("could not delete ephemeral message (id: {}) from database: {}", eph.id, e);
              }
            },
            Err(e) => warn!("could not delete ephemeral message (id: {}): {}", eph.id, e)
          }
        }
      });

      info!("Done checking for ephemeral messages");
    }
  }
}
<|endoftext|>


<commit_before>pub fn gcd(x: i64, y: i64) -> i64 {
    if y == 0 {
        x.abs()
    } else {
        gcd(y, x % y)
    }
}

/// Finds a solution of Bézout's identity ``s * x + t * y = d``
/// where `d` is the greatest common divisor of `x`, `y`.
pub fn gcd_ext(x: i64, y: i64) -> (i64, i64, i64) {
    if y == 0 {
        (x.abs(), x.signum(), 0)
    } else {
        let (d, s, t) = gcd_ext(y, x % y);
        (d, t, s - x / y * t)
    }
}

pub fn lcm(x: i64, y: i64) -> i64 {
    x / gcd(x, y) * y
}

#[cfg(test)]
mod tests {
    use super::{gcd, gcd_ext, lcm};

    #[test]
    fn test_gcd() {
        fn case(x: i64, y: i64, d: i64) {
            assert_eq!(gcd(x, y), d, "{:?}", (x, y, d));
            assert_eq!(gcd(y, x), d, "{:?}", (y, x, d));
        }

        case(1, 1, 1);
        case(3, 5, 1);
        case(2, 2, 2);
        case(4, 6, 2);
        case(18, 24, 6);
        case(-1, -1, 1);
        case(4, -6, 2);
    }

    #[test]
    fn test_gcd_ext() {
        fn case(x: i64, y: i64, (d, s, t): (i64, i64, i64)) {
            assert_eq!(gcd_ext(x, y), (d, s, t), "{:?}", (x, y));
            assert_eq!(s * x + t * y, d);
        }

        case(3, 7, (1, -2, 1));
        case(4, 6, (2, -1, 1));
        case(360 * 5 * 5 * 11, 360 * 7 * 7 * 13, (360, -227, 98));
        case(4, -6, (2, -1, -1));
        case(-9, -12, (3, 1, -1));
    }

    #[test]
    fn test_lcm() {
        assert_eq!(lcm(1, 1), 1);
        assert_eq!(lcm(12, 18), 36);
    }
}
<commit_msg>feat: Add 'chinese_reminder' function<commit_after>pub fn gcd(x: i64, y: i64) -> i64 {
    if y == 0 {
        x.abs()
    } else {
        gcd(y, x % y)
    }
}

/// Finds a solution of Bézout's identity ``s * x + t * y = d``
/// where `d` is the greatest common divisor of `x`, `y`.
pub fn gcd_ext(x: i64, y: i64) -> (i64, i64, i64) {
    if y == 0 {
        (x.abs(), x.signum(), 0)
    } else {
        let (d, s, t) = gcd_ext(y, x % y);
        (d, t, s - x / y * t)
    }
}

pub fn lcm(x: i64, y: i64) -> i64 {
    x / gcd(x, y) * y
}

/// Solves the following simultaneous equations:
///     x ≡ a1 (mod n1),
///     x ≡ a2 (mod n2)
/// to get `(x, m)`
/// where m = `lcm(n1, n2)`
/// and x is the unique solution in range `0 <= x < m`.
/// Based on Chinese Reminder Theorem.
/// To reduce, use `Some((0, 1))` as unit.
pub fn chinese_reminder(a1: i64, n1: i64, a2: i64, n2: i64) -> Option<(i64, i64)> {
    let (d, x, _) = gcd_ext(n1, n2);
    if (a2 - a1).abs() % d != 0 {
        return None;
    }

    let m = n1 * (n2 / d);
    let t = (a2 - a1) / d * x % (n2 / d);
    let u = ((a1 + n1 * t) % m + m) % m;
    Some((u, m))
}

#[cfg(test)]
mod tests {
    use super::{chinese_reminder, gcd, gcd_ext, lcm};

    #[test]
    fn test_gcd() {
        fn case(x: i64, y: i64, d: i64) {
            assert_eq!(gcd(x, y), d, "{:?}", (x, y, d));
            assert_eq!(gcd(y, x), d, "{:?}", (y, x, d));
        }

        case(1, 1, 1);
        case(3, 5, 1);
        case(2, 2, 2);
        case(4, 6, 2);
        case(18, 24, 6);
        case(-1, -1, 1);
        case(4, -6, 2);
    }

    #[test]
    fn test_gcd_ext() {
        fn case(x: i64, y: i64, (d, s, t): (i64, i64, i64)) {
            assert_eq!(gcd_ext(x, y), (d, s, t), "{:?}", (x, y));
            assert_eq!(s * x + t * y, d);
        }

        case(3, 7, (1, -2, 1));
        case(4, 6, (2, -1, 1));
        case(360 * 5 * 5 * 11, 360 * 7 * 7 * 13, (360, -227, 98));
        case(4, -6, (2, -1, -1));
        case(-9, -12, (3, 1, -1));
    }

    #[test]
    fn test_lcm() {
        assert_eq!(lcm(1, 1), 1);
        assert_eq!(lcm(12, 18), 36);
    }

    #[test]
    fn test_chinese_reminder() {
        assert_eq!(
            vec![(2, 3), (3, 5), (2, 7)]
                .into_iter()
                .fold(Some((0, 1)), |acc, (y, m)| acc
                    .and_then(|(x, n)| chinese_reminder(x, n, y, m))),
            Some((23, 105))
        );
    }
}
<|endoftext|>


<commit_before>use super::error::{ErrorKind, Result};
use super::header;
use super::util::tcpconnection;
use super::{ServicePair, ServiceResult};
use byteorder::WriteBytesExt;
use serde_rosmsg::{from_reader, to_writer};
use std;
use std::collections::HashMap;
use std::net::TcpListener;
use std::sync::Arc;
use std::thread;

pub struct Service {
    pub api: String,
    pub msg_type: String,
    pub service: String,
    _raii: tcpconnection::Raii,
}

impl Service {
    pub fn new<T, F>(
        hostname: &str,
        port: u16,
        service: &str,
        node_name: &str,
        handler: F,
    ) -> Result<Service>
    where
        T: ServicePair,
        F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
    {
        let listener = TcpListener::bind((hostname, port))?;
        let socket_address = listener.local_addr()?;
        let api = format!("rosrpc://{}:{}", hostname, socket_address.port());
        let (raii, listener) = tcpconnection::iterate(listener, format!("service '{}'", service));
        Ok(Service::wrap_stream::<T, _, _, _>(
            service,
            node_name,
            handler,
            raii,
            listener,
            &api,
        ))
    }

    fn wrap_stream<T, U, V, F>(
        service: &str,
        node_name: &str,
        handler: F,
        raii: tcpconnection::Raii,
        listener: V,
        api: &str,
    ) -> Service
    where
        T: ServicePair,
        U: std::io::Read + std::io::Write + Send + 'static,
        V: Iterator<Item = U> + Send + 'static,
        F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
    {
        let service_name = String::from(service);
        let node_name = String::from(node_name);
        thread::spawn(move || {
            listen_for_clients::<T, _, _, _>(&service_name, &node_name, handler, listener)
        });
        Service {
            api: String::from(api),
            msg_type: T::msg_type(),
            service: String::from(service),
            _raii: raii,
        }
    }
}

fn listen_for_clients<T, U, V, F>(service: &str, node_name: &str, handler: F, connections: V)
where
    T: ServicePair,
    U: std::io::Read + std::io::Write + Send + 'static,
    V: Iterator<Item = U>,
    F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
{
    let handler = Arc::new(handler);
    for mut stream in connections {
        // Service request starts by exchanging connection headers
        if let Err(err) = exchange_headers::<T, _>(&mut stream, service, node_name) {
            error!(
                "Failed to exchange headers for service '{}': {}",
                service, err
            );
            continue;
        }

        // Spawn a thread for handling requests
        spawn_request_handler::<T, U, F>(stream, Arc::clone(&handler));
    }
}

fn exchange_headers<T, U>(stream: &mut U, service: &str, node_name: &str) -> Result<()>
where
    T: ServicePair,
    U: std::io::Write + std::io::Read,
{
    read_request::<T, U>(stream, service)?;
    write_response::<T, U>(stream, node_name)
}

fn read_request<T: ServicePair, U: std::io::Read>(stream: &mut U, service: &str) -> Result<()> {
    let fields = header::decode(stream)?;
    header::match_field(&fields, "service", service)?;
    if fields.get("callerid").is_none() {
        bail!(ErrorKind::HeaderMissingField("callerid".into()));
    }
    if header::match_field(&fields, "probe", "1").is_ok() {
        return Ok(());
    }
    header::match_field(&fields, "md5sum", &T::md5sum())
}

fn write_response<T, U>(stream: &mut U, node_name: &str) -> Result<()>
where
    T: ServicePair,
    U: std::io::Write,
{
    let mut fields = HashMap::<String, String>::new();
    fields.insert(String::from("callerid"), String::from(node_name));
    fields.insert(String::from("md5sum"), T::md5sum());
    fields.insert(String::from("type"), T::msg_type());
    header::encode(stream, &fields)?;
    Ok(())
}

fn spawn_request_handler<T, U, F>(stream: U, handler: Arc<F>)
where
    T: ServicePair,
    U: std::io::Read + std::io::Write + Send + 'static,
    F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
{
    thread::spawn(move || {
        if let Err(err) = handle_request_loop::<T, U, F>(stream, &handler) {
            if !err.is_closed_connection() {
                let info = err.iter()
                    .map(|v| format!("{}", v))
                    .collect::<Vec<_>>()
                    .join("\nCaused by:");
                error!("{}", info);
            }
        }
    });
}

fn handle_request_loop<T, U, F>(mut stream: U, handler: &F) -> Result<()>
where
    T: ServicePair,
    U: std::io::Read + std::io::Write,
    F: Fn(T::Request) -> ServiceResult<T::Response>,
{
    // Receive request from client
    // Break out of loop in case of failure to read request
    while let Ok(req) = from_reader(&mut stream) {
        // Call function that handles request and returns response
        match handler(req) {
            Ok(res) => {
                // Send True flag and response in case of success
                stream.write_u8(1)?;
                to_writer(&mut stream, &res)?;
            }
            Err(message) => {
                // Send False flag and error message string in case of failure
                stream.write_u8(0)?;
                to_writer(&mut stream, &message)?;
            }
        };
    }

    // Upon failure to read request, send client failure message
    // This can be caused by actual issues or by the client stopping the connection
    stream.write_u8(0)?;
    to_writer(&mut stream, &"Failed to parse passed arguments")?;
    Ok(())
}
<commit_msg>do not emit error when rospy.wait_for_service() called<commit_after>use super::error::{ErrorKind, Result};
use super::header;
use super::util::tcpconnection;
use super::{ServicePair, ServiceResult};
use byteorder::WriteBytesExt;
use serde_rosmsg::{from_reader, to_writer};
use std;
use std::collections::HashMap;
use std::net::TcpListener;
use std::sync::Arc;
use std::thread;

pub struct Service {
    pub api: String,
    pub msg_type: String,
    pub service: String,
    _raii: tcpconnection::Raii,
}

impl Service {
    pub fn new<T, F>(
        hostname: &str,
        port: u16,
        service: &str,
        node_name: &str,
        handler: F,
    ) -> Result<Service>
    where
        T: ServicePair,
        F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
    {
        let listener = TcpListener::bind((hostname, port))?;
        let socket_address = listener.local_addr()?;
        let api = format!("rosrpc://{}:{}", hostname, socket_address.port());
        let (raii, listener) = tcpconnection::iterate(listener, format!("service '{}'", service));
        Ok(Service::wrap_stream::<T, _, _, _>(
            service,
            node_name,
            handler,
            raii,
            listener,
            &api,
        ))
    }

    fn wrap_stream<T, U, V, F>(
        service: &str,
        node_name: &str,
        handler: F,
        raii: tcpconnection::Raii,
        listener: V,
        api: &str,
    ) -> Service
    where
        T: ServicePair,
        U: std::io::Read + std::io::Write + Send + 'static,
        V: Iterator<Item = U> + Send + 'static,
        F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
    {
        let service_name = String::from(service);
        let node_name = String::from(node_name);
        thread::spawn(move || {
            listen_for_clients::<T, _, _, _>(&service_name, &node_name, handler, listener)
        });
        Service {
            api: String::from(api),
            msg_type: T::msg_type(),
            service: String::from(service),
            _raii: raii,
        }
    }
}

fn listen_for_clients<T, U, V, F>(service: &str, node_name: &str, handler: F, connections: V)
where
    T: ServicePair,
    U: std::io::Read + std::io::Write + Send + 'static,
    V: Iterator<Item = U>,
    F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
{
    let handler = Arc::new(handler);
    for mut stream in connections {
        // Service request starts by exchanging connection headers
        if let Err(err) = exchange_headers::<T, _>(&mut stream, service, node_name) {
            // Cconnection can be closed when a client checks for a service.
            if !err.is_closed_connection() {
                error!(
                    "Failed to exchange headers for service '{}': {}",
                    service, err
                );
            }
            continue;
        }

        // Spawn a thread for handling requests
        spawn_request_handler::<T, U, F>(stream, Arc::clone(&handler));
    }
}

fn exchange_headers<T, U>(stream: &mut U, service: &str, node_name: &str) -> Result<()>
where
    T: ServicePair,
    U: std::io::Write + std::io::Read,
{
    read_request::<T, U>(stream, service)?;
    write_response::<T, U>(stream, node_name)
}

fn read_request<T: ServicePair, U: std::io::Read>(stream: &mut U, service: &str) -> Result<()> {
    let fields = header::decode(stream)?;
    header::match_field(&fields, "service", service)?;
    if fields.get("callerid").is_none() {
        bail!(ErrorKind::HeaderMissingField("callerid".into()));
    }
    if header::match_field(&fields, "probe", "1").is_ok() {
        return Ok(());
    }
    header::match_field(&fields, "md5sum", &T::md5sum())
}

fn write_response<T, U>(stream: &mut U, node_name: &str) -> Result<()>
where
    T: ServicePair,
    U: std::io::Write,
{
    let mut fields = HashMap::<String, String>::new();
    fields.insert(String::from("callerid"), String::from(node_name));
    fields.insert(String::from("md5sum"), T::md5sum());
    fields.insert(String::from("type"), T::msg_type());
    header::encode(stream, &fields)?;
    Ok(())
}

fn spawn_request_handler<T, U, F>(stream: U, handler: Arc<F>)
where
    T: ServicePair,
    U: std::io::Read + std::io::Write + Send + 'static,
    F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static,
{
    thread::spawn(move || {
        if let Err(err) = handle_request_loop::<T, U, F>(stream, &handler) {
            if !err.is_closed_connection() {
                let info = err.iter()
                    .map(|v| format!("{}", v))
                    .collect::<Vec<_>>()
                    .join("\nCaused by:");
                error!("{}", info);
            }
        }
    });
}

fn handle_request_loop<T, U, F>(mut stream: U, handler: &F) -> Result<()>
where
    T: ServicePair,
    U: std::io::Read + std::io::Write,
    F: Fn(T::Request) -> ServiceResult<T::Response>,
{
    // Receive request from client
    // Break out of loop in case of failure to read request
    while let Ok(req) = from_reader(&mut stream) {
        // Call function that handles request and returns response
        match handler(req) {
            Ok(res) => {
                // Send True flag and response in case of success
                stream.write_u8(1)?;
                to_writer(&mut stream, &res)?;
            }
            Err(message) => {
                // Send False flag and error message string in case of failure
                stream.write_u8(0)?;
                to_writer(&mut stream, &message)?;
            }
        };
    }

    // Upon failure to read request, send client failure message
    // This can be caused by actual issues or by the client stopping the connection
    stream.write_u8(0)?;
    to_writer(&mut stream, &"Failed to parse passed arguments")?;
    Ok(())
}
<|endoftext|>


<commit_before>extern crate ssz;

use self::ssz::{
    SszStream,
};
use std::sync::Arc;
use super::{
    BlockValidationContext,
    SszBlockValidationError,
    BlockStatus,
    AttesterMap,
    ProposerMap,
};
use super::db::stores::{
    BlockStore,
    PoWChainStore,
    ValidatorStore,
};
use super::db::{
    MemoryDB,
};
use super::utils::hash::canonical_hash;
use super::utils::types::{
    Hash256,
    Bitfield,
};
use super::SszBlock;
use super::super::Block;
use super::super::attestation_record::AttestationRecord;
use super::super::super::bls::{
    Keypair,
    Signature,
    AggregateSignature,
};

pub struct TestStore {
    pub db: Arc<MemoryDB>,
    pub block: Arc<BlockStore<MemoryDB>>,
    pub pow_chain: Arc<PoWChainStore<MemoryDB>>,
    pub validator: Arc<ValidatorStore<MemoryDB>>,
}

impl TestStore {
    pub fn new() -> Self {
        let db = Arc::new(MemoryDB::open());
        let block = Arc::new(BlockStore::new(db.clone()));
        let pow_chain = Arc::new(PoWChainStore::new(db.clone()));
        let validator = Arc::new(ValidatorStore::new(db.clone()));
        Self {
            db,
            block,
            pow_chain,
            validator,
        }
    }
}

#[derive(Debug)]
pub struct TestParams {
    pub total_validators: usize,
    pub cycle_length: u8,
    pub shard_count: u16,
    pub shards_per_slot: u16,
    pub validators_per_shard: usize,
    pub block_slot: u64,
    pub attestations_justified_slot: u64,
}
type ParentHashes = Vec<Hash256>;

/// Setup for a block validation function, without actually executing the
/// block validation function.
pub fn setup_block_validation_scenario(params: &TestParams)
    -> (Block, ParentHashes, AttesterMap, ProposerMap, TestStore)
{
    let stores = TestStore::new();

    let cycle_length = params.cycle_length;
    let shards_per_slot = params.shards_per_slot;
    let validators_per_shard = params.validators_per_shard;
    let block_slot = params.block_slot;
    let attestations_justified_slot = params.attestations_justified_slot;

    let parent_hashes: Vec<Hash256> = (0..(cycle_length * 2))
        .map(|i| Hash256::from(i as u64))
        .collect();
    let parent_hash = Hash256::from("parent_hash".as_bytes());
    let randao_reveal = Hash256::from("randao_reveal".as_bytes());
    let justified_block_hash = Hash256::from("justified_hash".as_bytes());
    let pow_chain_ref = Hash256::from("pow_chain".as_bytes());
    let active_state_root = Hash256::from("active_state".as_bytes());
    let crystallized_state_root = Hash256::from("cry_state".as_bytes());
    let shard_block_hash = Hash256::from("shard_block_hash".as_bytes());

    stores.pow_chain.put_block_hash(pow_chain_ref.as_ref()).unwrap();
    stores.block.put_serialized_block(justified_block_hash.as_ref(), &vec![42]).unwrap();

    /*
     * Generate a minimum viable parent block and store it in the database.
     */
    let mut parent_block = Block::zero();
    let parent_attestation = AttestationRecord::zero();
    parent_block.slot_number = block_slot - 1;
    parent_block.attestations.push(parent_attestation);
    let parent_block_ssz = serialize_block(&parent_block);
    stores.block.put_serialized_block(parent_hash.as_ref(), &parent_block_ssz).unwrap();

    let validator_index: usize = 0;
    let proposer_map = {
        let mut proposer_map = ProposerMap::new();
        proposer_map.insert(parent_block.slot_number, validator_index);
        proposer_map
    };

    let (attester_map, attestations, _keypairs) = {
        let mut i = 0;
        let attestation_slot = block_slot - 1;
        let mut attester_map = AttesterMap::new();
        let mut attestations = vec![];
        let mut keypairs = vec![];
        for shard in 0..shards_per_slot {
            let mut attesters = vec![];
            let mut attester_bitfield = Bitfield::new();
            let mut aggregate_sig = AggregateSignature::new();

            let parent_hashes_slice = {
                let distance: usize = (block_slot - attestation_slot) as usize;
                let last: usize = parent_hashes.len() - distance;
                let first: usize = last - usize::from(cycle_length);
                &parent_hashes[first..last]
            };

            let attestation_message = {
                let mut stream = SszStream::new();
                stream.append(&attestation_slot);
                stream.append_vec(&parent_hashes_slice.to_vec());
                stream.append(&shard);
                stream.append(&shard_block_hash);
                stream.append(&attestations_justified_slot);
                let bytes = stream.drain();
                canonical_hash(&bytes)
            };

            for attestation_index in 0..validators_per_shard {
               /*
                * Add the attester to the attestation indices for this shard.
                */
               attesters.push(i);
               /*
                * Set the voters bit on the bitfield to true.
                */
               attester_bitfield.set_bit(attestation_index, true);
               /*
                * Generate a random keypair for this validatior and clone it into the
                * list of keypairs.
                */
               let keypair = Keypair::random();
               keypairs.push(keypair.clone());
               /*
                * Store the validators public key in the database.
                */
               stores.validator.put_public_key_by_index(i, &keypair.pk).unwrap();
               /*
                * Generate a new signature and aggregate it on the rolling signature.
                */
               let sig = Signature::new(&attestation_message, &keypair.sk);
               aggregate_sig.add(&sig);
               /*
                * Increment the validator counter to monotonically assign validators.
                */
               i += 1;
            }

            attester_map.insert((attestation_slot, shard), attesters);
            let attestation = AttestationRecord {
                slot: attestation_slot,
                shard_id: shard,
                oblique_parent_hashes: vec![],
                shard_block_hash,
                attester_bitfield,
                justified_slot: attestations_justified_slot,
                justified_block_hash,
                aggregate_sig,
            };
            attestations.push(attestation);
        }
        (attester_map, attestations, keypairs)
    };

    let block = Block {
        parent_hash,
        slot_number: block_slot,
        randao_reveal,
        attestations,
        pow_chain_ref,
        active_state_root,
        crystallized_state_root,
    };

    (block,
     parent_hashes,
     attester_map,
     proposer_map,
     stores)
}

/// Helper function to take some Block and SSZ serialize it.
pub fn serialize_block(b: &Block) -> Vec<u8> {
    let mut stream = SszStream::new();
    stream.append(b);
    stream.drain()
}

/// Setup and run a block validation scenario, given some parameters.
///
/// Returns the Result returned from the block validation function.
pub fn run_block_validation_scenario<F>(
    validation_slot: u64,
    validation_last_justified_slot: u64,
    validation_last_finalized_slot: u64,
    params: &TestParams,
    mutator_func: F)
    -> Result<(BlockStatus, Option<Block>), SszBlockValidationError>
    where F: FnOnce(Block, AttesterMap, ProposerMap, TestStore)
                -> (Block, AttesterMap, ProposerMap, TestStore)
{
    let (block,
     parent_hashes,
     attester_map,
     proposer_map,
     stores) = setup_block_validation_scenario(&params);

    let (block,
         attester_map,
         proposer_map,
         stores) = mutator_func(block, attester_map, proposer_map, stores);

    let ssz_bytes = serialize_block(&block);
    let ssz_block = SszBlock::from_slice(&ssz_bytes[..])
        .unwrap();

    let context = BlockValidationContext {
        present_slot: validation_slot,
        cycle_length: params.cycle_length,
        last_justified_slot: validation_last_justified_slot,
        last_finalized_slot: validation_last_finalized_slot,
        parent_hashes: Arc::new(parent_hashes),
        proposer_map: Arc::new(proposer_map),
        attester_map: Arc::new(attester_map),
        block_store: stores.block.clone(),
        validator_store: stores.validator.clone(),
        pow_store: stores.pow_chain.clone()
    };
    context.validate_ssz_block(&ssz_block)
}

fn get_simple_params() -> TestParams {
    let validators_per_shard: usize = 5;
    let cycle_length: u8 = 2;
    let shard_count: u16 = 4;
    let shards_per_slot: u16 = shard_count / u16::from(cycle_length);
    let total_validators: usize = validators_per_shard * shard_count as usize;
    let block_slot = u64::from(cycle_length) * 10000;
    let attestations_justified_slot = block_slot - u64::from(cycle_length);

    TestParams {
        total_validators,
        cycle_length,
        shard_count,
        shards_per_slot,
        validators_per_shard,
        block_slot,
        attestations_justified_slot,
    }
}

#[test]
fn test_block_validation_simple_scenario_valid() {
    let params = get_simple_params();

    let validation_slot = params.block_slot;
    let validation_last_justified_slot = params.attestations_justified_slot;
    let validation_last_finalized_slot = 0;

    let no_mutate = |block, attester_map, proposer_map, stores| {
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        validation_slot,
        validation_last_justified_slot,
        validation_last_finalized_slot,
        &params,
        no_mutate);

    assert_eq!(status.unwrap().0, BlockStatus::NewBlock);
}

#[test]
fn test_block_validation_simple_scenario_invalid_unknown_parent_block() {
    let params = get_simple_params();

    let validation_slot = params.block_slot;
    let validation_last_justified_slot = params.attestations_justified_slot;
    let validation_last_finalized_slot = 0;

    let no_mutate = |mut block: Block, attester_map, proposer_map, stores| {
        block.parent_hash = Hash256::from("unknown parent block".as_bytes());
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        validation_slot,
        validation_last_justified_slot,
        validation_last_finalized_slot,
        &params,
        no_mutate);

    assert_eq!(status, Err(SszBlockValidationError::UnknownParentHash));
}

#[test]
fn test_block_validation_simple_scenario_invalid_2nd_attestation() {
    let params = get_simple_params();

    let validation_slot = params.block_slot;
    let validation_last_justified_slot = params.attestations_justified_slot;
    let validation_last_finalized_slot = 0;

    let mutator = |mut block: Block, attester_map, proposer_map, stores| {
        /*
         * Set the second attestaion record to have an invalid signature.
         */
        block.attestations[1].aggregate_sig = AggregateSignature::new();
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        validation_slot,
        validation_last_justified_slot,
        validation_last_finalized_slot,
        &params,
        mutator);

    assert_eq!(status, Err(SszBlockValidationError::AttestationSignatureFailed));
}
<commit_msg>Simplify block validation tests<commit_after>extern crate ssz;

use self::ssz::{
    SszStream,
};
use std::sync::Arc;
use super::{
    BlockValidationContext,
    SszBlockValidationError,
    BlockStatus,
    AttesterMap,
    ProposerMap,
};
use super::db::stores::{
    BlockStore,
    PoWChainStore,
    ValidatorStore,
};
use super::db::{
    MemoryDB,
};
use super::utils::hash::canonical_hash;
use super::utils::types::{
    Hash256,
    Bitfield,
};
use super::SszBlock;
use super::super::Block;
use super::super::attestation_record::AttestationRecord;
use super::super::super::bls::{
    Keypair,
    Signature,
    AggregateSignature,
};

pub struct TestStore {
    pub db: Arc<MemoryDB>,
    pub block: Arc<BlockStore<MemoryDB>>,
    pub pow_chain: Arc<PoWChainStore<MemoryDB>>,
    pub validator: Arc<ValidatorStore<MemoryDB>>,
}

impl TestStore {
    pub fn new() -> Self {
        let db = Arc::new(MemoryDB::open());
        let block = Arc::new(BlockStore::new(db.clone()));
        let pow_chain = Arc::new(PoWChainStore::new(db.clone()));
        let validator = Arc::new(ValidatorStore::new(db.clone()));
        Self {
            db,
            block,
            pow_chain,
            validator,
        }
    }
}

#[derive(Debug)]
pub struct TestParams {
    pub total_validators: usize,
    pub cycle_length: u8,
    pub shard_count: u16,
    pub shards_per_slot: u16,
    pub validators_per_shard: usize,
    pub block_slot: u64,
    pub attestations_justified_slot: u64,
    pub validation_context_slot: u64,
    pub validation_context_justified_slot: u64,
    pub validation_context_finalized_slot: u64,
}
type ParentHashes = Vec<Hash256>;

/// Setup for a block validation function, without actually executing the
/// block validation function.
pub fn setup_block_validation_scenario(params: &TestParams)
    -> (Block, ParentHashes, AttesterMap, ProposerMap, TestStore)
{
    let stores = TestStore::new();

    let cycle_length = params.cycle_length;
    let shards_per_slot = params.shards_per_slot;
    let validators_per_shard = params.validators_per_shard;
    let block_slot = params.block_slot;
    let attestations_justified_slot = params.attestations_justified_slot;

    let parent_hashes: Vec<Hash256> = (0..(cycle_length * 2))
        .map(|i| Hash256::from(i as u64))
        .collect();
    let parent_hash = Hash256::from("parent_hash".as_bytes());
    let randao_reveal = Hash256::from("randao_reveal".as_bytes());
    let justified_block_hash = Hash256::from("justified_hash".as_bytes());
    let pow_chain_ref = Hash256::from("pow_chain".as_bytes());
    let active_state_root = Hash256::from("active_state".as_bytes());
    let crystallized_state_root = Hash256::from("cry_state".as_bytes());
    let shard_block_hash = Hash256::from("shard_block_hash".as_bytes());

    stores.pow_chain.put_block_hash(pow_chain_ref.as_ref()).unwrap();
    stores.block.put_serialized_block(justified_block_hash.as_ref(), &vec![42]).unwrap();

    /*
     * Generate a minimum viable parent block and store it in the database.
     */
    let mut parent_block = Block::zero();
    let parent_attestation = AttestationRecord::zero();
    parent_block.slot_number = block_slot - 1;
    parent_block.attestations.push(parent_attestation);
    let parent_block_ssz = serialize_block(&parent_block);
    stores.block.put_serialized_block(parent_hash.as_ref(), &parent_block_ssz).unwrap();

    let validator_index: usize = 0;
    let proposer_map = {
        let mut proposer_map = ProposerMap::new();
        proposer_map.insert(parent_block.slot_number, validator_index);
        proposer_map
    };

    let (attester_map, attestations, _keypairs) = {
        let mut i = 0;
        let attestation_slot = block_slot - 1;
        let mut attester_map = AttesterMap::new();
        let mut attestations = vec![];
        let mut keypairs = vec![];
        for shard in 0..shards_per_slot {
            let mut attesters = vec![];
            let mut attester_bitfield = Bitfield::new();
            let mut aggregate_sig = AggregateSignature::new();

            let parent_hashes_slice = {
                let distance: usize = (block_slot - attestation_slot) as usize;
                let last: usize = parent_hashes.len() - distance;
                let first: usize = last - usize::from(cycle_length);
                &parent_hashes[first..last]
            };

            let attestation_message = {
                let mut stream = SszStream::new();
                stream.append(&attestation_slot);
                stream.append_vec(&parent_hashes_slice.to_vec());
                stream.append(&shard);
                stream.append(&shard_block_hash);
                stream.append(&attestations_justified_slot);
                let bytes = stream.drain();
                canonical_hash(&bytes)
            };

            for attestation_index in 0..validators_per_shard {
               /*
                * Add the attester to the attestation indices for this shard.
                */
               attesters.push(i);
               /*
                * Set the voters bit on the bitfield to true.
                */
               attester_bitfield.set_bit(attestation_index, true);
               /*
                * Generate a random keypair for this validatior and clone it into the
                * list of keypairs.
                */
               let keypair = Keypair::random();
               keypairs.push(keypair.clone());
               /*
                * Store the validators public key in the database.
                */
               stores.validator.put_public_key_by_index(i, &keypair.pk).unwrap();
               /*
                * Generate a new signature and aggregate it on the rolling signature.
                */
               let sig = Signature::new(&attestation_message, &keypair.sk);
               aggregate_sig.add(&sig);
               /*
                * Increment the validator counter to monotonically assign validators.
                */
               i += 1;
            }

            attester_map.insert((attestation_slot, shard), attesters);
            let attestation = AttestationRecord {
                slot: attestation_slot,
                shard_id: shard,
                oblique_parent_hashes: vec![],
                shard_block_hash,
                attester_bitfield,
                justified_slot: attestations_justified_slot,
                justified_block_hash,
                aggregate_sig,
            };
            attestations.push(attestation);
        }
        (attester_map, attestations, keypairs)
    };

    let block = Block {
        parent_hash,
        slot_number: block_slot,
        randao_reveal,
        attestations,
        pow_chain_ref,
        active_state_root,
        crystallized_state_root,
    };

    (block,
     parent_hashes,
     attester_map,
     proposer_map,
     stores)
}

/// Helper function to take some Block and SSZ serialize it.
pub fn serialize_block(b: &Block) -> Vec<u8> {
    let mut stream = SszStream::new();
    stream.append(b);
    stream.drain()
}

/// Setup and run a block validation scenario, given some parameters.
///
/// Returns the Result returned from the block validation function.
pub fn run_block_validation_scenario<F>(
    params: &TestParams,
    mutator_func: F)
    -> Result<(BlockStatus, Option<Block>), SszBlockValidationError>
    where F: FnOnce(Block, AttesterMap, ProposerMap, TestStore)
                -> (Block, AttesterMap, ProposerMap, TestStore)
{
    let (block,
     parent_hashes,
     attester_map,
     proposer_map,
     stores) = setup_block_validation_scenario(&params);

    let (block,
         attester_map,
         proposer_map,
         stores) = mutator_func(block, attester_map, proposer_map, stores);

    let ssz_bytes = serialize_block(&block);
    let ssz_block = SszBlock::from_slice(&ssz_bytes[..])
        .unwrap();

    let context = BlockValidationContext {
        present_slot: params.validation_context_slot,
        cycle_length: params.cycle_length,
        last_justified_slot: params.validation_context_justified_slot,
        last_finalized_slot: params.validation_context_finalized_slot,
        parent_hashes: Arc::new(parent_hashes),
        proposer_map: Arc::new(proposer_map),
        attester_map: Arc::new(attester_map),
        block_store: stores.block.clone(),
        validator_store: stores.validator.clone(),
        pow_store: stores.pow_chain.clone()
    };
    context.validate_ssz_block(&ssz_block)
}

fn get_simple_params() -> TestParams {
    let validators_per_shard: usize = 5;
    let cycle_length: u8 = 2;
    let shard_count: u16 = 4;
    let shards_per_slot: u16 = shard_count / u16::from(cycle_length);
    let total_validators: usize = validators_per_shard * shard_count as usize;
    let block_slot = u64::from(cycle_length) * 10000;
    let attestations_justified_slot = block_slot - u64::from(cycle_length);

    let validation_context_slot = block_slot;
    let validation_context_justified_slot = attestations_justified_slot;
    let validation_context_finalized_slot = 0;

    TestParams {
        total_validators,
        cycle_length,
        shard_count,
        shards_per_slot,
        validators_per_shard,
        block_slot,
        attestations_justified_slot,
        validation_context_slot,
        validation_context_justified_slot,
        validation_context_finalized_slot,
    }
}

#[test]
fn test_block_validation_simple_scenario_valid() {
    let params = get_simple_params();

    let no_mutate = |block, attester_map, proposer_map, stores| {
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        &params,
        no_mutate);

    assert_eq!(status.unwrap().0, BlockStatus::NewBlock);
}

#[test]
fn test_block_validation_simple_scenario_invalid_unknown_parent_block() {
    let params = get_simple_params();

    let no_mutate = |mut block: Block, attester_map, proposer_map, stores| {
        block.parent_hash = Hash256::from("unknown parent block".as_bytes());
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        &params,
        no_mutate);

    assert_eq!(status, Err(SszBlockValidationError::UnknownParentHash));
}

#[test]
fn test_block_validation_simple_scenario_invalid_2nd_attestation() {
    let params = get_simple_params();

    let mutator = |mut block: Block, attester_map, proposer_map, stores| {
        /*
         * Set the second attestaion record to have an invalid signature.
         */
        block.attestations[1].aggregate_sig = AggregateSignature::new();
        (block, attester_map, proposer_map, stores)
    };

    let status = run_block_validation_scenario(
        &params,
        mutator);

    assert_eq!(status, Err(SszBlockValidationError::AttestationSignatureFailed));
}
<|endoftext|>