<commit_before>extern crate core;

use std::mem;
use std::cmp::min;
use std::iter::FromIterator;
use self::core::str::next_code_point;
use super::token::Token;
use super::token::Category;

pub struct StateFunction(pub fn(&mut Tokenizer) -> Option<StateFunction>);

/// The Tokenizer type is used to produce and store
/// tokens for the various language and format lexers.
pub struct Tokenizer {
    pub data: String,
    binary_data: Vec<u8>,
    head: usize,
    tail: usize,
    char_count: usize,
    pub token_start: usize,
    pub token_position: usize,
    tokens: Vec<Token>,
    pub states: Vec<StateFunction>,
}

/// Initializes a new tokenizer with the given data.
///
/// # Examples
///
/// ```
/// let lexer = luthor::tokenizer::new("luthor");
/// ```
pub fn new(data: &str) -> Tokenizer {
    Tokenizer{
      data: data.to_string(),
      binary_data: data.to_string().into_bytes(),
      head: 0,
      tail: 0,
      char_count: data.chars().count(),
      token_start: 0,
      token_position: 0,
      tokens: vec![],
      states: vec![]
    }
}

impl Tokenizer {
    /// Returns a copy of the tokens processed to date.
    ///
    /// # Examples
    ///
    /// ```
    /// let lexer = luthor::tokenizer::new("luthor");
    /// lexer.tokens();
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        self.tokens.clone()
    }

    /// Moves to the next character in the data.
    /// Does nothing if there is no more data to process.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut lexer = luthor::tokenizer::new("luthor");
    /// assert_eq!(lexer.current_char().unwrap(), 'l');
    /// lexer.advance();
    /// assert_eq!(lexer.current_char().unwrap(), 'u');
    /// ```
    pub fn advance(&mut self) {
        if self.has_more_data() {
            // Get the current character so that we can calculate
            // its byte length and advance the head appropriately.
            match self.current_char() {
                Some(c) => {
                    self.token_position += 1;
                    self.head += c.len_utf8();
                },
                None => (),
            }
        }
    }

    /// Determines whether or not there is more unprocessed data.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut lexer = luthor::tokenizer::new("l");
    /// assert_eq!(lexer.has_more_data(), true);
    /// lexer.advance();
    /// assert_eq!(lexer.has_more_data(), false);
    /// ```
    pub fn has_more_data(&self) -> bool {
        self.token_position < self.char_count
    }

    /// Returns the character at the current position,
    /// unless all of the data has been processed.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut lexer = luthor::tokenizer::new("l");
    /// assert_eq!(lexer.current_char().unwrap(), 'l');
    /// lexer.advance();
    /// assert_eq!(lexer.current_char(), None);
    /// ```
    pub fn current_char(&self) -> Option<char> {
        if self.has_more_data() {
            // Create an iterator for the remaining data.
            let mut remaining_data = self.binary_data[self.head..].iter();

            // Pull the first UTF-8 byte sequence and convert it to a char.
            next_code_point(&mut remaining_data).map(|code_point| {
                unsafe {
                    mem::transmute(code_point)
                }
            })
        } else {
            None
        }
    }

    /// Creates and stores a token with the given category containing any
    /// data processed using `advance` since the last call to this method.
    ///
    /// # Examples
    ///
    /// ```
    /// use luthor::token::Category;
    /// let mut lexer = luthor::tokenizer::new("luthor");
    /// lexer.advance();
    /// lexer.advance();
    /// lexer.tokenize(Category::Text);
    /// assert_eq!(lexer.tokens()[0].lexeme, "lu");
    /// ```
    pub fn tokenize(&mut self, category: Category) {
        if self.head > self.tail {
            // Build the lexeme by slicing the currently
            // selected range out of the buffer.
            let lexeme = unsafe {
                String::from_utf8_unchecked(
                    self.binary_data[self.tail..self.head].to_vec()
                )
            };
            let token = Token{
                lexeme: lexeme,
                category: category,
            };
            self.tokens.push(token);
            self.token_start = self.token_position;
            self.tail = self.head;
        }
    }

    /// Creates and stores a token with the given category and the
    /// next `amount` characters of the data. Before doing this, it
    /// tokenizes any previously processed characters with the generic
    /// Category::Text category.
    ///
    /// # Examples
    ///
    /// ```
    /// use luthor::token::Category;
    /// use luthor::token::Token;
    ///
    /// let mut lexer = luthor::tokenizer::new("luthor");
    /// lexer.advance();
    /// lexer.tokenize_next(5, Category::Keyword);
    /// assert_eq!(lexer.tokens()[0], Token{ lexeme: "l".to_string(), category: Category::Text});
    /// assert_eq!(lexer.tokens()[1], Token{ lexeme: "uthor".to_string(), category: Category::Keyword});
    /// ```
    pub fn tokenize_next(&mut self, amount: usize, category: Category) {
        // If there's any data that has yet
        // to be tokenized, take care of that.
        self.tokenize(Category::Text);

        // Mark the next series of characters.
        for _ in 0..amount { self.advance(); }

        // Tokenize the marked characters.
        self.tokenize(category);
    }
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::super::token::Token;
    use super::super::token::Category;

    #[test]
    fn advance_does_nothing_when_there_is_no_more_data() {
        let data = "élégant";
        let mut tokenizer = new(data);

        // Try to go beyond the last character.
        for _ in 0..10 {
            tokenizer.advance();
        }

        assert!(!tokenizer.has_more_data())
    }

    #[test]
    fn has_more_data_works() {
        let data = "él";
        let mut tokenizer = new(data);

        tokenizer.advance();
        assert!(tokenizer.has_more_data());

        tokenizer.advance();
        assert!(!tokenizer.has_more_data());
    }

    #[test]
    fn current_char_returns_the_char_at_token_position() {
        let data = "él";
        let tokenizer = new(data);

        assert_eq!(tokenizer.current_char().unwrap(), 'é');
    }

    #[test]
    fn current_char_returns_none_if_at_the_end() {
        let data = "él";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();

        assert_eq!(tokenizer.current_char(), None);
    }

    #[test]
    fn tokenize_creates_the_correct_token() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize(Category::Text);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "él".to_string(), category: Category::Text};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_does_nothing_if_range_is_empty() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.tokenize(Category::Text);

        assert_eq!(tokenizer.tokens.len(), 0);
    }

    #[test]
    fn tokenize_next_tokenizes_previous_data_as_text() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(1, Category::Keyword);

        let token = tokenizer.tokens.remove(0);
        let expected_token = Token{ lexeme: "él".to_string(), category: Category::Text};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_next_tokenizes_next_x_chars() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(5, Category::Keyword);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "égant".to_string(), category: Category::Keyword};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_next_takes_at_most_what_is_left() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(15, Category::Keyword);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "égant".to_string(), category: Category::Keyword};
        assert_eq!(token, expected_token);
    }
}
<commit_msg>Correct documentation test variable names.<commit_after>extern crate core;

use std::mem;
use std::cmp::min;
use std::iter::FromIterator;
use self::core::str::next_code_point;
use super::token::Token;
use super::token::Category;

pub struct StateFunction(pub fn(&mut Tokenizer) -> Option<StateFunction>);

/// The Tokenizer type is used to produce and store
/// tokens for the various language and format lexers.
pub struct Tokenizer {
    pub data: String,
    binary_data: Vec<u8>,
    head: usize,
    tail: usize,
    char_count: usize,
    pub token_start: usize,
    pub token_position: usize,
    tokens: Vec<Token>,
    pub states: Vec<StateFunction>,
}

/// Initializes a new tokenizer with the given data.
///
/// # Examples
///
/// ```
/// let tokenizer = luthor::tokenizer::new("luthor");
/// ```
pub fn new(data: &str) -> Tokenizer {
    Tokenizer{
      data: data.to_string(),
      binary_data: data.to_string().into_bytes(),
      head: 0,
      tail: 0,
      char_count: data.chars().count(),
      token_start: 0,
      token_position: 0,
      tokens: vec![],
      states: vec![]
    }
}

