<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 https://mozilla.org/MPL/2.0/.

//! This module provides function to "clean" a text typographically.
//!
//! # Example
//!
//! ```
//! use crowbook_text_processing::clean;
//! let input = "Some  'text'  whose formatting  could be enhanced...";
//! let output = clean::quotes(clean::ellipsis(clean::whitespaces(input)));
//! assert_eq!(&output, "Some ‘text’ whose formatting could be enhanced…");
//! ```

use common::is_whitespace;

use regex::Regex;
use std::borrow::Cow;

/// Removes unnecessary whitespaces from a String.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::whitespaces("  A  string   with   more   whitespaces  than  needed   ");
/// assert_eq!(&s, " A string with more whitespaces than needed ");
/// ```
pub fn whitespaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"[  \x{202F}\x{2002}]{2,}?").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        new_s.push_str(&input[0..first]);
        let mut previous_space = false;
        for c in input[first..].chars() {
            if is_whitespace(c) {
                if previous_space {
                    // previous char already a space, don't copy it
                } else {
                    new_s.push(c);
                    previous_space = true;
                }
            } else {
                previous_space = false;
                new_s.push(c);
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

/// Class of a character
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
enum CharClass {
    Whitespace = 0,
    Punctuation,
    Alphanumeric,
}

/// Get class of a character
fn char_class(c: char) -> CharClass {
    if c.is_alphanumeric() {
        CharClass::Alphanumeric
    } else if c.is_whitespace() {
        CharClass::Whitespace
    } else {
        CharClass::Punctuation
    }
}

/// Replace ellipsis (...) with the appropriate unicode character
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::ellipsis("foo...");
/// assert_eq!(&s, "foo…");
/// let s = clean::ellipsis("foo. . . ");
/// assert_eq!(&s, "foo.\u{a0}.\u{a0}. "); // non breaking spaces
/// ```
pub fn ellipsis<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"\.\.\.|\. \. \. ").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut output: Vec<u8> = Vec::with_capacity(input.len());
        output.extend_from_slice(input[0..first].as_bytes());
        let rest = input[first..].bytes().collect::<Vec<_>>();
        let len = rest.len();
        let mut i = 0;
        while i < len {
            if i + 3 <= len && &rest[i..(i + 3)] == &[b'.', b'.', b'.'] {
                output.extend_from_slice("…".as_bytes());
                i += 3;
            } else if i + 6 <= len && &rest[i..(i + 6)] == &[b'.', b' ', b'.', b' ', b'.', b' '] {
                if i + 6 == len || rest[i + 6] != b'.' {
                    output.extend_from_slice(". . . ".as_bytes());
                } else {
                    output.extend_from_slice(". . . ".as_bytes());
                }
                i += 6;
            } else {
                output.push(rest[i]);
                i += 1;
            }
        }
        Cow::Owned(String::from_utf8(output).unwrap())
    } else {
        input
    }
}