impl Tokenizer {
    /// Returns a copy of the tokens processed to date.
    ///
    /// # Examples
    ///
    /// ```
    /// let tokenizer = luthor::tokenizer::new("luthor");
    /// tokenizer.tokens();
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        self.tokens.clone()
    }

    /// Moves to the next character in the data.
    /// Does nothing if there is no more data to process.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut tokenizer = luthor::tokenizer::new("luthor");
    /// assert_eq!(tokenizer.current_char().unwrap(), 'l');
    /// tokenizer.advance();
    /// assert_eq!(tokenizer.current_char().unwrap(), 'u');
    /// ```
    pub fn advance(&mut self) {
        if self.has_more_data() {
            // Get the current character so that we can calculate
            // its byte length and advance the head appropriately.
            match self.current_char() {
                Some(c) => {
                    self.token_position += 1;
                    self.head += c.len_utf8();
                },
                None => (),
            }
        }
    }

    /// Determines whether or not there is more unprocessed data.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut tokenizer = luthor::tokenizer::new("l");
    /// assert_eq!(tokenizer.has_more_data(), true);
    /// tokenizer.advance();
    /// assert_eq!(tokenizer.has_more_data(), false);
    /// ```
    pub fn has_more_data(&self) -> bool {
        self.token_position < self.char_count
    }

    /// Returns the character at the current position,
    /// unless all of the data has been processed.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut tokenizer = luthor::tokenizer::new("l");
    /// assert_eq!(tokenizer.current_char().unwrap(), 'l');
    /// tokenizer.advance();
    /// assert_eq!(tokenizer.current_char(), None);
    /// ```
    pub fn current_char(&self) -> Option<char> {
        if self.has_more_data() {
            // Create an iterator for the remaining data.
            let mut remaining_data = self.binary_data[self.head..].iter();

            // Pull the first UTF-8 byte sequence and convert it to a char.
            next_code_point(&mut remaining_data).map(|code_point| {
                unsafe {
                    mem::transmute(code_point)
                }
            })
        } else {
            None
        }
    }

    /// Creates and stores a token with the given category containing any
    /// data processed using `advance` since the last call to this method.
    ///
    /// # Examples
    ///
    /// ```
    /// use luthor::token::Category;
    /// let mut tokenizer = luthor::tokenizer::new("luthor");
    /// tokenizer.advance();
    /// tokenizer.advance();
    /// tokenizer.tokenize(Category::Text);
    /// assert_eq!(tokenizer.tokens()[0].lexeme, "lu");
    /// ```
    pub fn tokenize(&mut self, category: Category) {
        if self.head > self.tail {
            // Build the lexeme by slicing the currently
            // selected range out of the buffer.
            let lexeme = unsafe {
                String::from_utf8_unchecked(
                    self.binary_data[self.tail..self.head].to_vec()
                )
            };
            let token = Token{
                lexeme: lexeme,
                category: category,
            };
            self.tokens.push(token);
            self.token_start = self.token_position;
            self.tail = self.head;
        }
    }

    /// Creates and stores a token with the given category and the
    /// next `amount` characters of the data. Before doing this, it
    /// tokenizes any previously processed characters with the generic
    /// Category::Text category.
    ///
    /// # Examples
    ///
    /// ```
    /// use luthor::token::Category;
    /// use luthor::token::Token;
    ///
    /// let mut tokenizer = luthor::tokenizer::new("luthor");
    /// tokenizer.advance();
    /// tokenizer.tokenize_next(5, Category::Keyword);
    /// assert_eq!(tokenizer.tokens()[0], Token{ lexeme: "l".to_string(), category: Category::Text});
    /// assert_eq!(tokenizer.tokens()[1], Token{ lexeme: "uthor".to_string(), category: Category::Keyword});
    /// ```
    pub fn tokenize_next(&mut self, amount: usize, category: Category) {
        // If there's any data that has yet
        // to be tokenized, take care of that.
        self.tokenize(Category::Text);

        // Mark the next series of characters.
        for _ in 0..amount { self.advance(); }

        // Tokenize the marked characters.
        self.tokenize(category);
    }
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::super::token::Token;
    use super::super::token::Category;

    #[test]
    fn advance_does_nothing_when_there_is_no_more_data() {
        let data = "élégant";
        let mut tokenizer = new(data);

        // Try to go beyond the last character.
        for _ in 0..10 {
            tokenizer.advance();
        }

        assert!(!tokenizer.has_more_data())
    }

    #[test]
    fn has_more_data_works() {
        let data = "él";
        let mut tokenizer = new(data);

        tokenizer.advance();
        assert!(tokenizer.has_more_data());

        tokenizer.advance();
        assert!(!tokenizer.has_more_data());
    }

    #[test]
    fn current_char_returns_the_char_at_token_position() {
        let data = "él";
        let tokenizer = new(data);

        assert_eq!(tokenizer.current_char().unwrap(), 'é');
    }

    #[test]
    fn current_char_returns_none_if_at_the_end() {
        let data = "él";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();

        assert_eq!(tokenizer.current_char(), None);
    }

    #[test]
    fn tokenize_creates_the_correct_token() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize(Category::Text);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "él".to_string(), category: Category::Text};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_does_nothing_if_range_is_empty() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.tokenize(Category::Text);

        assert_eq!(tokenizer.tokens.len(), 0);
    }

    #[test]
    fn tokenize_next_tokenizes_previous_data_as_text() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(1, Category::Keyword);

        let token = tokenizer.tokens.remove(0);
        let expected_token = Token{ lexeme: "él".to_string(), category: Category::Text};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_next_tokenizes_next_x_chars() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(5, Category::Keyword);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "égant".to_string(), category: Category::Keyword};
        assert_eq!(token, expected_token);
    }

    #[test]
    fn tokenize_next_takes_at_most_what_is_left() {
        let data = "élégant";
        let mut tokenizer = new(data);
        tokenizer.advance();
        tokenizer.advance();
        tokenizer.tokenize_next(15, Category::Keyword);

        let token = tokenizer.tokens.pop().unwrap();
        let expected_token = Token{ lexeme: "égant".to_string(), category: Category::Keyword};
        assert_eq!(token, expected_token);
    }
}
<|endoftext|>


<commit_before>use futures::future;
use futures::Future;
use hyper::Uri;
use hyper::client::{Client, Request};
use hyper::Method::Post;
use mktemp::Temp;
use std::io::prelude::*;
use std::path::Path;
use std::process::Command;
use tokio_core::reactor::{Handle, Remote};
use tokio_process::CommandExt;
use tera::Tera;

use http_client;
use papers::DocumentSpec;
use error::Error;

pub struct Workspace {
    dir: Temp,
    document_spec: DocumentSpec,
    handle: Handle,
    template_path: ::std::path::PathBuf,
}

/// We ignore file names that end with a slash for now, and always determine the file name from the Uri
/// TODO: investigate Content-Disposition: attachment
fn extract_file_name_from_uri(uri: &Uri) -> Option<String> {
    match uri.path().split('/').last() {
        Some(name) => {
            if !name.is_empty() {
                Some(name.to_string())
            } else {
                None
            }
        },
        None => None,
    }
}

/// Since mktemp::Temp implements Drop by deleting the directory, we don't need to worry about
/// leaving files or directories behind.
impl Workspace {

    pub fn new(remote: Remote, document_spec: DocumentSpec) -> Result<Workspace, Error> {
        let dir = Temp::new_dir()?;
        let mut template_path = dir.to_path_buf();
        template_path.push(Path::new("out.tex"));
        Ok(Workspace {
            document_spec,
            handle: remote.handle().unwrap(),
            template_path,
            dir,
        })
    }

    pub fn preview(self) -> Box<Future<Item=String, Error=Error>> {
        let Workspace {
            handle,
            document_spec,
            template_path: _,
            dir: _,
        } = self;

        let DocumentSpec {
            assets_urls: _,
            callback_url: _,
            template_url,
            variables,
        } = document_spec;

        // Download the template and populate it
        let work = http_client::download_file(&handle.clone(), template_url.0)
            .and_then(|bytes| {
                future::result(::std::string::String::from_utf8(bytes))
                    .map_err(Error::from)
            }).and_then(move |template_string| {
                Tera::one_off(&template_string, &variables, false)
                    .map_err(Error::from)
            });

        Box::new(work)
    }

    pub fn execute(self) -> Box<Future<Item=(), Error=Error>> {
        let Workspace {
            handle,
            document_spec,
            template_path,
            dir,
        } = self;

        let DocumentSpec {
            assets_urls,
            callback_url,
            template_url,
            variables,
        } = document_spec;

        let out_tex_path = dir.to_path_buf();
        let out_pdf_path = dir.to_path_buf();

        // First download the template and populate it
        let work = http_client::download_file(&handle.clone(), template_url.0)
                .and_then(|bytes| {
                    future::result(::std::string::String::from_utf8(bytes))
                        .map_err(Error::from)
                        .map(|template_string| (handle, template_string))
                }).and_then(move |(handle, template_string)| {
                    Tera::one_off(&template_string, &variables, false)
                        .map_err(Error::from)
                        .map(|latex_string| (handle, latex_string))
                }).and_then(|(handle, latex_string)| {
                    let mut file = ::std::fs::File::open(template_path.clone()).unwrap();
                    file.write_all(latex_string.as_bytes()).unwrap();
                    future::ok((handle, template_path))
                })

        // Then download the assets and save them in the temporary directory
                .and_then(move |(handle, template_path)| {
                    let named = assets_urls.into_iter().filter_map(|url| {
                        let inner = url.0;
                        extract_file_name_from_uri(&inner).map(|file_name| (file_name, inner))
                    });

                    let inner_handle = handle.clone();

                    let download_named = move |(name, url)| {
                        let mut path = out_tex_path.clone();
                        path.push(name);

                        http_client::download_file(&inner_handle, url)
                            .map(move |bytes| ::std::fs::File::open(&path).unwrap().write_all(&bytes))
                            .map_err(Error::from)
                    };
                    future::join_all(named.map(download_named))
                        .map(|result| (handle, template_path, result))
                })

        // Then run latex
                .and_then(move |(handle, template_path, _)| {
                    let inner_handle = handle.clone();
                    Command::new("pdflatex")
                        .arg(template_path.clone().to_str().unwrap())
                        .status_async(&inner_handle)
                        .map(|exit_status| (exit_status, handle))
                        .map_err(Error::from)
                }).and_then(|(exit_status, handle)| {
                    if exit_status.success() {
                        future::ok(handle)
                    } else {
                        future::err(Error::LatexFailed)
                    }
                })

        // Then construct the path to the generated PDF
                .and_then(move |handle| {
                    let mut path = out_pdf_path;
                    path.push(Path::new("out.pdf"));
                    future::ok((handle, path))
                })

        // Then get the bytes from the generated PDF path
                .and_then(|(handle, pdf_path)| {
                    let pdf_file = ::std::fs::File::open(pdf_path).unwrap();
                    let pdf_bytes: Vec<u8> = pdf_file.bytes().collect::<Result<Vec<u8>, _>>().unwrap();
                    future::ok((handle, pdf_bytes))
                })

        // Finally, post the PDF to the callback URL
                .and_then(move |(handle, pdf_bytes)| {
                    let mut request = Request::new(Post, callback_url.0);
                    request.set_body(pdf_bytes);

                    let client = Client::new(&handle);
                    client.request(request);
                    future::ok(())
                });

        Box::new(work)
    }
}

#[cfg(test)]
mod tests {
    use hyper::Uri;
    use super::extract_file_name_from_uri;

    #[test]
    fn test_extract_file_name_from_uri_works() {
        let assert_extracted = |input: &'static str, expected_output: Option<&'static str>| {
            let uri = input.parse::<Uri>().unwrap();
            assert_eq!(extract_file_name_from_uri(&uri), expected_output.map(|o| o.to_string()));
        };

        assert_extracted("/logo.png", Some("logo.png"));
        assert_extracted("/assets/", None);
        assert_extracted("/assets/icon", Some("icon"));
        assert_extracted("/", None);
        assert_extracted("http://www.store2be.com", None);
    }
}
<commit_msg>:hammer: Return request from last then<commit_after>use futures::future;
use futures::Future;
use hyper::Uri;
use hyper::client::{Client, Request};
use hyper::Method::Post;
use mktemp::Temp;
use std::io::prelude::*;
use std::path::Path;
use std::process::Command;
use tokio_core::reactor::{Handle, Remote};
use tokio_process::CommandExt;
use tera::Tera;

use http_client;
use papers::DocumentSpec;
use error::Error;

pub struct Workspace {
    dir: Temp,
    document_spec: DocumentSpec,
    handle: Handle,
    template_path: ::std::path::PathBuf,
}

/// We ignore file names that end with a slash for now, and always determine the file name from the Uri
/// TODO: investigate Content-Disposition: attachment
fn extract_file_name_from_uri(uri: &Uri) -> Option<String> {
    match uri.path().split('/').last() {
        Some(name) => {
            if !name.is_empty() {
                Some(name.to_string())
            } else {
                None
            }
        },
        None => None,
    }
}

/// Since mktemp::Temp implements Drop by deleting the directory, we don't need to worry about
/// leaving files or directories behind.
impl Workspace {

    pub fn new(remote: Remote, document_spec: DocumentSpec) -> Result<Workspace, Error> {
        let dir = Temp::new_dir()?;
        let mut template_path = dir.to_path_buf();
        template_path.push(Path::new("out.tex"));
        Ok(Workspace {
            document_spec,
            handle: remote.handle().unwrap(),
            template_path,
            dir,
        })
    }

    pub fn preview(self) -> Box<Future<Item=String, Error=Error>> {
        let Workspace {
            handle,
            document_spec,
            template_path: _,
            dir: _,
        } = self;

        let DocumentSpec {
            assets_urls: _,
            callback_url: _,
            template_url,
            variables,
        } = document_spec;

        // Download the template and populate it
        let work = http_client::download_file(&handle.clone(), template_url.0)
            .and_then(|bytes| {
                future::result(::std::string::String::from_utf8(bytes))
                    .map_err(Error::from)
            }).and_then(move |template_string| {
                Tera::one_off(&template_string, &variables, false)
                    .map_err(Error::from)
            });

        Box::new(work)
    }

    pub fn execute(self) -> Box<Future<Item=(), Error=Error>> {
        let Workspace {
            handle,
            document_spec,
            template_path,
            dir,
        } = self;

        let DocumentSpec {
            assets_urls,
            callback_url,
            template_url,
            variables,
        } = document_spec;

        let out_tex_path = dir.to_path_buf();
        let out_pdf_path = dir.to_path_buf();

        // First download the template and populate it
        let work = http_client::download_file(&handle.clone(), template_url.0)
                .and_then(|bytes| {
                    future::result(::std::string::String::from_utf8(bytes))
                        .map_err(Error::from)
                        .map(|template_string| (handle, template_string))
                }).and_then(move |(handle, template_string)| {
                    Tera::one_off(&template_string, &variables, false)
                        .map_err(Error::from)
                        .map(|latex_string| (handle, latex_string))
                }).and_then(|(handle, latex_string)| {
                    let mut file = ::std::fs::File::open(template_path.clone()).unwrap();
                    file.write_all(latex_string.as_bytes()).unwrap();
                    future::ok((handle, template_path))
                })

        // Then download the assets and save them in the temporary directory
                .and_then(move |(handle, template_path)| {
                    let named = assets_urls.into_iter().filter_map(|url| {
                        let inner = url.0;
                        extract_file_name_from_uri(&inner).map(|file_name| (file_name, inner))
                    });

                    let inner_handle = handle.clone();

                    let download_named = move |(name, url)| {
                        let mut path = out_tex_path.clone();
                        path.push(name);

                        http_client::download_file(&inner_handle, url)
                            .map(move |bytes| ::std::fs::File::open(&path).unwrap().write_all(&bytes))
                            .map_err(Error::from)
                    };
                    future::join_all(named.map(download_named))
                        .map(|result| (handle, template_path, result))
                })

        // Then run latex
                .and_then(move |(handle, template_path, _)| {
                    let inner_handle = handle.clone();
                    Command::new("pdflatex")
                        .arg(template_path.clone().to_str().unwrap())
                        .status_async(&inner_handle)
                        .map(|exit_status| (exit_status, handle))
                        .map_err(Error::from)
                }).and_then(|(exit_status, handle)| {
                    if exit_status.success() {
                        future::ok(handle)
                    } else {
                        future::err(Error::LatexFailed)
                    }
                })

        // Then construct the path to the generated PDF
                .and_then(move |handle| {
                    let mut path = out_pdf_path;
                    path.push(Path::new("out.pdf"));
                    future::ok((handle, path))
                })

        // Then get the bytes from the generated PDF path
                .and_then(|(handle, pdf_path)| {
                    let pdf_file = ::std::fs::File::open(pdf_path).unwrap();
                    let pdf_bytes: Vec<u8> = pdf_file.bytes().collect::<Result<Vec<u8>, _>>().unwrap();
                    future::ok((handle, pdf_bytes))
                })

        // Finally, post the PDF to the callback URL
                .and_then(move |(handle, pdf_bytes)| {
                    let mut request = Request::new(Post, callback_url.0);
                    request.set_body(pdf_bytes);

                    Client::new(&handle).request(request)
                        .map(|_| ())
                        .map_err(Error::from)
                });

        Box::new(work)
    }
}

#[cfg(test)]
mod tests {
    use hyper::Uri;
    use super::extract_file_name_from_uri;

    #[test]
    fn test_extract_file_name_from_uri_works() {
        let assert_extracted = |input: &'static str, expected_output: Option<&'static str>| {
            let uri = input.parse::<Uri>().unwrap();
            assert_eq!(extract_file_name_from_uri(&uri), expected_output.map(|o| o.to_string()));
        };

        assert_extracted("/logo.png", Some("logo.png"));
        assert_extracted("/assets/", None);
        assert_extracted("/assets/icon", Some("icon"));
        assert_extracted("/", None);
        assert_extracted("http://www.store2be.com", None);
    }
}
<|endoftext|>


<commit_before>use buffer;
use buffer::Buffer;
use std::old_io::IoError;

pub struct Workspace {
    path: Path,
    buffers: Vec<Buffer>,
    current_buffer_index: Option<usize>,
}

impl Workspace {
    pub fn add_buffer(&mut self, buf: Buffer) {
        self.buffers.push(buf);
        self.current_buffer_index = Some(self.buffers.len()-1);
    }

    pub fn current_buffer(&mut self) -> Option<&mut Buffer> {
        match self.current_buffer_index {
            Some(index) => Some(&mut self.buffers[index]),
            None => None,
        }
    }

    pub fn close_current_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                self.buffers.remove(index);

                if self.buffers.is_empty() {
                    self.current_buffer_index = None;
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            }
            None => return,
        };
    }

    pub fn previous_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                if index > 0 {
                    self.current_buffer_index = Some(index-1);
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            },
            None => return,
        }
    }
}

pub fn new(path: Path) -> Workspace {
    Workspace{ path: path, buffers: Vec::new(), current_buffer_index: None }
}

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

    #[test]
    fn add_buffer_adds_and_selects_the_passed_buffer() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);

        assert_eq!(workspace.buffers.len(), 1);
        assert_eq!(workspace.current_buffer().unwrap().data(), "it works!\n");
    }

    #[test]
    fn current_buffer_returns_none_when_there_are_no_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn current_buffer_returns_one_when_there_are_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);
        assert!(workspace.current_buffer().is_some());
    }

    #[test]
    fn close_current_buffer_does_nothing_when_none_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn close_current_buffer_cleans_up_when_only_one_buffer_is_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.add_buffer(buffer::new());
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
        assert!(workspace.current_buffer_index.is_none());
    }

    #[test]
    fn close_current_buffer_when_two_are_open_selects_the_other() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);

        // Ensure that the second buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        workspace.close_current_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");
    }

    #[test]
    fn previous_buffer_does_nothing_when_no_buffers_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.previous_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn previous_buffer_when_three_are_open_selects_previous_wrapping_to_last() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        let mut third_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        third_buffer.insert("third buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);
        workspace.add_buffer(third_buffer);

        // Ensure that the third buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");

        // Ensure that the second buffer is returned.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        // Ensure that the first buffer is returned.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");

        // Ensure that it wraps back to the third buffer.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");
    }
}
<commit_msg>Implement next_buffer method for workspace type.<commit_after>use buffer;
use buffer::Buffer;
use std::old_io::IoError;