/// Replace straight quotes with more typographic variants
///
/// While it should work pretty well for double quotes (`"`), the rules for single
/// quote (`'`) are more ambiguous, as it can be a quote or an apostrophe and it's not
/// that easy (and, in some circumstances, impossible without understanding the meaning
/// of the text) to get right.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::quotes("\"foo\"");
/// assert_eq!(&s, "“foo”");
/// let s = clean::quotes("'foo'");
/// assert_eq!(&s, "‘foo’");
/// ```
pub fn quotes<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new("[\"\']").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((mut first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        if first > 0 {
            // Move one step backward since we might need to know if previous char was
            // a letter or not
            first -= 1;
        }
        new_s.push_str(&input[0..first]);
        let mut chars = input[first..].chars().collect::<Vec<_>>();
        let mut closing_quote = None;
        let mut opened_doubles = 0;
        for i in 0..chars.len() {
            let c = chars[i];
            let has_opened_quote = if let Some(n) = closing_quote {
                i <= n
            } else {
                false
            };
            match c {
                '"' => {
                    let prev = if i > 0 {
                        char_class(chars[i - 1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i + 1])
                    } else {
                        CharClass::Whitespace
                    };

                    if prev < next {
                        opened_doubles += 1;
                        new_s.push('“');
                    } else if opened_doubles > 0 {
                        opened_doubles -= 1;
                        new_s.push('”');
                    } else {
                        new_s.push('"');
                    }
                }
                '\'' => {
                    let prev = if i > 0 {
                        char_class(chars[i - 1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i + 1])
                    } else {
                        CharClass::Whitespace
                    };

                    let replacement = match (prev, next) {
                        // Elision or possessive
                        (CharClass::Alphanumeric, CharClass::Alphanumeric)
//                            | (CharClass::Punctuation, CharClass::Alphanumeric)
                            => '’',

                        // Beginning of word, it's opening (not always though)
                        (x, y) if x < y
                            => {
                                let mut is_next_closing = false;
                                for j in (i + 1)..chars.len() {
                                    if chars[j] == '\'' {
                                        if chars[j-1].is_whitespace() {
                                            continue;
                                        } else if j >= chars.len() - 1
                                            || char_class(chars[j+1]) != CharClass::Alphanumeric {
                                                is_next_closing = true;
                                                closing_quote = Some(j);
                                                chars[j] = '’';
                                                break;
                                            }
                                    }
                                }
                                if is_next_closing && !has_opened_quote {
                                    '‘'
                                } else {
                                    '’'
                                }
                            }

                        // Apostrophe at end of word, it's closing
                        (x, y) if x > y
                            => {
                                '’'
                            },
                        _ => '\'',
                    };
                    new_s.push(replacement);
                    // if i > 0 && !chars[i - 1].is_whitespace() {
                    //     new_s.push('’');
                    // } else if i < chars.len() - 1 && !chars[i + 1].is_whitespace() {
                    //     new_s.push('‘');
                    // } else {
                    //     new_s.push('\'');
                    // }
                }
                _ => new_s.push(c),
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

#[test]
fn whitespaces_1() {
    let s = "   Remove    supplementary   spaces    but    don't    trim     either   ";
    let res = whitespaces(s);
    assert_eq!(&res, " Remove supplementary spaces but don't trim either ");
}

#[test]
fn quotes_1() {
    let s = "Some string without ' typographic ' quotes";
    let res = quotes(s);
    assert_eq!(&res, s);
}

#[test]
fn quotes_2() {
    let s = quotes("\"foo\"");
    assert_eq!(&s, "“foo”");
    let s = quotes("'foo'");
    assert_eq!(&s, "‘foo’");
}

#[test]
fn quotes_3() {
    let s = quotes("\'mam, how are you?");
    assert_eq!(&s, "’mam, how are you?");
}

#[test]
fn quotes_4() {
    let s = quotes("some char: 'c', '4', '&'");
    assert_eq!(&s, "some char: ‘c’, ‘4’, ‘&’");
}

#[test]
fn quotes_5() {
    let s = quotes("It's a good day to say 'hi'");
    assert_eq!(&s, "It’s a good day to say ‘hi’");
}

#[test]
fn quotes_6() {
    let s = quotes("The '60s were nice, weren't they?");
    assert_eq!(&s, "The ’60s were nice, weren’t they?");
}

#[test]
fn quotes_7() {
    let s = quotes("Plurals' possessive");
    assert_eq!(&s, "Plurals’ possessive");
}

#[test]
fn quotes_8() {
    let s = quotes("\"I like 'That '70s show'\", she said");
    assert_eq!(&s, "“I like ‘That ’70s show’”, she said");
}

#[test]
fn quotes_9() {
    let s = quotes("some char: '!', '?', ','");
    assert_eq!(&s, "some char: ‘!’, ‘?’, ‘,’");
}

#[test]
fn quotes_10() {
    let s = quotes("\"'Let's try \"nested\" quotes,' he said.\"");
    assert_eq!(&s, "“‘Let’s try “nested” quotes,’ he said.”");
}

#[test]
fn quotes_11() {
    let s = quotes("Enhanced \"quotes\"'s heuristics");
    assert_eq!(&s, "Enhanced “quotes”’s heuristics");
}

#[test]
fn ellipsis_0() {
    let s = ellipsis("Foo...");
    assert_eq!(&s, "Foo…");
}

#[test]
fn ellipsis_1() {
    let s = ellipsis("Foo... Bar");
    assert_eq!(&s, "Foo… Bar");
}

#[test]
fn ellipsis_2() {
    let s = ellipsis("foo....");
    assert_eq!(&s, "foo….");
}

#[test]
fn ellipsis_3() {
    let s = ellipsis("foo. . . ");
    assert_eq!(&s, "foo. . . ");
}

#[test]
fn ellipsis_4() {
    let s = ellipsis("foo. . . .");
    assert_eq!(&s, "foo. . . .");
}
<commit_msg>Add clean::dashes<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 https://mozilla.org/MPL/2.0/.

//! This module provides function to "clean" a text typographically.
//!
//! # Example
//!
//! ```
//! use crowbook_text_processing::clean;
//! let input = "Some  'text'  whose formatting  could be enhanced...";
//! let output = clean::quotes(clean::ellipsis(clean::whitespaces(input)));
//! assert_eq!(&output, "Some ‘text’ whose formatting could be enhanced…");
//! ```

use common::is_whitespace;

use regex::Regex;
use std::borrow::Cow;

/// Removes unnecessary whitespaces from a String.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::whitespaces("  A  string   with   more   whitespaces  than  needed   ");
/// assert_eq!(&s, " A string with more whitespaces than needed ");
/// ```
pub fn whitespaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"[  \x{202F}\x{2002}]{2,}?").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        new_s.push_str(&input[0..first]);
        let mut previous_space = false;
        for c in input[first..].chars() {
            if is_whitespace(c) {
                if previous_space {
                    // previous char already a space, don't copy it
                } else {
                    new_s.push(c);
                    previous_space = true;
                }
            } else {
                previous_space = false;
                new_s.push(c);
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

/// Class of a character
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
enum CharClass {
    Whitespace = 0,
    Punctuation,
    Alphanumeric,
}

/// Get class of a character
fn char_class(c: char) -> CharClass {
    if c.is_alphanumeric() {
        CharClass::Alphanumeric
    } else if c.is_whitespace() {
        CharClass::Whitespace
    } else {
        CharClass::Punctuation
    }
}

/// Replace ellipsis (...) with the appropriate unicode character
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::ellipsis("foo...");
/// assert_eq!(&s, "foo…");
/// let s = clean::ellipsis("foo. . . ");
/// assert_eq!(&s, "foo.\u{a0}.\u{a0}. "); // non breaking spaces
/// ```
pub fn ellipsis<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"\.\.\.|\. \. \. ").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut output: Vec<u8> = Vec::with_capacity(input.len());
        output.extend_from_slice(input[0..first].as_bytes());
        let rest = input[first..].bytes().collect::<Vec<_>>();
        let len = rest.len();
        let mut i = 0;
        while i < len {
            if i + 3 <= len && &rest[i..(i + 3)] == &[b'.', b'.', b'.'] {
                output.extend_from_slice("…".as_bytes());
                i += 3;
            } else if i + 6 <= len && &rest[i..(i + 6)] == &[b'.', b' ', b'.', b' ', b'.', b' '] {
                if i + 6 == len || rest[i + 6] != b'.' {
                    output.extend_from_slice(". . . ".as_bytes());
                } else {
                    output.extend_from_slice(". . . ".as_bytes());
                }
                i += 6;
            } else {
                output.push(rest[i]);
                i += 1;
            }
        }
        Cow::Owned(String::from_utf8(output).unwrap())
    } else {
        input
    }
}

/// Replace straight quotes with more typographic variants
///
/// While it should work pretty well for double quotes (`"`), the rules for single
/// quote (`'`) are more ambiguous, as it can be a quote or an apostrophe and it's not
/// that easy (and, in some circumstances, impossible without understanding the meaning
/// of the text) to get right.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::quotes("\"foo\"");
/// assert_eq!(&s, "“foo”");
/// let s = clean::quotes("'foo'");
/// assert_eq!(&s, "‘foo’");
/// ```
pub fn quotes<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new("[\"\']").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((mut first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        if first > 0 {
            // Move one step backward since we might need to know if previous char was
            // a letter or not
            first -= 1;
        }
        new_s.push_str(&input[0..first]);
        let mut chars = input[first..].chars().collect::<Vec<_>>();
        let mut closing_quote = None;
        let mut opened_doubles = 0;
        for i in 0..chars.len() {
            let c = chars[i];
            let has_opened_quote = if let Some(n) = closing_quote {
                i <= n
            } else {
                false
            };
            match c {
                '"' => {
                    let prev = if i > 0 {
                        char_class(chars[i - 1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i + 1])
                    } else {
                        CharClass::Whitespace
                    };

                    if prev < next {
                        opened_doubles += 1;
                        new_s.push('“');
                    } else if opened_doubles > 0 {
                        opened_doubles -= 1;
                        new_s.push('”');
                    } else {
                        new_s.push('"');
                    }
                }
                '\'' => {
                    let prev = if i > 0 {
                        char_class(chars[i - 1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i + 1])
                    } else {
                        CharClass::Whitespace
                    };

                    let replacement = match (prev, next) {
                        // Elision or possessive
                        (CharClass::Alphanumeric, CharClass::Alphanumeric)
//                            | (CharClass::Punctuation, CharClass::Alphanumeric)
                            => '’',

                        // Beginning of word, it's opening (not always though)
                        (x, y) if x < y
                            => {
                                let mut is_next_closing = false;
                                for j in (i + 1)..chars.len() {
                                    if chars[j] == '\'' {
                                        if chars[j-1].is_whitespace() {
                                            continue;
                                        } else if j >= chars.len() - 1
                                            || char_class(chars[j+1]) != CharClass::Alphanumeric {
                                                is_next_closing = true;
                                                closing_quote = Some(j);
                                                chars[j] = '’';
                                                break;
                                            }
                                    }
                                }
                                if is_next_closing && !has_opened_quote {
                                    '‘'
                                } else {
                                    '’'
                                }
                            }

                        // Apostrophe at end of word, it's closing
                        (x, y) if x > y
                            => {
                                '’'
                            },
                        _ => '\'',
                    };
                    new_s.push(replacement);
                    // if i > 0 && !chars[i - 1].is_whitespace() {
                    //     new_s.push('’');
                    // } else if i < chars.len() - 1 && !chars[i + 1].is_whitespace() {
                    //     new_s.push('‘');
                    // } else {
                    //     new_s.push('\'');
                    // }
                }
                _ => new_s.push(c),
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

/// Replace double dashes (`--`) and triple dashes (`---`) to en dash and em dash, respectively.
///
/// This function can be useful when writing literaty texts, but should be used with caution
/// as double and triple dashes can have special meanings.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean;
/// let s = clean::dashes("--- Hi, he said -- unexpectedly");
/// assert_eq!(&s, "— Hi, he said – unexpectedly");
/// ```
pub fn dashes<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"\x2D\x2D").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut output: Vec<u8> = Vec::with_capacity(input.len());
        output.extend_from_slice(input[0..first].as_bytes());
        let rest = input[first..].bytes().collect::<Vec<_>>();
        let len = rest.len();
        let mut i = 0;
        println!("found first dash at {}", first);
        while i < len {
            if i + 2 <= len && &rest[i..(i + 2)] == &[b'-', b'-'] {
                if i + 2 < len && rest[i + 2] == b'-' {
                    output.extend_from_slice("—".as_bytes());
                    i += 3;
                } else {
                    output.extend_from_slice("–".as_bytes());
                    i += 2;
                }
            } else {
                output.push(rest[i]);
                i += 1;
            }
        }
        Cow::Owned(String::from_utf8(output).unwrap())
    } else {
        input
    }
}

#[test]
fn whitespaces_1() {
    let s = "   Remove    supplementary   spaces    but    don't    trim     either   ";
    let res = whitespaces(s);
    assert_eq!(&res, " Remove supplementary spaces but don't trim either ");
}

#[test]
fn quotes_1() {
    let s = "Some string without ' typographic ' quotes";
    let res = quotes(s);
    assert_eq!(&res, s);
}

#[test]
fn quotes_2() {
    let s = quotes("\"foo\"");
    assert_eq!(&s, "“foo”");
    let s = quotes("'foo'");
    assert_eq!(&s, "‘foo’");
}

#[test]
fn quotes_3() {
    let s = quotes("\'mam, how are you?");
    assert_eq!(&s, "’mam, how are you?");
}

#[test]
fn quotes_4() {
    let s = quotes("some char: 'c', '4', '&'");
    assert_eq!(&s, "some char: ‘c’, ‘4’, ‘&’");
}

#[test]
fn quotes_5() {
    let s = quotes("It's a good day to say 'hi'");
    assert_eq!(&s, "It’s a good day to say ‘hi’");
}

#[test]
fn quotes_6() {
    let s = quotes("The '60s were nice, weren't they?");
    assert_eq!(&s, "The ’60s were nice, weren’t they?");
}

#[test]
fn quotes_7() {
    let s = quotes("Plurals' possessive");
    assert_eq!(&s, "Plurals’ possessive");
}

#[test]
fn quotes_8() {
    let s = quotes("\"I like 'That '70s show'\", she said");
    assert_eq!(&s, "“I like ‘That ’70s show’”, she said");
}

#[test]
fn quotes_9() {
    let s = quotes("some char: '!', '?', ','");
    assert_eq!(&s, "some char: ‘!’, ‘?’, ‘,’");
}

#[test]
fn quotes_10() {
    let s = quotes("\"'Let's try \"nested\" quotes,' he said.\"");
    assert_eq!(&s, "“‘Let’s try “nested” quotes,’ he said.”");
}

#[test]
fn quotes_11() {
    let s = quotes("Enhanced \"quotes\"'s heuristics");
    assert_eq!(&s, "Enhanced “quotes”’s heuristics");
}

#[test]
fn ellipsis_0() {
    let s = ellipsis("Foo...");
    assert_eq!(&s, "Foo…");
}

#[test]
fn ellipsis_1() {
    let s = ellipsis("Foo... Bar");
    assert_eq!(&s, "Foo… Bar");
}

#[test]
fn ellipsis_2() {
    let s = ellipsis("foo....");
    assert_eq!(&s, "foo….");
}

#[test]
fn ellipsis_3() {
    let s = ellipsis("foo. . . ");
    assert_eq!(&s, "foo. . . ");
}

#[test]
fn ellipsis_4() {
    let s = ellipsis("foo. . . .");
    assert_eq!(&s, "foo. . . .");
}

#[test]
fn ellipsis_5() {
    let s = ellipsis("foo..");
    assert_eq!(&s, "foo..");
}

#[test]
fn dashes_0() {
    let s = dashes("foo - bar");
    assert_eq!(&s, "foo - bar");
}

#[test]
fn dashes_1() {
    let s = dashes("foo -- bar");
    assert_eq!(&s, "foo – bar");
}

#[test]
fn dashes_2() {
    let s = dashes("foo --- bar");
    assert_eq!(&s, "foo — bar");
}

#[test]
fn dashes_3() {
    let s = dashes("foo --- bar--");
    assert_eq!(&s, "foo — bar–");
}

<|endoftext|>


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

//! Implements the frame layer as defined in /experimental/oak_baremetal_channel/SPEC.MD

extern crate alloc;

use crate::Channel;
use alloc::{boxed::Box, format, vec, vec::Vec};
use bitflags::bitflags;

pub const PADDING_SIZE: usize = 4;

type Length = u16;
pub const LENGTH_SIZE: usize = 2;
static_assertions::assert_eq_size!([u8; LENGTH_SIZE], Length);

bitflags! {
    #[derive(Default)]
    pub struct Flags: u16 {
        const START = 1;
        const END = 2;
    }
}
pub const FLAGS_SIZE: usize = 2;
static_assertions::assert_eq_size!([u8; FLAGS_SIZE], Flags);

pub const BODY_OFFSET: usize = 8;
static_assertions::assert_eq_size!(
    [u8; BODY_OFFSET],
    [u8; PADDING_SIZE + LENGTH_SIZE + FLAGS_SIZE]
);

pub const MAX_SIZE: usize = 4000;
pub const MAX_BODY_SIZE: usize = MAX_SIZE - BODY_OFFSET;

/// Rust implementation of the Frame structure defined in
/// /experimental/oak_baremetal_channel/SPEC.MD
#[derive(Clone, Default, Debug)]
pub struct Frame {
    pub flags: Flags,
    pub body: Vec<u8>,
}

impl TryFrom<Frame> for Vec<u8> {
    type Error = anyhow::Error;
    fn try_from(frame: Frame) -> Result<Self, Self::Error> {
        let length = BODY_OFFSET + frame.body.len();
        let mut frame_bytes: Vec<u8> = Vec::with_capacity(length);

        frame_bytes.extend_from_slice(&[0; PADDING_SIZE]);
        {
            let frame_length = Length::try_from(length).map_err(|_error| {
                anyhow::Error::msg(format!(
                    "could convert the frame length usize to {:?}",
                    core::any::type_name::<Length>()
                ))
            })?;
            frame_bytes.extend_from_slice(&frame_length.to_le_bytes());
        }
        frame_bytes.extend_from_slice(&frame.flags.bits.to_le_bytes());
        frame_bytes.extend_from_slice(&frame.body);

        Ok(frame_bytes)
    }
}

pub struct Framed {
    inner: Box<dyn Channel>,
}

impl Framed {
    pub fn new(socket: Box<dyn Channel>) -> Self {
        Self { inner: socket }
    }

    pub fn read_frame(&mut self) -> anyhow::Result<Frame> {
        {
            let mut padding_bytes = [0; PADDING_SIZE];
            self.inner.read(&mut padding_bytes)?;
        };
        let length: usize = {
            let mut length_bytes = [0; LENGTH_SIZE];
            self.inner.read(&mut length_bytes)?;
            let length = Length::from_le_bytes(length_bytes).into();
            if length > MAX_SIZE {
                return Err(anyhow::Error::msg("frame exceeds the maximum frame size"));
            };
            length
        };
        let flags = {
            let mut flags_bytes = [0; FLAGS_SIZE];
            self.inner.read(&mut flags_bytes)?;
            Flags::from_bits_truncate(u16::from_le_bytes(flags_bytes))
        };

        let body = {
            let body_length: usize = length - BODY_OFFSET;
            let mut body: Vec<u8> = vec![0; body_length];
            self.inner.read(&mut body)?;
            body
        };

        Ok(Frame { flags, body })
    }

    pub fn write_frame(&mut self, frame: Frame) -> anyhow::Result<()> {
        let frame_bytes: Vec<u8> = frame.try_into()?;
        self.inner.write(&frame_bytes)?;
        self.inner.flush()?;
        Ok(())
    }
}

pub fn bytes_into_frames(data: Vec<u8>) -> anyhow::Result<Vec<Frame>> {
    if data.is_empty() {
        anyhow::bail!("cannot convert empty payloads into frames")
    }

    let mut frames: Vec<Frame> = data
        .chunks(MAX_BODY_SIZE)
        // TODO(#2848): It'd be nice if we didn't have to reallocate here.
        // We may be able to avoid doing so by making [`Frame`] use
        // slices and lifetimes. Or alternatively use the Bytes crate for
        // reference counting.
        .map(|frame_body| Frame {
            flags: Flags::default(),
            body: frame_body.to_vec(),
        })
        .collect();

    frames
        .first_mut()
        .expect("there should always be at least one frame")
        .flags
        .set(Flags::START, true);

    frames
        .last_mut()
        .expect("there should always be at least one frame")
        .flags
        .set(Flags::END, true);

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

//! Implements the frame layer as defined in /experimental/oak_baremetal_channel/SPEC.MD

extern crate alloc;

use crate::Channel;
use alloc::{boxed::Box, format, vec, vec::Vec};
use bitflags::bitflags;

pub const PADDING_SIZE: usize = 4;

type Length = u16;
pub const LENGTH_SIZE: usize = 2;
static_assertions::assert_eq_size!([u8; LENGTH_SIZE], Length);

bitflags! {
    #[derive(Default)]
    pub struct Flags: u16 {
        const START = 1;
        const END = 2;
    }
}
pub const FLAGS_SIZE: usize = 2;
static_assertions::assert_eq_size!([u8; FLAGS_SIZE], Flags);

pub const BODY_OFFSET: usize = 8;
static_assertions::assert_eq_size!(
    [u8; BODY_OFFSET],
    [u8; PADDING_SIZE + LENGTH_SIZE + FLAGS_SIZE]
);

pub const MAX_SIZE: usize = 4000;
pub const MAX_BODY_SIZE: usize = MAX_SIZE - BODY_OFFSET;

/// Rust implementation of the Frame structure defined in
/// /experimental/oak_baremetal_channel/SPEC.MD
#[derive(Clone, Default, Debug)]
pub struct Frame {
    pub flags: Flags,
    pub body: Vec<u8>,
}

impl TryFrom<Frame> for Vec<u8> {
    type Error = anyhow::Error;
    fn try_from(frame: Frame) -> Result<Self, Self::Error> {
        let length = BODY_OFFSET + frame.body.len();
        let mut frame_bytes: Vec<u8> = Vec::with_capacity(length);

        frame_bytes.extend_from_slice(&[0; PADDING_SIZE]);
        {
            let frame_length = Length::try_from(length).map_err(|_error| {
                anyhow::Error::msg(format!(
                    "could convert the frame length usize to {:?}",
                    core::any::type_name::<Length>()
                ))
            })?;
            frame_bytes.extend_from_slice(&frame_length.to_le_bytes());
        }
        frame_bytes.extend_from_slice(&frame.flags.bits.to_le_bytes());
        frame_bytes.extend_from_slice(&frame.body);

        Ok(frame_bytes)
    }
}

pub struct Framed {
    inner: Box<dyn Channel>,
}

impl Framed {
    pub fn new(socket: Box<dyn Channel>) -> Self {
        Self { inner: socket }
    }

    pub fn read_frame(&mut self) -> anyhow::Result<Frame> {
        {
            let mut padding_bytes = [0; PADDING_SIZE];
            self.inner.read(&mut padding_bytes)?;
        };
        let length: usize = {
            let mut length_bytes = [0; LENGTH_SIZE];
            self.inner.read(&mut length_bytes)?;
            let length = Length::from_le_bytes(length_bytes).into();
            if length <= BODY_OFFSET {
                return Err(anyhow::Error::msg("frame is too small"));
            };
            if length > MAX_SIZE {
                return Err(anyhow::Error::msg("frame exceeds the maximum frame size"));
            };
            length
        };
        let flags = {
            let mut flags_bytes = [0; FLAGS_SIZE];
            self.inner.read(&mut flags_bytes)?;
            Flags::from_bits_truncate(u16::from_le_bytes(flags_bytes))
        };

        let body = {
            let body_length: usize = length - BODY_OFFSET;
            let mut body: Vec<u8> = vec![0; body_length];
            self.inner.read(&mut body)?;
            body
        };

        Ok(Frame { flags, body })
    }

    pub fn write_frame(&mut self, frame: Frame) -> anyhow::Result<()> {
        let frame_bytes: Vec<u8> = frame.try_into()?;
        self.inner.write(&frame_bytes)?;
        self.inner.flush()?;
        Ok(())
    }
}

pub fn bytes_into_frames(data: Vec<u8>) -> anyhow::Result<Vec<Frame>> {
    if data.is_empty() {
        anyhow::bail!("cannot convert empty payloads into frames")
    }

    let mut frames: Vec<Frame> = data
        .chunks(MAX_BODY_SIZE)
        // TODO(#2848): It'd be nice if we didn't have to reallocate here.
        // We may be able to avoid doing so by making [`Frame`] use
        // slices and lifetimes. Or alternatively use the Bytes crate for
        // reference counting.
        .map(|frame_body| Frame {
            flags: Flags::default(),
            body: frame_body.to_vec(),
        })
        .collect();

    frames
        .first_mut()
        .expect("there should always be at least one frame")
        .flags
        .set(Flags::START, true);

    frames
        .last_mut()
        .expect("there should always be at least one frame")
        .flags
        .set(Flags::END, true);

    Ok(frames)
}
<|endoftext|>


<commit_before>use super::{BeaconChain, ClientDB, DBError, SlotClock};
use bls::{AggregatePublicKey, AggregateSignature, PublicKey, Signature};
use boolean_bitfield::BooleanBitfield;
use hashing::hash;
use slot_clock::{SystemTimeSlotClockError, TestingSlotClockError};
use ssz::ssz_encode;
use types::{
    beacon_state::SlotProcessingError, readers::BeaconBlockReader, AttestationData,
    AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Exit, Fork, Hash256,
    PendingAttestation,
};

// TODO: define elsehwere.
const DOMAIN_ATTESTATION: u64 = 1;
const DOMAIN_PROPOSAL: u64 = 2;
const DOMAIN_EXIT: u64 = 3;
const DOMAIN_RANDAO: u64 = 4;

macro_rules! ensure {
    ($condition: expr, $result: expr) => {
        if !$condition {
            return Err($result);
        }
    };
}

#[derive(Debug, PartialEq)]
pub enum Error {
    DBError(String),
    StateAlreadyTransitioned,
    NotImplemented,
    PresentSlotIsNone,
    UnableToDecodeBlock,
    MissingParentState(Hash256),
    InvalidParentState(Hash256),
    MissingBeaconBlock(Hash256),
    InvalidBeaconBlock(Hash256),
    MissingParentBlock(Hash256),
    NoBlockProducer,
    StateSlotMismatch,
    BadBlockSignature,
    BadRandaoSignature,
    MaxProposerSlashingsExceeded,
    BadProposerSlashing,
    MaxAttestationsExceeded,
    BadAttestation,
    NoBlockRoot,
    MaxDepositsExceeded,
    MaxExitsExceeded,
    BadExit,
    BadCustodyReseeds,
    BadCustodyChallenges,
    BadCustodyResponses,
    SlotClockError(SystemTimeSlotClockError),
}

impl<T, U> BeaconChain<T, U>
where
    T: ClientDB,
    U: SlotClock,
{
    pub fn state_transition(
        &self,
        state: BeaconState,
        block: &BeaconBlock,
    ) -> Result<BeaconState, Error> {
        self.internal_state_transition(state, block, true)
    }

    pub fn state_transition_without_verifying_block_signature(
        &self,
        state: BeaconState,
        block: &BeaconBlock,
    ) -> Result<BeaconState, Error> {
        self.internal_state_transition(state, block, false)
    }

    fn internal_state_transition(
        &self,
        mut state: BeaconState,
        block: &BeaconBlock,
        verify_block_signature: bool,
    ) -> Result<BeaconState, Error> {
        ensure!(state.slot < block.slot, Error::StateAlreadyTransitioned);

        for _ in state.slot..block.slot {
            state.per_slot_processing(block.parent_root.clone(), &self.spec)?;
        }

        /*
         * Slot
         */

        ensure!(block.slot() == state.slot, Error::StateSlotMismatch);

        /*
         * Proposer Signature
         */

        let block_proposer_index = state
            .get_beacon_proposer_index(block.slot, &self.spec)
            .ok_or(Error::NoBlockProducer)?;
        let block_proposer = &state.validator_registry[block_proposer_index];

        if verify_block_signature {
            ensure!(
                bls_verify(
                    &block_proposer.pubkey,
                    &block.proposal_root(&self.spec)[..],
                    &block.signature,
                    get_domain(&state.fork_data, state.slot, DOMAIN_PROPOSAL)
                ),
                Error::BadBlockSignature
            );
        }

        /*
         * RANDAO
         */

        ensure!(
            bls_verify(
                &block_proposer.pubkey,
                &ssz_encode(&block_proposer.proposer_slots),
                &block.randao_reveal,
                get_domain(&state.fork_data, state.slot, DOMAIN_RANDAO)
            ),
            Error::BadRandaoSignature
        );

        // TODO: check this is correct.
        let new_mix = {
            let mut mix = state.latest_randao_mixes
                [(state.slot % self.spec.latest_randao_mixes_length) as usize]
                .to_vec();
            mix.append(&mut ssz_encode(&block.randao_reveal));
            Hash256::from(&hash(&mix)[..])
        };

        state.latest_randao_mixes[(state.slot % self.spec.latest_randao_mixes_length) as usize] =
            new_mix;

        /*
         * Eth1 data
         */

        // TODO: Eth1 data stuff.

        /*
         * OPERATIONS
         */

        /*
         * Proposer slashings
         */

        ensure!(
            block.body.proposer_slashings.len() as u64 <= self.spec.max_proposer_slashings,
            Error::MaxProposerSlashingsExceeded
        );
        for proposer_slashing in &block.body.proposer_slashings {
            let proposer = state
                .validator_registry
                .get(proposer_slashing.proposer_index as usize)
                .ok_or(Error::BadProposerSlashing)?;
            ensure!(
                proposer_slashing.proposal_data_1.slot == proposer_slashing.proposal_data_2.slot,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer_slashing.proposal_data_1.shard == proposer_slashing.proposal_data_2.shard,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer_slashing.proposal_data_1.block_root
                    != proposer_slashing.proposal_data_2.block_root,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer.penalized_slot > state.slot,
                Error::BadProposerSlashing
            );
            ensure!(
                bls_verify(
                    &proposer.pubkey,
                    &hash_tree_root(&proposer_slashing.proposal_data_1),
                    &proposer_slashing.proposal_signature_1,
                    get_domain(
                        &state.fork_data,
                        proposer_slashing.proposal_data_1.slot,
                        DOMAIN_PROPOSAL
                    )
                ),
                Error::BadProposerSlashing
            );
            ensure!(
                bls_verify(
                    &proposer.pubkey,
                    &hash_tree_root(&proposer_slashing.proposal_data_2),
                    &proposer_slashing.proposal_signature_2,
                    get_domain(
                        &state.fork_data,
                        proposer_slashing.proposal_data_2.slot,
                        DOMAIN_PROPOSAL
                    )
                ),
                Error::BadProposerSlashing
            );
            penalize_validator(&state, proposer_slashing.proposer_index as usize);
        }

        /*
         * Attestations
         */
        ensure!(
            block.body.attestations.len() as u64 <= self.spec.max_attestations,
            Error::MaxAttestationsExceeded
        );

        for attestation in &block.body.attestations {
            ensure!(
                attestation.data.slot + self.spec.min_attestation_inclusion_delay <= state.slot,
                Error::BadAttestation
            );
            ensure!(
                attestation.data.slot + self.spec.epoch_length >= state.slot,
                Error::BadAttestation
            );
            if state.justified_slot >= state.slot - (state.slot % self.spec.epoch_length) {
                ensure!(
                    attestation.data.justified_slot == state.justified_slot,
                    Error::BadAttestation
                );
            } else {
                ensure!(
                    attestation.data.justified_slot == state.previous_justified_slot,
                    Error::BadAttestation
                );
            }
            ensure!(
                attestation.data.justified_block_root
                    == *get_block_root(
                        &state,
                        attestation.data.justified_slot,
                        self.spec.latest_block_roots_length
                    )
                    .ok_or(Error::NoBlockRoot)?,
                Error::BadAttestation
            );
            ensure!(
                (attestation.data.latest_crosslink_root
                    == state.latest_crosslinks[attestation.data.shard as usize].shard_block_root)
                    || (attestation.data.shard_block_root
                        == state.latest_crosslinks[attestation.data.shard as usize]
                            .shard_block_root),
                Error::BadAttestation
            );
            let participants = get_attestation_participants(
                &state,
                &attestation.data,
                &attestation.aggregation_bitfield,
            );
            let mut group_public_key = AggregatePublicKey::new();
            for participant in participants {
                group_public_key.add(
                    state.validator_registry[participant as usize]
                        .pubkey
                        .as_raw(),
                )
            }
            let attestation_message = {
                let attestation_data_and_custody_bit = AttestationDataAndCustodyBit {
                    data: attestation.data.clone(),
                    custody_bit: false,
                };
                hash_tree_root(&attestation_data_and_custody_bit).to_vec()
            };
            // Signature verification.
            ensure!(
                bls_verify_aggregate(
                    &group_public_key,
                    &attestation_message[..],
                    &attestation.aggregate_signature,
                    get_domain(&state.fork_data, attestation.data.slot, DOMAIN_ATTESTATION)
                ),
                Error::BadProposerSlashing
            );
            ensure!(
                attestation.data.shard_block_root == self.spec.zero_hash,
                Error::BadAttestation
            );
            let pending_attestation = PendingAttestation {
                data: attestation.data.clone(),
                aggregation_bitfield: attestation.aggregation_bitfield.clone(),
                custody_bitfield: attestation.custody_bitfield.clone(),
                slot_included: state.slot,
            };
            state.latest_attestations.push(pending_attestation);
        }

        /*
         * Deposits
         */
        ensure!(
            block.body.deposits.len() as u64 <= self.spec.max_deposits,
            Error::MaxDepositsExceeded
        );

        // TODO: process deposits.

        /*
         * Exits
         */

        ensure!(
            block.body.exits.len() as u64 <= self.spec.max_exits,
            Error::MaxExitsExceeded
        );

        for exit in &block.body.exits {
            let validator = state
                .validator_registry
                .get(exit.validator_index as usize)
                .ok_or(Error::BadExit)?;
            ensure!(
                validator.exit_slot > state.slot + self.spec.entry_exit_delay,
                Error::BadExit
            );
            ensure!(state.slot >= exit.slot, Error::BadExit);
            let exit_message = {
                let exit_struct = Exit {
                    slot: exit.slot,
                    validator_index: exit.validator_index,
                    signature: self.spec.empty_signature.clone(),
                };
                hash_tree_root(&exit_struct)
            };
            ensure!(
                bls_verify(
                    &validator.pubkey,
                    &exit_message,
                    &exit.signature,
                    get_domain(&state.fork_data, exit.slot, DOMAIN_EXIT)
                ),
                Error::BadProposerSlashing
            );
            initiate_validator_exit(&state, exit.validator_index);
        }

        /*
         * Custody
         */
        ensure!(
            block.body.custody_reseeds.is_empty(),
            Error::BadCustodyReseeds
        );
        ensure!(
            block.body.custody_challenges.is_empty(),
            Error::BadCustodyChallenges
        );
        ensure!(
            block.body.custody_responses.is_empty(),
            Error::BadCustodyResponses
        );

        Ok(state)
    }
}

fn initiate_validator_exit(_state: &BeaconState, _index: u32) {
    // TODO: stubbed out.
}

fn get_attestation_participants(
    _state: &BeaconState,
    _attestation_data: &AttestationData,
    _aggregation_bitfield: &BooleanBitfield,
) -> Vec<usize> {
    vec![0, 1]
}

fn get_block_root(
    state: &BeaconState,
    slot: u64,
    latest_block_roots_length: u64,
) -> Option<&Hash256> {
    // TODO: test
    if state.slot <= slot + latest_block_roots_length && slot <= state.slot {
        state
            .latest_block_roots
            .get((slot % latest_block_roots_length) as usize)
    } else {
        None
    }
}

fn penalize_validator(_state: &BeaconState, _proposer_index: usize) {
    // TODO: stubbed out.
}

fn get_domain(_fork: &Fork, _slot: u64, _domain_type: u64) -> u64 {
    // TODO: stubbed out.
    0
}

fn bls_verify(pubkey: &PublicKey, message: &[u8], signature: &Signature, _domain: u64) -> bool {
    // TODO: add domain
    signature.verify(message, pubkey)
}

fn bls_verify_aggregate(
    pubkey: &AggregatePublicKey,
    message: &[u8],
    signature: &AggregateSignature,
    _domain: u64,
) -> bool {
    // TODO: add domain
    signature.verify(message, pubkey)
}

fn hash_tree_root<T>(_input: &T) -> Hash256 {
    // TODO: stubbed out.
    Hash256::zero()
}

impl From<DBError> for Error {
    fn from(e: DBError) -> Error {
        Error::DBError(e.message)
    }
}

impl From<TestingSlotClockError> for Error {
    fn from(_: TestingSlotClockError) -> Error {
        unreachable!(); // Testing clock never throws an error.
    }
}

impl From<SystemTimeSlotClockError> for Error {
    fn from(e: SystemTimeSlotClockError) -> Error {
        Error::SlotClockError(e)
    }
}

impl From<SlotProcessingError> for Error {
    fn from(e: SlotProcessingError) -> Error {
        match e {
            SlotProcessingError::UnableToDetermineProducer => Error::NoBlockProducer,
        }
    }
}
<commit_msg>Remove old hash_tree_root stub from state trans<commit_after>use super::{BeaconChain, ClientDB, DBError, SlotClock};
use bls::{AggregatePublicKey, AggregateSignature, PublicKey, Signature};
use boolean_bitfield::BooleanBitfield;
use hashing::hash;
use slot_clock::{SystemTimeSlotClockError, TestingSlotClockError};
use ssz::{ssz_encode, TreeHash};
use types::{
    beacon_state::SlotProcessingError, readers::BeaconBlockReader, AttestationData,
    AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Exit, Fork, Hash256,
    PendingAttestation,
};

// TODO: define elsehwere.
const DOMAIN_ATTESTATION: u64 = 1;
const DOMAIN_PROPOSAL: u64 = 2;
const DOMAIN_EXIT: u64 = 3;
const DOMAIN_RANDAO: u64 = 4;

macro_rules! ensure {
    ($condition: expr, $result: expr) => {
        if !$condition {
            return Err($result);
        }
    };
}

#[derive(Debug, PartialEq)]
pub enum Error {
    DBError(String),
    StateAlreadyTransitioned,
    NotImplemented,
    PresentSlotIsNone,
    UnableToDecodeBlock,
    MissingParentState(Hash256),
    InvalidParentState(Hash256),
    MissingBeaconBlock(Hash256),
    InvalidBeaconBlock(Hash256),
    MissingParentBlock(Hash256),
    NoBlockProducer,
    StateSlotMismatch,
    BadBlockSignature,
    BadRandaoSignature,
    MaxProposerSlashingsExceeded,
    BadProposerSlashing,
    MaxAttestationsExceeded,
    BadAttestation,
    NoBlockRoot,
    MaxDepositsExceeded,
    MaxExitsExceeded,
    BadExit,
    BadCustodyReseeds,
    BadCustodyChallenges,
    BadCustodyResponses,
    SlotClockError(SystemTimeSlotClockError),
}

impl<T, U> BeaconChain<T, U>
where
    T: ClientDB,
    U: SlotClock,
{
    pub fn state_transition(
        &self,
        state: BeaconState,
        block: &BeaconBlock,
    ) -> Result<BeaconState, Error> {
        self.internal_state_transition(state, block, true)
    }

    pub fn state_transition_without_verifying_block_signature(
        &self,
        state: BeaconState,
        block: &BeaconBlock,
    ) -> Result<BeaconState, Error> {
        self.internal_state_transition(state, block, false)
    }

    fn internal_state_transition(
        &self,
        mut state: BeaconState,
        block: &BeaconBlock,
        verify_block_signature: bool,
    ) -> Result<BeaconState, Error> {
        ensure!(state.slot < block.slot, Error::StateAlreadyTransitioned);

        for _ in state.slot..block.slot {
            state.per_slot_processing(block.parent_root.clone(), &self.spec)?;
        }

        /*
         * Slot
         */

        ensure!(block.slot() == state.slot, Error::StateSlotMismatch);

        /*
         * Proposer Signature
         */

        let block_proposer_index = state
            .get_beacon_proposer_index(block.slot, &self.spec)
            .ok_or(Error::NoBlockProducer)?;
        let block_proposer = &state.validator_registry[block_proposer_index];

        if verify_block_signature {
            ensure!(
                bls_verify(
                    &block_proposer.pubkey,
                    &block.proposal_root(&self.spec)[..],
                    &block.signature,
                    get_domain(&state.fork_data, state.slot, DOMAIN_PROPOSAL)
                ),
                Error::BadBlockSignature
            );
        }

        /*
         * RANDAO
         */

        ensure!(
            bls_verify(
                &block_proposer.pubkey,
                &ssz_encode(&block_proposer.proposer_slots),
                &block.randao_reveal,
                get_domain(&state.fork_data, state.slot, DOMAIN_RANDAO)
            ),
            Error::BadRandaoSignature
        );

        // TODO: check this is correct.
        let new_mix = {
            let mut mix = state.latest_randao_mixes
                [(state.slot % self.spec.latest_randao_mixes_length) as usize]
                .to_vec();
            mix.append(&mut ssz_encode(&block.randao_reveal));
            Hash256::from(&hash(&mix)[..])
        };

        state.latest_randao_mixes[(state.slot % self.spec.latest_randao_mixes_length) as usize] =
            new_mix;

        /*
         * Eth1 data
         */

        // TODO: Eth1 data stuff.

        /*
         * OPERATIONS
         */

        /*
         * Proposer slashings
         */

        ensure!(
            block.body.proposer_slashings.len() as u64 <= self.spec.max_proposer_slashings,
            Error::MaxProposerSlashingsExceeded
        );
        for proposer_slashing in &block.body.proposer_slashings {
            let proposer = state
                .validator_registry
                .get(proposer_slashing.proposer_index as usize)
                .ok_or(Error::BadProposerSlashing)?;
            ensure!(
                proposer_slashing.proposal_data_1.slot == proposer_slashing.proposal_data_2.slot,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer_slashing.proposal_data_1.shard == proposer_slashing.proposal_data_2.shard,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer_slashing.proposal_data_1.block_root
                    != proposer_slashing.proposal_data_2.block_root,
                Error::BadProposerSlashing
            );
            ensure!(
                proposer.penalized_slot > state.slot,
                Error::BadProposerSlashing
            );
            ensure!(
                bls_verify(
                    &proposer.pubkey,
                    &proposer_slashing.proposal_data_1.hash_tree_root(),
                    &proposer_slashing.proposal_signature_1,
                    get_domain(
                        &state.fork_data,
                        proposer_slashing.proposal_data_1.slot,
                        DOMAIN_PROPOSAL
                    )
                ),
                Error::BadProposerSlashing
            );
            ensure!(
                bls_verify(
                    &proposer.pubkey,
                    &proposer_slashing.proposal_data_2.hash_tree_root(),
                    &proposer_slashing.proposal_signature_2,
                    get_domain(
                        &state.fork_data,
                        proposer_slashing.proposal_data_2.slot,
                        DOMAIN_PROPOSAL
                    )
                ),
                Error::BadProposerSlashing
            );
            penalize_validator(&state, proposer_slashing.proposer_index as usize);
        }

        /*
         * Attestations
         */
        ensure!(
            block.body.attestations.len() as u64 <= self.spec.max_attestations,
            Error::MaxAttestationsExceeded
        );

        for attestation in &block.body.attestations {
            ensure!(
                attestation.data.slot + self.spec.min_attestation_inclusion_delay <= state.slot,
                Error::BadAttestation
            );
            ensure!(
                attestation.data.slot + self.spec.epoch_length >= state.slot,
                Error::BadAttestation
            );
            if state.justified_slot >= state.slot - (state.slot % self.spec.epoch_length) {
                ensure!(
                    attestation.data.justified_slot == state.justified_slot,
                    Error::BadAttestation
                );
            } else {
                ensure!(
                    attestation.data.justified_slot == state.previous_justified_slot,
                    Error::BadAttestation
                );
            }
            ensure!(
                attestation.data.justified_block_root
                    == *get_block_root(
                        &state,
                        attestation.data.justified_slot,
                        self.spec.latest_block_roots_length
                    )
                    .ok_or(Error::NoBlockRoot)?,
                Error::BadAttestation
            );
            ensure!(
                (attestation.data.latest_crosslink_root
                    == state.latest_crosslinks[attestation.data.shard as usize].shard_block_root)
                    || (attestation.data.shard_block_root
                        == state.latest_crosslinks[attestation.data.shard as usize]
                            .shard_block_root),
                Error::BadAttestation
            );
            let participants = get_attestation_participants(
                &state,
                &attestation.data,
                &attestation.aggregation_bitfield,
            );
            let mut group_public_key = AggregatePublicKey::new();
            for participant in participants {
                group_public_key.add(
                    state.validator_registry[participant as usize]
                        .pubkey
                        .as_raw(),
                )
            }
            let attestation_message = {
                let attestation_data_and_custody_bit = AttestationDataAndCustodyBit {
                    data: attestation.data.clone(),
                    custody_bit: false,
                };
                &attestation_data_and_custody_bit.hash_tree_root()
            };
            // Signature verification.
            ensure!(
                bls_verify_aggregate(
                    &group_public_key,
                    &attestation_message[..],
                    &attestation.aggregate_signature,
                    get_domain(&state.fork_data, attestation.data.slot, DOMAIN_ATTESTATION)
                ),
                Error::BadProposerSlashing
            );
            ensure!(
                attestation.data.shard_block_root == self.spec.zero_hash,
                Error::BadAttestation
            );
            let pending_attestation = PendingAttestation {
                data: attestation.data.clone(),
                aggregation_bitfield: attestation.aggregation_bitfield.clone(),
                custody_bitfield: attestation.custody_bitfield.clone(),
                slot_included: state.slot,
            };
            state.latest_attestations.push(pending_attestation);
        }

        /*
         * Deposits
         */
        ensure!(
            block.body.deposits.len() as u64 <= self.spec.max_deposits,
            Error::MaxDepositsExceeded
        );

        // TODO: process deposits.

        /*
         * Exits
         */

        ensure!(
            block.body.exits.len() as u64 <= self.spec.max_exits,
            Error::MaxExitsExceeded
        );

        for exit in &block.body.exits {
            let validator = state
                .validator_registry
                .get(exit.validator_index as usize)
                .ok_or(Error::BadExit)?;
            ensure!(
                validator.exit_slot > state.slot + self.spec.entry_exit_delay,
                Error::BadExit
            );
            ensure!(state.slot >= exit.slot, Error::BadExit);
            let exit_message = {
                let exit_struct = Exit {
                    slot: exit.slot,
                    validator_index: exit.validator_index,
                    signature: self.spec.empty_signature.clone(),
                };
                exit_struct.hash_tree_root()
            };
            ensure!(
                bls_verify(
                    &validator.pubkey,
                    &exit_message,
                    &exit.signature,
                    get_domain(&state.fork_data, exit.slot, DOMAIN_EXIT)
                ),
                Error::BadProposerSlashing
            );
            initiate_validator_exit(&state, exit.validator_index);
        }

        /*
         * Custody
         */
        ensure!(
            block.body.custody_reseeds.is_empty(),
            Error::BadCustodyReseeds
        );
        ensure!(
            block.body.custody_challenges.is_empty(),
            Error::BadCustodyChallenges
        );
        ensure!(
            block.body.custody_responses.is_empty(),
            Error::BadCustodyResponses
        );

        Ok(state)
    }
}

fn initiate_validator_exit(_state: &BeaconState, _index: u32) {
    // TODO: stubbed out.
}

fn get_attestation_participants(
    _state: &BeaconState,
    _attestation_data: &AttestationData,
    _aggregation_bitfield: &BooleanBitfield,
) -> Vec<usize> {
    vec![0, 1]
}

fn get_block_root(
    state: &BeaconState,
    slot: u64,
    latest_block_roots_length: u64,
) -> Option<&Hash256> {
    // TODO: test
    if state.slot <= slot + latest_block_roots_length && slot <= state.slot {
        state
            .latest_block_roots
            .get((slot % latest_block_roots_length) as usize)
    } else {
        None
    }
}

fn penalize_validator(_state: &BeaconState, _proposer_index: usize) {
    // TODO: stubbed out.
}

fn get_domain(_fork: &Fork, _slot: u64, _domain_type: u64) -> u64 {
    // TODO: stubbed out.
    0
}

fn bls_verify(pubkey: &PublicKey, message: &[u8], signature: &Signature, _domain: u64) -> bool {
    // TODO: add domain
    signature.verify(message, pubkey)
}

fn bls_verify_aggregate(
    pubkey: &AggregatePublicKey,
    message: &[u8],
    signature: &AggregateSignature,
    _domain: u64,
) -> bool {
    // TODO: add domain
    signature.verify(message, pubkey)
}

impl From<DBError> for Error {
    fn from(e: DBError) -> Error {
        Error::DBError(e.message)
    }
}

impl From<TestingSlotClockError> for Error {
    fn from(_: TestingSlotClockError) -> Error {
        unreachable!(); // Testing clock never throws an error.
    }
}

impl From<SystemTimeSlotClockError> for Error {
    fn from(e: SystemTimeSlotClockError) -> Error {
        Error::SlotClockError(e)
    }
}

impl From<SlotProcessingError> for Error {
    fn from(e: SlotProcessingError) -> Error {
        match e {
            SlotProcessingError::UnableToDetermineProducer => Error::NoBlockProducer,
        }
    }
}
<|endoftext|>


<commit_before>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<commit_msg>futures_cpupool: Implement `Debug` for `Builder`<commit_after>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<|endoftext|>