pub struct Workspace {
    path: Path,
    buffers: Vec<Buffer>,
    current_buffer_index: Option<usize>,
}

impl Workspace {
    pub fn add_buffer(&mut self, buf: Buffer) {
        self.buffers.push(buf);
        self.current_buffer_index = Some(self.buffers.len()-1);
    }

    pub fn current_buffer(&mut self) -> Option<&mut Buffer> {
        match self.current_buffer_index {
            Some(index) => Some(&mut self.buffers[index]),
            None => None,
        }
    }

    pub fn close_current_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                self.buffers.remove(index);

                if self.buffers.is_empty() {
                    self.current_buffer_index = None;
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            }
            None => return,
        };
    }

    pub fn previous_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                if index > 0 {
                    self.current_buffer_index = Some(index-1);
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            },
            None => return,
        }
    }

    pub fn next_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                if index == self.buffers.len()-1 {
                    self.current_buffer_index = Some(0);
                } else {
                    self.current_buffer_index = Some(index+1);
                }
            },
            None => return,
        }
    }
}

pub fn new(path: Path) -> Workspace {
    Workspace{ path: path, buffers: Vec::new(), current_buffer_index: None }
}

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

    #[test]
    fn add_buffer_adds_and_selects_the_passed_buffer() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);

        assert_eq!(workspace.buffers.len(), 1);
        assert_eq!(workspace.current_buffer().unwrap().data(), "it works!\n");
    }

    #[test]
    fn current_buffer_returns_none_when_there_are_no_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn current_buffer_returns_one_when_there_are_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);
        assert!(workspace.current_buffer().is_some());
    }

    #[test]
    fn close_current_buffer_does_nothing_when_none_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn close_current_buffer_cleans_up_when_only_one_buffer_is_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.add_buffer(buffer::new());
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
        assert!(workspace.current_buffer_index.is_none());
    }

    #[test]
    fn close_current_buffer_when_two_are_open_selects_the_other() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);

        // Ensure that the second buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        workspace.close_current_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");
    }

    #[test]
    fn previous_buffer_does_nothing_when_no_buffers_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.previous_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn previous_buffer_when_three_are_open_selects_previous_wrapping_to_last() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        let mut third_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        third_buffer.insert("third buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);
        workspace.add_buffer(third_buffer);

        // Ensure that the third buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");

        // Ensure that the second buffer is returned.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        // Ensure that the first buffer is returned.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");

        // Ensure that it wraps back to the third buffer.
        workspace.previous_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");
    }

    #[test]
    fn next_buffer_does_nothing_when_no_buffers_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.next_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn next_buffer_when_three_are_open_selects_next_wrapping_to_first() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        let mut third_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        third_buffer.insert("third buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);
        workspace.add_buffer(third_buffer);

        // Ensure that the third buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");

        // Ensure that it wraps back to the first buffer.
        workspace.next_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");

        // Ensure that the second buffer is returned.
        workspace.next_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        // Ensure that the third buffer is returned.
        workspace.next_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "third buffer");
    }
}
<|endoftext|>


<commit_before>use super::frontend_prelude::*;

use crate::models::ApiToken;
use crate::schema::api_tokens;
use crate::util::read_fill;
use crate::views::EncodableApiTokenWithToken;

use conduit::{Body, Response};
use serde_json as json;

/// Handles the `GET /me/tokens` route.
pub fn list(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let conn = req.db_conn()?;
    let user = authenticated_user.user();

    let tokens = ApiToken::belonging_to(&user)
        .filter(api_tokens::revoked.eq(false))
        .order(api_tokens::created_at.desc())
        .load(&*conn)?;
    #[derive(Serialize)]
    struct R {
        api_tokens: Vec<ApiToken>,
    }
    Ok(req.json(&R { api_tokens: tokens }))
}

/// Handles the `PUT /me/tokens` route.
pub fn new(req: &mut dyn RequestExt) -> EndpointResult {
    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiToken {
        name: String,
    }

    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiTokenRequest {
        api_token: NewApiToken,
    }

    let max_size = 2000;
    let length = req
        .content_length()
        .chain_error(|| bad_request("missing header: Content-Length"))?;

    if length > max_size {
        return Err(bad_request(&format!("max content length is: {}", max_size)));
    }

    let mut json = vec![0; length as usize];
    read_fill(req.body(), &mut json)?;

    let json =
        String::from_utf8(json).map_err(|_| bad_request(&"json body was not valid utf-8"))?;

    let new: NewApiTokenRequest = json::from_str(&json)
        .map_err(|e| bad_request(&format!("invalid new token request: {:?}", e)))?;

    let name = &new.api_token.name;
    if name.is_empty() {
        return Err(bad_request("name must have a value"));
    }

    let authenticated_user = req.authenticate()?;
    if authenticated_user.api_token_id().is_some() {
        return Err(bad_request(
            "cannot use an API token to create a new API token",
        ));
    }

    let conn = req.db_conn()?;
    let user = authenticated_user.user();

    let max_token_per_user = 500;
    let count: i64 = ApiToken::belonging_to(&user).count().get_result(&*conn)?;
    if count >= max_token_per_user {
        return Err(bad_request(&format!(
            "maximum tokens per user is: {}",
            max_token_per_user
        )));
    }

    let api_token = ApiToken::insert(&*conn, user.id, name)?;

    #[derive(Serialize)]
    struct R {
        api_token: EncodableApiTokenWithToken,
    }
    Ok(req.json(&R {
        api_token: api_token.into(),
    }))
}

/// Handles the `DELETE /me/tokens/:id` route.
pub fn revoke(req: &mut dyn RequestExt) -> EndpointResult {
    let id = req.params()["id"]
        .parse::<i32>()
        .map_err(|e| bad_request(&format!("invalid token id: {:?}", e)))?;

    let authenticated_user = req.authenticate()?;
    let conn = req.db_conn()?;
    let user = authenticated_user.user();
    diesel::update(ApiToken::belonging_to(&user).find(id))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    #[derive(Serialize)]
    struct R {}
    Ok(req.json(&R {}))
}

/// Handles the `DELETE /tokens/current` route.
pub fn revoke_current(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?;
    let api_token_id = authenticated_user
        .api_token_id()
        .ok_or_else(|| bad_request("token not provided"))?;

    let conn = req.db_conn()?;
    diesel::update(api_tokens::table.filter(api_tokens::id.eq(api_token_id)))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(Response::builder().status(204).body(Body::empty()).unwrap())
}
<commit_msg>controllers::token: Use `json!` macro to simplify JSON serialization code<commit_after>use super::frontend_prelude::*;

use crate::models::ApiToken;
use crate::schema::api_tokens;
use crate::util::read_fill;
use crate::views::EncodableApiTokenWithToken;

use conduit::{Body, Response};
use serde_json as json;

/// Handles the `GET /me/tokens` route.
pub fn list(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let conn = req.db_conn()?;
    let user = authenticated_user.user();

    let tokens: Vec<ApiToken> = ApiToken::belonging_to(&user)
        .filter(api_tokens::revoked.eq(false))
        .order(api_tokens::created_at.desc())
        .load(&*conn)?;

    Ok(req.json(&json!({ "api_tokens": tokens })))
}

/// Handles the `PUT /me/tokens` route.
pub fn new(req: &mut dyn RequestExt) -> EndpointResult {
    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiToken {
        name: String,
    }

    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiTokenRequest {
        api_token: NewApiToken,
    }

    let max_size = 2000;
    let length = req
        .content_length()
        .chain_error(|| bad_request("missing header: Content-Length"))?;

    if length > max_size {
        return Err(bad_request(&format!("max content length is: {}", max_size)));
    }

    let mut json = vec![0; length as usize];
    read_fill(req.body(), &mut json)?;

    let json =
        String::from_utf8(json).map_err(|_| bad_request(&"json body was not valid utf-8"))?;

    let new: NewApiTokenRequest = json::from_str(&json)
        .map_err(|e| bad_request(&format!("invalid new token request: {:?}", e)))?;

    let name = &new.api_token.name;
    if name.is_empty() {
        return Err(bad_request("name must have a value"));
    }

    let authenticated_user = req.authenticate()?;
    if authenticated_user.api_token_id().is_some() {
        return Err(bad_request(
            "cannot use an API token to create a new API token",
        ));
    }

    let conn = req.db_conn()?;
    let user = authenticated_user.user();

    let max_token_per_user = 500;
    let count: i64 = ApiToken::belonging_to(&user).count().get_result(&*conn)?;
    if count >= max_token_per_user {
        return Err(bad_request(&format!(
            "maximum tokens per user is: {}",
            max_token_per_user
        )));
    }

    let api_token = ApiToken::insert(&*conn, user.id, name)?;
    let api_token = EncodableApiTokenWithToken::from(api_token);

    Ok(req.json(&json!({ "api_token": api_token })))
}

/// Handles the `DELETE /me/tokens/:id` route.
pub fn revoke(req: &mut dyn RequestExt) -> EndpointResult {
    let id = req.params()["id"]
        .parse::<i32>()
        .map_err(|e| bad_request(&format!("invalid token id: {:?}", e)))?;

    let authenticated_user = req.authenticate()?;
    let conn = req.db_conn()?;
    let user = authenticated_user.user();
    diesel::update(ApiToken::belonging_to(&user).find(id))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(req.json(&json!({})))
}

/// Handles the `DELETE /tokens/current` route.
pub fn revoke_current(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?;
    let api_token_id = authenticated_user
        .api_token_id()
        .ok_or_else(|| bad_request("token not provided"))?;

    let conn = req.db_conn()?;
    diesel::update(api_tokens::table.filter(api_tokens::id.eq(api_token_id)))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(Response::builder().status(204).body(Body::empty()).unwrap())
}
<|endoftext|>


<commit_before>//! Contains the ThreadPool and the individual Threads.

use std::heap::{Alloc, Layout, Global, oom};
use std::sync::atomic::{AtomicBool,Ordering};
use std::thread::{JoinHandle,self};
use std::sync::{Once, ONCE_INIT};
use std::ptr::NonNull;
use std::ptr;
use std::cell::UnsafeCell;

use crossbeam_utils::scoped;

use pleco::MoveList;
use pleco::tools::pleco_arc::Arc;
use pleco::board::*;
use pleco::core::piece_move::BitMove;

use sync::LockLatch;
use time::uci_timer::*;
use search::Searcher;

use consts::*;

const KILOBYTE: usize = 1000;
const THREAD_STACK_SIZE: usize = 18000 * KILOBYTE;

// The Global threadpool!
pub static mut THREADPOOL: NonNull<ThreadPool> = unsafe {NonNull::new_unchecked(ptr::null_mut())};

static THREADPOOL_INIT: Once = ONCE_INIT;

pub fn init_threadpool() {
    THREADPOOL_INIT.call_once(|| {
        unsafe {
            // We have a spawned thread create all structures, as a stack overflow can
            // occur otherwise
            let builder = thread::Builder::new()
                .name("Starter".to_string())
                .stack_size(THREAD_STACK_SIZE);
            let handle = scoped::builder_spawn_unsafe(builder, move || {
                let layout = Layout::new::<ThreadPool>();
                let result = Global.alloc_zeroed(layout);
                let new_ptr: *mut ThreadPool = match result {
                    Ok(ptr) => ptr.cast().as_ptr() as *mut ThreadPool,
                    Err(_err) => oom(layout),
                };
                ptr::write(new_ptr, ThreadPool::new());
                THREADPOOL = NonNull::new_unchecked(new_ptr);
            });
            handle.unwrap().join().unwrap();
        }
    });
}

/// Returns access to the global thread pool.
pub fn threadpool() -> &'static mut ThreadPool {
    unsafe {
        THREADPOOL.as_mut()
    }
}

// Dummy struct to allow us to pass a pointer into a spawned thread.
struct SearcherPtr {
    ptr: UnsafeCell<*mut Searcher>
}

unsafe impl Sync for SearcherPtr {}
unsafe impl Send for SearcherPtr {}

/// The thread-pool for the chess engine.
pub struct ThreadPool {
    /// Access to each thread's Structure
    pub threads: Vec<UnsafeCell<*mut Searcher>>,
    /// Handles of each thread
    handles: Vec<JoinHandle<()>>,
    /// Condition for the main thread to start.
    pub main_cond: Arc<LockLatch>,
    /// Condition for all non-main threads
    pub thread_cond: Arc<LockLatch>,
    /// Stop condition, if true the threads should halt.
    pub stop: AtomicBool
}

// Okay, this all looks like madness, but there is some reason to it all.
// Basically, `ThreadPool` manages spawning and despawning threads, as well
// as passing state to / from those threads, telling them to stop, go, drop,
// and lastly determining the "best move" from all the threads.
///
// While we spawn all the other threads, We mostly communicate with the
// MainThread to do anything useful. The mainthread handles anything fun.
// The goal of the ThreadPool is to be NON BLOCKING, unless we want to await a
// result.
impl ThreadPool {

    /// Creates a new `ThreadPool`
    pub fn new() -> Self {
        let mut pool: ThreadPool = ThreadPool {
            threads: Vec::new(),
            handles: Vec::new(),
            main_cond: Arc::new(LockLatch::new()),
            thread_cond: Arc::new(LockLatch::new()),
            stop: AtomicBool::new(true)
        };
        // Lock both the cond variables
        pool.main_cond.lock();
        pool.thread_cond.lock();

        // spawn the main thread
        pool.attach_thread();
        pool
    }

    /// Spawns a new thread and appends it to our vector of JoinHandles.
    fn attach_thread(&mut self) {
         unsafe {
             let thread_ptr: SearcherPtr = self.create_thread();
             let builder = thread::Builder::new()
                 .name(self.size().to_string())
                 .stack_size(THREAD_STACK_SIZE);

             let handle = scoped::builder_spawn_unsafe(builder,
                move || {
                    let thread = &mut **thread_ptr.ptr.get();
                    thread.cond.lock();
                    thread.idle_loop();
             }).unwrap();
             self.handles.push(handle);
        };
    }

    /// Allocates a thread structure and pushes it to the threadstack.
    ///
    /// This does not spawn a thread, just creates the structure in the heap for the thread.
    ///
    /// Only to be called by `attach_thread()`.
    fn create_thread(&mut self) -> SearcherPtr {
        let len: usize = self.threads.len();
        let layout = Layout::new::<Searcher>();
        let cond = if len == 0 {self.main_cond.clone()} else {self.thread_cond.clone()};
        unsafe {
            let result = Global.alloc_zeroed(layout);
            let new_ptr: *mut Searcher = match result {
                Ok(ptr) => ptr.cast().as_ptr() as *mut Searcher,
                Err(_err) => oom(layout),
            };
            ptr::write(new_ptr, Searcher::new(len, cond));
            self.threads.push(UnsafeCell::new(new_ptr));
            SearcherPtr {ptr: UnsafeCell::new(new_ptr)}
        }
    }

    /// Returns the number of threads
    #[inline(always)]
    pub fn size(&self) -> usize {
        self.threads.len()
    }

    /// Returns a pointer to the main thread.
    fn main(&mut self) -> &mut Searcher {
        unsafe {
            let main_thread: *mut Searcher = *self.threads.get_unchecked(0).get();
            return &mut *main_thread;
        }
    }

    /// Sets the use of standard out. This can be changed mid search as well.
    #[inline(always)]
    pub fn stdout(&mut self, use_stdout: bool) {
        USE_STDOUT.store(use_stdout, Ordering::Relaxed);
    }

    /// Sets the thread count of the pool. If num is less than 1, nothing will happen.
    ///
    /// # Safety
    ///
    /// Completely unsafe to use when the pool is searching.
    pub fn set_thread_count(&mut self, mut num: usize) {
        if num >= 1 {
            num = num.min(MAX_THREADS);
            self.wait_for_finish();
            self.kill_all();
            while self.size() < num {
                self.attach_thread();
            }
        }
    }

    /// Kills and de-allocates all the threads that are running. This function will also
    /// block on waiting for the search to finish.
    pub fn kill_all(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        self.wait_for_finish();
        let mut join_handles = Vec::with_capacity(self.size());
        unsafe {
            // tell each thread to drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| {
                    s.kill.store(true, Ordering::SeqCst)
                });

            // If awaiting a signal, wake up each thread so each can drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| {
                    s.cond.set();
                });

            // Start connecting each join handle. We don't unwrap here, as if one of the
            // threads fail, the other threads remain un-joined.
            while let Some(handle) = self.handles.pop() {
                join_handles.push(handle.join());
            }

            // De-allocate each thread.
            while let Some(unc) = self.threads.pop() {
                let th: *mut Searcher = *unc.get();
                let layout = Layout::new::<Searcher>();
                Global.dealloc(NonNull::new_unchecked(th).as_opaque(), layout);
            }
        }

        // Unwrap the results from each `thread::join`,
        while let Some(handle_result) = join_handles.pop() {
            handle_result.unwrap_or_else(|e| println!("Thread failed: {:?}",e));
        }
    }

    /// Sets the threads to stop (or not!).
    #[inline(always)]
    pub fn set_stop(&mut self, stop: bool) {
        self.stop.store(stop, Ordering::Relaxed);
    }

    /// Waits for all the threads to finish
    pub fn wait_for_finish(&self) {
        unsafe {
            self.threads.iter()
                .map(|t| &**t.get())
                .for_each(|t| t.searching.await(false));
        }
    }

    /// Waits for all the threads to start.
    pub fn wait_for_start(&self) {
        unsafe {
            self.threads.iter()
                .map(|t| &**t.get())
                .for_each(|t| t.searching.await(true));
        }
    }

    /// Waits for all non-main threads to finish.
    pub fn wait_for_non_main(&self) {
        unsafe {
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|t: &Searcher|{
                    if t.id != 0 {
                        t.searching.await(false);
                    }
            });
        }
    }

    /// Waits for all the non-main threads to start running
    pub fn wait_for_main_start(&self) {
        unsafe {
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|t: &Searcher|{
                    if t.id == 0 {
                        t.searching.await(true);
                    }
                });
        }
    }

    pub fn clear_all(&mut self) {
        for thread_ptr in self.threads.iter_mut() {
            let mut thread: &mut Searcher = unsafe { &mut **(*thread_ptr).get() };
            thread.clear();
        }
    }

    /// Starts a UCI search. The result will be printed to stdout if the stdout setting
    /// is true.
    pub fn uci_search(&mut self, board: &Board, limits: &Limits) {

        // Start the timer!
        if let Some(uci_timer) = limits.use_time_management() {
            timer().init(limits.start, &uci_timer, board.turn(), board.moves_played());
        } else {
            timer().start_timer(limits.start);
        }

        let root_moves: MoveList = board.generate_moves();

        assert!(!root_moves.is_empty());
        self.wait_for_finish();
        self.stop.store(false, Ordering::Relaxed);

        for thread_ptr in self.threads.iter_mut() {
            let mut thread: &mut Searcher = unsafe {&mut **(*thread_ptr).get()};
            thread.nodes.store(0, Ordering::Relaxed);
            thread.depth_completed = 0;
            thread.board = board.shallow_clone();
            thread.limit = limits.clone();
            thread.root_moves().replace(&root_moves);
        }

        self.main_cond.set();
        self.wait_for_main_start();
        self.main_cond.lock();
    }

    /// Performs a standard search, and blocks waiting for a returned `BitMove`.
    pub fn search(&mut self, board: &Board, limits: &Limits) -> BitMove {
        self.uci_search(board, limits);
        self.wait_for_finish();
        self.best_move()
    }

    /// Returns the best move of a search
    pub fn best_move(&mut self) -> BitMove {
        self.main().root_moves().get(0).unwrap().bit_move
    }

    /// Returns total number of nodes searched so far.
    pub fn nodes(&self) -> u64 {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .map(|s: &Searcher| s.nodes.load(Ordering::Relaxed))
            .sum()
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        self.kill_all();
    }
}

<commit_msg>Experiments for having the global threadpool being a fixed location in memory.<commit_after>//! Contains the ThreadPool and the individual Threads.

use std::heap::{Alloc, Layout, Global, oom};
use std::sync::atomic::{AtomicBool,Ordering};
use std::thread::{JoinHandle,self};
use std::sync::{Once, ONCE_INIT};
use std::ptr::NonNull;
use std::{ptr,mem};
use std::cell::UnsafeCell;

use crossbeam_utils::scoped;

use pleco::MoveList;
use pleco::tools::pleco_arc::Arc;
use pleco::board::*;
use pleco::core::piece_move::BitMove;

use sync::LockLatch;
use time::uci_timer::*;
use search::Searcher;

use consts::*;

const KILOBYTE: usize = 1000;
const THREAD_STACK_SIZE: usize = 18000 * KILOBYTE;

// The Global threadpool!
const POOL_SIZE: usize = mem::size_of::<ThreadPool>();
type DummyThreadPool = [u8; POOL_SIZE];

pub static mut THREADPOOL: DummyThreadPool = [0; POOL_SIZE];

static THREADPOOL_INIT: Once = ONCE_INIT;

pub fn init_threadpool() {
    THREADPOOL_INIT.call_once(|| {
        unsafe {
            // We have a spawned thread create all structures, as a stack overflow can
            // occur otherwise
            let builder = thread::Builder::new()
                .name("Starter".to_string())
                .stack_size(THREAD_STACK_SIZE);

            let handle = scoped::builder_spawn_unsafe(builder, move || {
                let pool: *mut ThreadPool = mem::transmute(&mut THREADPOOL);
                ptr::write(pool, ThreadPool::new());

            });
            handle.unwrap().join().unwrap();
        }
    });
}

/// Returns access to the global thread pool.
#[inline(always)]
pub fn threadpool() -> &'static mut ThreadPool {
    unsafe {
        mem::transmute::<&mut DummyThreadPool, &'static mut ThreadPool>(&mut THREADPOOL)
    }
}

// Dummy struct to allow us to pass a pointer into a spawned thread.
struct SearcherPtr {
    ptr: UnsafeCell<*mut Searcher>
}

unsafe impl Sync for SearcherPtr {}
unsafe impl Send for SearcherPtr {}

/// The thread-pool for the chess engine.
pub struct ThreadPool {
    /// Access to each thread's Structure
    pub threads: Vec<UnsafeCell<*mut Searcher>>,
    /// Handles of each thread
    handles: Vec<JoinHandle<()>>,
    /// Condition for the main thread to start.
    pub main_cond: Arc<LockLatch>,
    /// Condition for all non-main threads
    pub thread_cond: Arc<LockLatch>,
    /// Stop condition, if true the threads should halt.
    pub stop: AtomicBool
}

// Okay, this all looks like madness, but there is some reason to it all.
// Basically, `ThreadPool` manages spawning and despawning threads, as well
// as passing state to / from those threads, telling them to stop, go, drop,
// and lastly determining the "best move" from all the threads.
///
// While we spawn all the other threads, We mostly communicate with the
// MainThread to do anything useful. The mainthread handles anything fun.
// The goal of the ThreadPool is to be NON BLOCKING, unless we want to await a
// result.
impl ThreadPool {

    /// Creates a new `ThreadPool`
    pub fn new() -> Self {
        let mut pool: ThreadPool = ThreadPool {
            threads: Vec::new(),
            handles: Vec::new(),
            main_cond: Arc::new(LockLatch::new()),
            thread_cond: Arc::new(LockLatch::new()),
            stop: AtomicBool::new(true)
        };
        // Lock both the cond variables
        pool.main_cond.lock();
        pool.thread_cond.lock();

        // spawn the main thread
        pool.attach_thread();
        pool
    }

    /// Spawns a new thread and appends it to our vector of JoinHandles.
    fn attach_thread(&mut self) {
         unsafe {
             let thread_ptr: SearcherPtr = self.create_thread();
             let builder = thread::Builder::new()
                 .name(self.size().to_string())
                 .stack_size(THREAD_STACK_SIZE);

             let handle = scoped::builder_spawn_unsafe(builder,
                move || {
                    let thread = &mut **thread_ptr.ptr.get();
                    thread.cond.lock();
                    thread.idle_loop();
             }).unwrap();
             self.handles.push(handle);
        };
    }

    /// Allocates a thread structure and pushes it to the threadstack.
    ///
    /// This does not spawn a thread, just creates the structure in the heap for the thread.
    ///
    /// Only to be called by `attach_thread()`.
    fn create_thread(&mut self) -> SearcherPtr {
        let len: usize = self.threads.len();
        let layout = Layout::new::<Searcher>();
        let cond = if len == 0 {self.main_cond.clone()} else {self.thread_cond.clone()};
        unsafe {
            let result = Global.alloc_zeroed(layout);
            let new_ptr: *mut Searcher = match result {
                Ok(ptr) => ptr.cast().as_ptr() as *mut Searcher,
                Err(_err) => oom(layout),
            };
            ptr::write(new_ptr, Searcher::new(len, cond));
            self.threads.push(UnsafeCell::new(new_ptr));
            SearcherPtr {ptr: UnsafeCell::new(new_ptr)}
        }
    }

    /// Returns the number of threads
    #[inline(always)]
    pub fn size(&self) -> usize {
        self.threads.len()
    }

    /// Returns a pointer to the main thread.
    fn main(&mut self) -> &mut Searcher {
        unsafe {
            let main_thread: *mut Searcher = *self.threads.get_unchecked(0).get();
            return &mut *main_thread;
        }
    }

    /// Sets the use of standard out. This can be changed mid search as well.
    #[inline(always)]
    pub fn stdout(&mut self, use_stdout: bool) {
        USE_STDOUT.store(use_stdout, Ordering::Relaxed);
    }

    /// Sets the thread count of the pool. If num is less than 1, nothing will happen.
    ///
    /// # Safety
    ///
    /// Completely unsafe to use when the pool is searching.
    pub fn set_thread_count(&mut self, mut num: usize) {
        if num >= 1 {
            num = num.min(MAX_THREADS);
            self.wait_for_finish();
            self.kill_all();
            while self.size() < num {
                self.attach_thread();
            }
        }
    }

    /// Kills and de-allocates all the threads that are running. This function will also
    /// block on waiting for the search to finish.
    pub fn kill_all(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        self.wait_for_finish();
        let mut join_handles = Vec::with_capacity(self.size());
        unsafe {
            // tell each thread to drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| {
                    s.kill.store(true, Ordering::SeqCst)
                });

            // If awaiting a signal, wake up each thread so each can drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| {
                    s.cond.set();
                });

            // Start connecting each join handle. We don't unwrap here, as if one of the
            // threads fail, the other threads remain un-joined.
            while let Some(handle) = self.handles.pop() {
                join_handles.push(handle.join());
            }

            // De-allocate each thread.
            while let Some(unc) = self.threads.pop() {
                let th: *mut Searcher = *unc.get();
                let layout = Layout::new::<Searcher>();
                Global.dealloc(NonNull::new_unchecked(th).as_opaque(), layout);
            }
        }

        // Unwrap the results from each `thread::join`,
        while let Some(handle_result) = join_handles.pop() {
            handle_result.unwrap_or_else(|e| println!("Thread failed: {:?}",e));
        }
    }

    /// Sets the threads to stop (or not!).
    #[inline(always)]
    pub fn set_stop(&mut self, stop: bool) {
        self.stop.store(stop, Ordering::Relaxed);
    }

    /// Waits for all the threads to finish
    pub fn wait_for_finish(&self) {
        unsafe {
            self.threads.iter()
                .map(|t| &**t.get())
                .for_each(|t| t.searching.await(false));
        }
    }

    /// Waits for all the threads to start.
    pub fn wait_for_start(&self) {
        unsafe {
            self.threads.iter()
                .map(|t| &**t.get())
                .for_each(|t| t.searching.await(true));
        }
    }

    /// Waits for all non-main threads to finish.
    pub fn wait_for_non_main(&self) {
        unsafe {
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|t: &Searcher|{
                    if t.id != 0 {
                        t.searching.await(false);
                    }
            });
        }
    }

    /// Waits for all the non-main threads to start running
    pub fn wait_for_main_start(&self) {
        unsafe {
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|t: &Searcher|{
                    if t.id == 0 {
                        t.searching.await(true);
                    }
                });
        }
    }

    pub fn clear_all(&mut self) {
        for thread_ptr in self.threads.iter_mut() {
            let mut thread: &mut Searcher = unsafe { &mut **(*thread_ptr).get() };
            thread.clear();
        }
    }

    /// Starts a UCI search. The result will be printed to stdout if the stdout setting
    /// is true.
    pub fn uci_search(&mut self, board: &Board, limits: &Limits) {

        // Start the timer!
        if let Some(uci_timer) = limits.use_time_management() {
            timer().init(limits.start, &uci_timer, board.turn(), board.moves_played());
        } else {
            timer().start_timer(limits.start);
        }

        let root_moves: MoveList = board.generate_moves();

        assert!(!root_moves.is_empty());
        self.wait_for_finish();
        self.stop.store(false, Ordering::Relaxed);

        for thread_ptr in self.threads.iter_mut() {
            let mut thread: &mut Searcher = unsafe {&mut **(*thread_ptr).get()};
            thread.nodes.store(0, Ordering::Relaxed);
            thread.depth_completed = 0;
            thread.board = board.shallow_clone();
            thread.limit = limits.clone();
            thread.root_moves().replace(&root_moves);
        }

        self.main_cond.set();
        self.wait_for_main_start();
        self.main_cond.lock();
    }

    /// Performs a standard search, and blocks waiting for a returned `BitMove`.
    pub fn search(&mut self, board: &Board, limits: &Limits) -> BitMove {
        self.uci_search(board, limits);
        self.wait_for_finish();
        self.best_move()
    }

    /// Returns the best move of a search
    pub fn best_move(&mut self) -> BitMove {
        self.main().root_moves().get(0).unwrap().bit_move
    }

    /// Returns total number of nodes searched so far.
    pub fn nodes(&self) -> u64 {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .map(|s: &Searcher| s.nodes.load(Ordering::Relaxed))
            .sum()
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        self.kill_all();
    }
}

<|endoftext|>


<commit_before>use byteorder::WriteBytesExt;
use std::io;

pub struct RangeEncoder<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W,
    range: u32,
    low: u64,
    cache: u8,
    cachesz: u32,
}

impl<'a, W> RangeEncoder<'a, W>
where
    W: io::Write,
{
    pub fn new(stream: &'a mut W) -> Self {
        let enc = Self {
            stream,
            range: 0xFFFF_FFFF,
            low: 0,
            cache: 0,
            cachesz: 1,
        };
        lzma_debug!("0 {{ range: {:08x}, low: {:010x} }}", enc.range, enc.low);
        enc
    }

    fn write_low(&mut self) -> io::Result<()> {
        if self.low < 0xFF00_0000 || self.low > 0xFFFF_FFFF {
            let mut tmp = self.cache;
            loop {
                let byte = tmp.wrapping_add((self.low >> 32) as u8);
                self.stream.write_u8(byte)?;
                lzma_debug!("> byte: {:02x}", byte);
                tmp = 0xFF;
                self.cachesz -= 1;
                if self.cachesz == 0 {
                    break;
                }
            }
            self.cache = (self.low >> 24) as u8;
        }

        self.cachesz += 1;
        self.low = (self.low << 8) & 0xFFFF_FFFF;
        Ok(())
    }

    pub fn finish(&mut self) -> io::Result<()> {
        for _ in 0..5 {
            self.write_low()?;

            lzma_debug!("$ {{ range: {:08x}, low: {:010x} }}", self.range, self.low);
        }
        Ok(())
    }

    fn normalize(&mut self) -> io::Result<()> {
        while self.range < 0x0100_0000 {
            lzma_debug!(
                "+ {{ range: {:08x}, low: {:010x}, cache: {:02x}, {} }}",
                self.range,
                self.low,
                self.cache,
                self.cachesz
            );
            self.range <<= 8;
            self.write_low()?;
            lzma_debug!(
                "* {{ range: {:08x}, low: {:010x}, cache: {:02x}, {} }}",
                self.range,
                self.low,
                self.cache,
                self.cachesz
            );
        }
        lzma_trace!("  {{ range: {:08x}, low: {:010x} }}", self.range, self.low);
        Ok(())
    }

    pub fn encode_bit(&mut self, prob: &mut u16, bit: bool) -> io::Result<()> {
        let bound: u32 = (self.range >> 11) * (*prob as u32);
        lzma_trace!(
            "  bound: {:08x}, prob: {:04x}, bit: {}",
            bound,
            prob,
            bit as u8
        );

        if bit {
            *prob -= *prob >> 5;
            self.low += bound as u64;
            self.range -= bound;
        } else {
            *prob += (0x800_u16 - *prob) >> 5;
            self.range = bound;
        }

        self.normalize()
    }

    #[cfg(test)]
    fn encode_bit_tree(
        &mut self,
        num_bits: usize,
        probs: &mut [u16],
        value: u32,
    ) -> io::Result<()> {
        debug_assert!(value.leading_zeros() as usize + num_bits >= 32);
        let mut tmp: usize = 1;
        for i in 0..num_bits {
            let bit = ((value >> (num_bits - i - 1)) & 1) != 0;
            self.encode_bit(&mut probs[tmp], bit)?;
            tmp = (tmp << 1) ^ (bit as usize);
        }
        Ok(())
    }

    #[cfg(test)]
    pub fn encode_reverse_bit_tree(
        &mut self,
        num_bits: usize,
        probs: &mut [u16],
        offset: usize,
        mut value: u32,
    ) -> io::Result<()> {
        debug_assert!(value.leading_zeros() as usize + num_bits >= 32);
        let mut tmp: usize = 1;
        for _ in 0..num_bits {
            let bit = (value & 1) != 0;
            value >>= 1;
            self.encode_bit(&mut probs[offset + tmp], bit)?;
            tmp = (tmp << 1) ^ (bit as usize);
        }
        Ok(())
    }
}

// TODO: parametrize by constant and use [u16; 1 << num_bits] as soon as Rust supports this
#[cfg(test)]
#[derive(Clone)]
pub struct BitTree {
    num_bits: usize,
    probs: Vec<u16>,
}

#[cfg(test)]
impl BitTree {
    pub fn new(num_bits: usize) -> Self {
        BitTree {
            num_bits,
            probs: vec![0x400; 1 << num_bits],
        }
    }

    pub fn encode<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        value: u32,
    ) -> io::Result<()> {
        rangecoder.encode_bit_tree(self.num_bits, self.probs.as_mut_slice(), value)
    }

    pub fn encode_reverse<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        value: u32,
    ) -> io::Result<()> {
        rangecoder.encode_reverse_bit_tree(self.num_bits, self.probs.as_mut_slice(), 0, value)
    }
}

#[cfg(test)]
pub struct LenEncoder {
    choice: u16,
    choice2: u16,
    low_coder: Vec<BitTree>,
    mid_coder: Vec<BitTree>,
    high_coder: BitTree,
}

#[cfg(test)]
impl LenEncoder {
    pub fn new() -> Self {
        LenEncoder {
            choice: 0x400,
            choice2: 0x400,
            low_coder: vec![BitTree::new(3); 16],
            mid_coder: vec![BitTree::new(3); 16],
            high_coder: BitTree::new(8),
        }
    }

    pub fn encode<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        pos_state: usize,
        value: u32,
    ) -> io::Result<()> {
        let is_low: bool = value < 8;
        rangecoder.encode_bit(&mut self.choice, !is_low)?;
        if is_low {
            return self.low_coder[pos_state].encode(rangecoder, value);
        }

        let is_middle: bool = value < 16;
        rangecoder.encode_bit(&mut self.choice2, !is_middle)?;
        if is_middle {
            return self.mid_coder[pos_state].encode(rangecoder, value - 8);
        }

        self.high_coder.encode(rangecoder, value - 16)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::decode::rangecoder::{LenDecoder, RangeDecoder};
    use crate::{decode, encode};
    use std::io::BufReader;

    fn encode_decode(prob_init: u16, bits: &[bool]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut prob = prob_init;
        for &b in bits {
            encoder.encode_bit(&mut prob, b).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut prob = prob_init;
        for &b in bits {
            assert_eq!(decoder.decode_bit(&mut prob).unwrap(), b);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_zeros() {
        encode_decode(0x400, &[false; 10000]);
    }

    #[test]
    fn test_encode_decode_ones() {
        encode_decode(0x400, &[true; 10000]);
    }

    fn encode_decode_bittree(num_bits: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut tree = encode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            tree.encode(&mut encoder, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut tree = decode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            assert_eq!(tree.parse(&mut decoder).unwrap(), v);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_bittree_zeros() {
        for num_bits in 0..16 {
            encode_decode_bittree(num_bits, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_bittree_ones() {
        for num_bits in 0..16 {
            encode_decode_bittree(num_bits, &[(1 << num_bits) - 1; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_bittree_all() {
        for num_bits in 0..16 {
            let max = 1 << num_bits;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_bittree(num_bits, &values);
        }
    }

    fn encode_decode_reverse_bittree(num_bits: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut tree = encode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            tree.encode_reverse(&mut encoder, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut tree = decode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            assert_eq!(tree.parse_reverse(&mut decoder).unwrap(), v);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_reverse_bittree_zeros() {
        for num_bits in 0..16 {
            encode_decode_reverse_bittree(num_bits, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_reverse_bittree_ones() {
        for num_bits in 0..16 {
            encode_decode_reverse_bittree(num_bits, &[(1 << num_bits) - 1; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_reverse_bittree_all() {
        for num_bits in 0..16 {
            let max = 1 << num_bits;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_reverse_bittree(num_bits, &values);
        }
    }

    fn encode_decode_length(pos_state: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut len_encoder = LenEncoder::new();
        for &v in values {
            len_encoder.encode(&mut encoder, pos_state, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut len_decoder = LenDecoder::new();
        for &v in values {
            assert_eq!(
                len_decoder.decode(&mut decoder, pos_state).unwrap(),
                v as usize
            );
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_length_zeros() {
        for pos_state in 0..16 {
            encode_decode_length(pos_state, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_length_all() {
        for pos_state in 0..16 {
            let max = (1 << 8) + 16;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_length(pos_state, &values);
        }
    }
}
<commit_msg>Blacklist clippy false positive due to empty macro.<commit_after>use byteorder::WriteBytesExt;
use std::io;

pub struct RangeEncoder<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W,
    range: u32,
    low: u64,
    cache: u8,
    cachesz: u32,
}

impl<'a, W> RangeEncoder<'a, W>
where
    W: io::Write,
{
    #[allow(clippy::let_and_return)]
    pub fn new(stream: &'a mut W) -> Self {
        let enc = Self {
            stream,
            range: 0xFFFF_FFFF,
            low: 0,
            cache: 0,
            cachesz: 1,
        };
        lzma_debug!("0 {{ range: {:08x}, low: {:010x} }}", enc.range, enc.low);
        enc
    }

    fn write_low(&mut self) -> io::Result<()> {
        if self.low < 0xFF00_0000 || self.low > 0xFFFF_FFFF {
            let mut tmp = self.cache;
            loop {
                let byte = tmp.wrapping_add((self.low >> 32) as u8);
                self.stream.write_u8(byte)?;
                lzma_debug!("> byte: {:02x}", byte);
                tmp = 0xFF;
                self.cachesz -= 1;
                if self.cachesz == 0 {
                    break;
                }
            }
            self.cache = (self.low >> 24) as u8;
        }

        self.cachesz += 1;
        self.low = (self.low << 8) & 0xFFFF_FFFF;
        Ok(())
    }

    pub fn finish(&mut self) -> io::Result<()> {
        for _ in 0..5 {
            self.write_low()?;

            lzma_debug!("$ {{ range: {:08x}, low: {:010x} }}", self.range, self.low);
        }
        Ok(())
    }

    fn normalize(&mut self) -> io::Result<()> {
        while self.range < 0x0100_0000 {
            lzma_debug!(
                "+ {{ range: {:08x}, low: {:010x}, cache: {:02x}, {} }}",
                self.range,
                self.low,
                self.cache,
                self.cachesz
            );
            self.range <<= 8;
            self.write_low()?;
            lzma_debug!(
                "* {{ range: {:08x}, low: {:010x}, cache: {:02x}, {} }}",
                self.range,
                self.low,
                self.cache,
                self.cachesz
            );
        }
        lzma_trace!("  {{ range: {:08x}, low: {:010x} }}", self.range, self.low);
        Ok(())
    }

    pub fn encode_bit(&mut self, prob: &mut u16, bit: bool) -> io::Result<()> {
        let bound: u32 = (self.range >> 11) * (*prob as u32);
        lzma_trace!(
            "  bound: {:08x}, prob: {:04x}, bit: {}",
            bound,
            prob,
            bit as u8
        );

        if bit {
            *prob -= *prob >> 5;
            self.low += bound as u64;
            self.range -= bound;
        } else {
            *prob += (0x800_u16 - *prob) >> 5;
            self.range = bound;
        }

        self.normalize()
    }

    #[cfg(test)]
    fn encode_bit_tree(
        &mut self,
        num_bits: usize,
        probs: &mut [u16],
        value: u32,
    ) -> io::Result<()> {
        debug_assert!(value.leading_zeros() as usize + num_bits >= 32);
        let mut tmp: usize = 1;
        for i in 0..num_bits {
            let bit = ((value >> (num_bits - i - 1)) & 1) != 0;
            self.encode_bit(&mut probs[tmp], bit)?;
            tmp = (tmp << 1) ^ (bit as usize);
        }
        Ok(())
    }

    #[cfg(test)]
    pub fn encode_reverse_bit_tree(
        &mut self,
        num_bits: usize,
        probs: &mut [u16],
        offset: usize,
        mut value: u32,
    ) -> io::Result<()> {
        debug_assert!(value.leading_zeros() as usize + num_bits >= 32);
        let mut tmp: usize = 1;
        for _ in 0..num_bits {
            let bit = (value & 1) != 0;
            value >>= 1;
            self.encode_bit(&mut probs[offset + tmp], bit)?;
            tmp = (tmp << 1) ^ (bit as usize);
        }
        Ok(())
    }
}

// TODO: parametrize by constant and use [u16; 1 << num_bits] as soon as Rust supports this
#[cfg(test)]
#[derive(Clone)]
pub struct BitTree {
    num_bits: usize,
    probs: Vec<u16>,
}

#[cfg(test)]
impl BitTree {
    pub fn new(num_bits: usize) -> Self {
        BitTree {
            num_bits,
            probs: vec![0x400; 1 << num_bits],
        }
    }

    pub fn encode<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        value: u32,
    ) -> io::Result<()> {
        rangecoder.encode_bit_tree(self.num_bits, self.probs.as_mut_slice(), value)
    }

    pub fn encode_reverse<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        value: u32,
    ) -> io::Result<()> {
        rangecoder.encode_reverse_bit_tree(self.num_bits, self.probs.as_mut_slice(), 0, value)
    }
}

#[cfg(test)]
pub struct LenEncoder {
    choice: u16,
    choice2: u16,
    low_coder: Vec<BitTree>,
    mid_coder: Vec<BitTree>,
    high_coder: BitTree,
}

#[cfg(test)]
impl LenEncoder {
    pub fn new() -> Self {
        LenEncoder {
            choice: 0x400,
            choice2: 0x400,
            low_coder: vec![BitTree::new(3); 16],
            mid_coder: vec![BitTree::new(3); 16],
            high_coder: BitTree::new(8),
        }
    }

    pub fn encode<W: io::Write>(
        &mut self,
        rangecoder: &mut RangeEncoder<W>,
        pos_state: usize,
        value: u32,
    ) -> io::Result<()> {
        let is_low: bool = value < 8;
        rangecoder.encode_bit(&mut self.choice, !is_low)?;
        if is_low {
            return self.low_coder[pos_state].encode(rangecoder, value);
        }

        let is_middle: bool = value < 16;
        rangecoder.encode_bit(&mut self.choice2, !is_middle)?;
        if is_middle {
            return self.mid_coder[pos_state].encode(rangecoder, value - 8);
        }

        self.high_coder.encode(rangecoder, value - 16)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::decode::rangecoder::{LenDecoder, RangeDecoder};
    use crate::{decode, encode};
    use std::io::BufReader;

    fn encode_decode(prob_init: u16, bits: &[bool]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut prob = prob_init;
        for &b in bits {
            encoder.encode_bit(&mut prob, b).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut prob = prob_init;
        for &b in bits {
            assert_eq!(decoder.decode_bit(&mut prob).unwrap(), b);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_zeros() {
        encode_decode(0x400, &[false; 10000]);
    }

    #[test]
    fn test_encode_decode_ones() {
        encode_decode(0x400, &[true; 10000]);
    }

    fn encode_decode_bittree(num_bits: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut tree = encode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            tree.encode(&mut encoder, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut tree = decode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            assert_eq!(tree.parse(&mut decoder).unwrap(), v);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_bittree_zeros() {
        for num_bits in 0..16 {
            encode_decode_bittree(num_bits, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_bittree_ones() {
        for num_bits in 0..16 {
            encode_decode_bittree(num_bits, &[(1 << num_bits) - 1; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_bittree_all() {
        for num_bits in 0..16 {
            let max = 1 << num_bits;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_bittree(num_bits, &values);
        }
    }

    fn encode_decode_reverse_bittree(num_bits: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut tree = encode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            tree.encode_reverse(&mut encoder, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut tree = decode::rangecoder::BitTree::new(num_bits);
        for &v in values {
            assert_eq!(tree.parse_reverse(&mut decoder).unwrap(), v);
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_reverse_bittree_zeros() {
        for num_bits in 0..16 {
            encode_decode_reverse_bittree(num_bits, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_reverse_bittree_ones() {
        for num_bits in 0..16 {
            encode_decode_reverse_bittree(num_bits, &[(1 << num_bits) - 1; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_reverse_bittree_all() {
        for num_bits in 0..16 {
            let max = 1 << num_bits;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_reverse_bittree(num_bits, &values);
        }
    }

    fn encode_decode_length(pos_state: usize, values: &[u32]) {
        let mut buf: Vec<u8> = Vec::new();

        let mut encoder = RangeEncoder::new(&mut buf);
        let mut len_encoder = LenEncoder::new();
        for &v in values {
            len_encoder.encode(&mut encoder, pos_state, v).unwrap();
        }
        encoder.finish().unwrap();

        let mut bufread = BufReader::new(buf.as_slice());
        let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
        let mut len_decoder = LenDecoder::new();
        for &v in values {
            assert_eq!(
                len_decoder.decode(&mut decoder, pos_state).unwrap(),
                v as usize
            );
        }
        assert!(decoder.is_finished_ok().unwrap());
    }

    #[test]
    fn test_encode_decode_length_zeros() {
        for pos_state in 0..16 {
            encode_decode_length(pos_state, &[0; 10000]);
        }
    }

    #[test]
    fn test_encode_decode_length_all() {
        for pos_state in 0..16 {
            let max = (1 << 8) + 16;
            let values: Vec<u32> = (0..max).collect();
            encode_decode_length(pos_state, &values);
        }
    }
}
<|endoftext|>