<commit_before>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published
// by the Free Software Foundation, either version 2.1 of the License, or
// (at your option) any later version.
//
// Caribon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

use book::Book;
use number::Number;
use error::{Error,Result, Source};
use token::Token;
use zipper::Zipper;
use escape::escape_tex;
use resource_handler::ResourceHandler;
use renderer::Renderer;

use std::iter::Iterator;
use std::fs::File;
use std::io::Read;
use std::borrow::Cow;

use mustache;

/// LaTeX renderer
pub struct LatexRenderer<'a> {
    book: &'a Book,
    current_chapter: Number,
    handler: ResourceHandler<'a>,
    source: Source,
    escape: bool,
    first_letter: bool,
    first_paragraph: bool,
}

impl<'a> LatexRenderer<'a> {
    /// Creates new LatexRenderer
    pub fn new(book: &'a Book) -> LatexRenderer<'a> {
        let mut handler = ResourceHandler::new(&book.logger);
        handler.set_images_mapping(true);
        LatexRenderer {
            book: book,
            current_chapter: Number::Default,
            handler: handler,
            source: Source::empty(),
            escape: true,
            first_letter: false,
            first_paragraph: true,
        }
    }

    /// Render pdf in a file
    pub fn render_pdf(&mut self) -> Result<String> {
        if let Ok(pdf_file) = self.book.options.get_path("output.pdf") {
            let content = try!(self.render_book());
            let mut zipper = try!(Zipper::new(&self.book.options.get_path("temp_dir").unwrap()));
            try!(zipper.write("result.tex", &content.as_bytes(), false));

            // write image files
            for (source, dest) in self.handler.images_mapping() {
                let mut f = try!(File::open(source).map_err(|_| Error::file_not_found(&self.source,
                                                                                "image",
                                                                                source.to_owned())));
                let mut content = vec!();
                try!(f.read_to_end(&mut content).map_err(|e| Error::render(&self.source,
                                                                               format!("error while reading image file: {}", e))));
                try!(zipper.write(dest, &content, true));
            }

            zipper.generate_pdf(&self.book.options.get_str("tex.command").unwrap(), "result.tex", &pdf_file)
        } else {
            Err(Error::render(&self.source,
                                  "no output pdf file specified in book config"))
        }
    }

    /// Render latex in a string
    pub fn render_book(&mut self) -> Result<String> {
        let mut content = String::from("");
        for (i, filename) in self.book.filenames.iter().enumerate() {
            self.handler.add_link(filename.clone(), format!("chapter-{}", i));
        }

        // set tex numbering and toc display to book's parameters
        let numbering = self.book.options.get_i32("numbering").unwrap() - 1;
        content.push_str(&format!("\\setcounter{{tocdepth}}{{{}}}
\\setcounter{{secnumdepth}}{{{}}}\n",
                                  numbering, numbering));

        if self.book.options.get_bool("display_toc").unwrap() {
            content.push_str("\\tableofcontents\n");
        }

        let mut i = 0;
        for &(n, ref v) in &self.book.chapters {
            self.source = Source::new(self.book.filenames[i].as_ref());
            content.push_str(&format!("\\label{{chapter-{}}}", i));
            self.current_chapter = n;
            content.push_str(&try!(self.render_vec(v)));
            i += 1;
        }
        self.source = Source::empty();

        let tex_lang = String::from(match self.book.options.get_str("lang").unwrap() {
            "af" => "afrikaans",
            "sq" => "albanian",
            "eu" => "basque",
            "bg" => "bulgarian",
            "ca" => "catalan",
            "hr" => "croatian",
            "cs" => "czech",
            "da" => "danish",
            "nl" => "dutch",
            "en" => "english",
            "eo" => "esperanto",
            "et" => "estonian",
            "fi" => "finnish",
            "fr" => "francais",
            "gl" => "galician",
            "el" => "greek",
            "de" => "ngerman",
            "he" => "hebrew",
            "hu" => "hungarian",
            "it" => "italian",
            "is" => "icelandic",
            "id" => "indonesian",
            "ga" => "irish",
            "la" => "latin",
            "ms" => "malay",
            "nn" => "norsk",
            "pl" => "polish",
            "pt" => "portuguese",
            "ro" => "romanian",
            "ru" => "russian",
            "gd" => "scottish",
            "sr" => "serbian",
            "sk" => "slovak",
            "sl" => "slovene",
            "es" => "spanish",
            "sw" => "swedish",
            "tr" => "turkish",
            "uk" => "ukrainian",
            "cy" => "welsh",
            _ => {
                self.book.logger.error(format!("LaTeX: can't find a tex equivalent for lang '{}', fallbacking on english", self.book.options.get_str("lang").unwrap()));
                "english"
            }
        });

        let template = mustache::compile_str(try!(self.book.get_template("tex.template")).as_ref());
        let mut data = self.book.get_mapbuilder("tex")
            .insert_str("content", content)
            .insert_str("tex_lang", tex_lang);
        if self.book.options.get_bool("tex.short") == Ok(true) {
            data = data.insert_bool("short", true);
        }
        if self.book.options.get_bool("use_initials") == Ok(true) {
            data = data.insert_bool("initials", true);
        }
        let data = data.build();
        let mut res:Vec<u8> = vec!();
        template.render_data(&mut res, &data);
        match String::from_utf8(res) {
            Err(_) => panic!("generated LaTeX was not valid utf-8"),
            Ok(res) => Ok(res)
        }
    }
}

impl<'a> Renderer for LatexRenderer<'a> {
    fn render_token(&mut self, token: &Token) -> Result<String> {
        match *token {
            Token::Str(ref text) => {
                let content = if self.escape {
                    self.book.clean(escape_tex(text.as_ref()), true)
                } else {
                    Cow::Borrowed(text.as_ref())
                };
                if self.first_letter {
                    self.first_letter = false;
                    if self.book.options.get_bool("use_initials").unwrap() {
                        let mut chars = content.chars().peekable();
                        let initial = try!(chars.next()
                                           .ok_or(Error::parser(&self.book.source,
                                                                    "empty str token, could not find initial")));
                        let mut first_word = String::new();
                        loop  {
                            let c = if let Some(next_char) = chars.peek() {
                                *next_char
                            } else {
                                break;
                            };
                            if !c.is_whitespace() {
                                first_word.push(c);
                                chars.next();
                            } else {
                                break;
                            }
                        }

                        let rest = chars.collect::<String>();

                        if initial.is_alphanumeric() {
                            Ok(format!("\\lettrine{{{}}}{{{}}}{}", initial, first_word, rest))
                        } else {
                            Ok(format!("{}{}{}", initial, first_word, rest))
                        }
                    } else {
                        Ok(content.into_owned())
                    }
                } else {
                    Ok(content.into_owned())
                }
            }
            Token::Paragraph(ref vec) => {
                if self.first_paragraph {
                    self.first_paragraph = false;
                    if !vec.is_empty() && vec[0].is_str() {
                        // Only use initials if first element is a Token::str
                        self.first_letter = true;
                    }
                }
                Ok(format!("{}\n\n",
                           try!(self.render_vec(vec))))
            },
            Token::Header(n, ref vec) => {
                let mut content = String::new();
                if n == 1 {
                    self.first_paragraph = true;
                    if self.current_chapter == Number::Hidden {
                        return Ok(r#"\chapter*{}"#.to_owned());
                    } else {
                        if let Number::Specified(n) = self.current_chapter {
                            content.push_str(r"\setcounter{chapter}{");
                            content.push_str(&format!("{}", n - 1));
                            content.push_str("}\n");
                        }
                    }
                }
                match n {
                    1 => content.push_str(r"\chapter"),
                    2 => content.push_str(r"\section"),
                    3 => content.push_str(r"\subsection"),
                    4 => content.push_str(r"\subsubsection"),
                    _ => content.push_str(r"\paragraph"),
                }
                if self.current_chapter == Number::Unnumbered {
                    content.push_str("*");
                }
                content.push_str(r"{");
                content.push_str(&try!(self.render_vec(vec)));
                content.push_str("}\n");
                Ok(content)
            },
            Token::Emphasis(ref vec) => Ok(format!("\\emph{{{}}}", try!(self.render_vec(vec)))),
            Token::Strong(ref vec) => Ok(format!("\\textbf{{{}}}", try!(self.render_vec(vec)))),
            Token::Code(ref vec) => Ok(format!("\\texttt{{{}}}", try!(self.render_vec(vec)))),
            Token::BlockQuote(ref vec) => Ok(format!("\\begin{{quotation}}\n{}\\end{{quotation}}\n", try!(self.render_vec(vec)))),
            Token::CodeBlock(_, ref vec) => {
                self.escape = false;
                let res = try!(self.render_vec(vec));
                self.escape = true;
                Ok(format!("\\begin{{spverbatim}}{}\\end{{spverbatim}}\n\\vspace{{1em}}\n", res))
            },
            Token::Rule => Ok(String::from("\\HRule\n")),
            Token::SoftBreak => Ok(String::from(" ")),
            Token::HardBreak => Ok(String::from("\n")),
            Token::List(ref vec) => Ok(format!("\\begin{{itemize}}\n{}\\end{{itemize}}", try!(self.render_vec(vec)))),
            Token::OrderedList(_, ref vec) => Ok(format!("\\begin{{enumerate}}\n{}\\end{{enumerate}}\n", try!(self.render_vec(vec)))),
            Token::Item(ref vec) => Ok(format!("\\item {}\n", try!(self.render_vec(vec)))),
            Token::Link(ref url, _, ref vec) => {
                let content = try!(self.render_vec(vec));

                if ResourceHandler::is_local(url) {
                    Ok(format!("\\hyperref[{}]{{{}}}",
                            self.handler.get_link(url),
                            content))
                }
                else {
                    let url = escape_tex(url.as_ref());
                    if &content == &url {
                        Ok(format!("\\url{{{}}}", content))
                    } else {
                        if self.book.options.get_bool("tex.links_as_footnotes").unwrap() {
                            Ok(format!("\\href{{{}}}{{{}}}\\footnote{{\\url{{{}}}}}", url, content, url))
                        } else {
                            Ok(format!("\\href{{{}}}{{{}}}", url, content))
                        }
                    }
                }
            },
            Token::StandaloneImage(ref url, _, _) => {
                if ResourceHandler::is_local(url) {
                    let img = try!(self.handler.map_image(&self.source,
                                                          url.as_ref()));
                    Ok(format!("\\begin{{center}}
  \\includegraphics[width=0.8\\linewidth]{{{}}}
\\end{{center}}",
                            img))

                } else {
                    self.book.logger.error(&format!("LaTeX: image '{}' doesn't seem to be local; ignoring it.", url));
                    Ok(String::new())
                }
            },
            Token::Image(ref url, _, _) => {
                if ResourceHandler::is_local(url) {
                    Ok(format!("\\includegraphics{{{}}}",
                            try!(self.handler.map_image(&self.source,
                                                        url.as_ref()))))
                } else {
                    self.book.logger.error(&format!("LaTeX: image '{}' doesn't seem to be local; ignoring it.", url));
                    Ok(String::new())
                }
            },
            Token::Footnote(ref vec) => Ok(format!("\\footnote{{{}}}", try!(self.render_vec(vec)))),
            Token::Table(n, ref vec) => {
                let mut cols = String::new();
                for _ in 0..n {
                    cols.push_str("|c");
                }
                cols.push_str("|");
                Ok(format!("\\begin{{center}}
\\begin{{tabular}}{{{}}}
\\hline
{}
\\hline
\\end{{tabular}}
\\end{{center}}\n\n", cols, try!(self.render_vec(vec))))
            },
            Token::TableRow(ref vec) | Token::TableHead(ref vec) => {
                let mut res:String = try!(vec.iter()
                    .map(|v| {self.render_token(v)})
                    .collect::<Result<Vec<_>>>())
                    .join(" & ");
                res.push_str("\\\\ \n");
                if let Token::TableHead(_) = *token {
                    res.push_str("\\hline\n");
                }
                Ok(res)
            }
            Token::TableCell(ref vec) => self.render_vec(vec),
        }
    }
}
<commit_msg>Write \section instead of \chapter if tex.short is set to true<commit_after>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published
// by the Free Software Foundation, either version 2.1 of the License, or
// (at your option) any later version.
//
// Caribon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

use book::Book;
use number::Number;
use error::{Error,Result, Source};
use token::Token;
use zipper::Zipper;
use escape::escape_tex;
use resource_handler::ResourceHandler;
use renderer::Renderer;

use std::iter::Iterator;
use std::fs::File;
use std::io::Read;
use std::borrow::Cow;

use mustache;

/// LaTeX renderer
pub struct LatexRenderer<'a> {
    book: &'a Book,
    current_chapter: Number,
    handler: ResourceHandler<'a>,
    source: Source,
    escape: bool,
    first_letter: bool,
    first_paragraph: bool,
    is_short: bool,
}

impl<'a> LatexRenderer<'a> {
    /// Creates new LatexRenderer
    pub fn new(book: &'a Book) -> LatexRenderer<'a> {
        let mut handler = ResourceHandler::new(&book.logger);
        handler.set_images_mapping(true);
        LatexRenderer {
            book: book,
            current_chapter: Number::Default,
            handler: handler,
            source: Source::empty(),
            escape: true,
            first_letter: false,
            first_paragraph: true,
            is_short: book.options.get_bool("tex.short").unwrap(),
        }
    }

    /// Render pdf in a file
    pub fn render_pdf(&mut self) -> Result<String> {
        if let Ok(pdf_file) = self.book.options.get_path("output.pdf") {
            let content = try!(self.render_book());
            let mut zipper = try!(Zipper::new(&self.book.options.get_path("temp_dir").unwrap()));
            try!(zipper.write("result.tex", &content.as_bytes(), false));

            // write image files
            for (source, dest) in self.handler.images_mapping() {
                let mut f = try!(File::open(source).map_err(|_| Error::file_not_found(&self.source,
                                                                                "image",
                                                                                source.to_owned())));
                let mut content = vec!();
                try!(f.read_to_end(&mut content).map_err(|e| Error::render(&self.source,
                                                                               format!("error while reading image file: {}", e))));
                try!(zipper.write(dest, &content, true));
            }

            zipper.generate_pdf(&self.book.options.get_str("tex.command").unwrap(), "result.tex", &pdf_file)
        } else {
            Err(Error::render(&self.source,
                                  "no output pdf file specified in book config"))
        }
    }

    /// Render latex in a string
    pub fn render_book(&mut self) -> Result<String> {
        let mut content = String::from("");
        for (i, filename) in self.book.filenames.iter().enumerate() {
            self.handler.add_link(filename.clone(), format!("chapter-{}", i));
        }

        // set tex numbering and toc display to book's parameters
        let numbering = self.book.options.get_i32("numbering").unwrap() - 1;
        content.push_str(&format!("\\setcounter{{tocdepth}}{{{}}}
\\setcounter{{secnumdepth}}{{{}}}\n",
                                  numbering, numbering));

        if self.book.options.get_bool("display_toc").unwrap() {
            content.push_str("\\tableofcontents\n");
        }

        let mut i = 0;
        for &(n, ref v) in &self.book.chapters {
            self.source = Source::new(self.book.filenames[i].as_ref());
            content.push_str(&format!("\\label{{chapter-{}}}", i));
            self.current_chapter = n;
            content.push_str(&try!(self.render_vec(v)));
            i += 1;
        }
        self.source = Source::empty();

        let tex_lang = String::from(match self.book.options.get_str("lang").unwrap() {
            "af" => "afrikaans",
            "sq" => "albanian",
            "eu" => "basque",
            "bg" => "bulgarian",
            "ca" => "catalan",
            "hr" => "croatian",
            "cs" => "czech",
            "da" => "danish",
            "nl" => "dutch",
            "en" => "english",
            "eo" => "esperanto",
            "et" => "estonian",
            "fi" => "finnish",
            "fr" => "francais",
            "gl" => "galician",
            "el" => "greek",
            "de" => "ngerman",
            "he" => "hebrew",
            "hu" => "hungarian",
            "it" => "italian",
            "is" => "icelandic",
            "id" => "indonesian",
            "ga" => "irish",
            "la" => "latin",
            "ms" => "malay",
            "nn" => "norsk",
            "pl" => "polish",
            "pt" => "portuguese",
            "ro" => "romanian",
            "ru" => "russian",
            "gd" => "scottish",
            "sr" => "serbian",
            "sk" => "slovak",
            "sl" => "slovene",
            "es" => "spanish",
            "sw" => "swedish",
            "tr" => "turkish",
            "uk" => "ukrainian",
            "cy" => "welsh",
            _ => {
                self.book.logger.error(format!("LaTeX: can't find a tex equivalent for lang '{}', fallbacking on english", self.book.options.get_str("lang").unwrap()));
                "english"
            }
        });

        let template = mustache::compile_str(try!(self.book.get_template("tex.template")).as_ref());
        let mut data = self.book.get_mapbuilder("tex")
            .insert_str("content", content)
            .insert_str("tex_lang", tex_lang);
        if self.book.options.get_bool("tex.short") == Ok(true) {
            data = data.insert_bool("short", true);
        }
        if self.book.options.get_bool("use_initials") == Ok(true) {
            data = data.insert_bool("initials", true);
        }
        let data = data.build();
        let mut res:Vec<u8> = vec!();
        template.render_data(&mut res, &data);
        match String::from_utf8(res) {
            Err(_) => panic!("generated LaTeX was not valid utf-8"),
            Ok(res) => Ok(res)
        }
    }
}

impl<'a> Renderer for LatexRenderer<'a> {
    fn render_token(&mut self, token: &Token) -> Result<String> {
        match *token {
            Token::Str(ref text) => {
                let content = if self.escape {
                    self.book.clean(escape_tex(text.as_ref()), true)
                } else {
                    Cow::Borrowed(text.as_ref())
                };
                if self.first_letter {
                    self.first_letter = false;
                    if self.book.options.get_bool("use_initials").unwrap() {
                        let mut chars = content.chars().peekable();
                        let initial = try!(chars.next()
                                           .ok_or(Error::parser(&self.book.source,
                                                                    "empty str token, could not find initial")));
                        let mut first_word = String::new();
                        loop  {
                            let c = if let Some(next_char) = chars.peek() {
                                *next_char
                            } else {
                                break;
                            };
                            if !c.is_whitespace() {
                                first_word.push(c);
                                chars.next();
                            } else {
                                break;
                            }
                        }

                        let rest = chars.collect::<String>();

                        if initial.is_alphanumeric() {
                            Ok(format!("\\lettrine{{{}}}{{{}}}{}", initial, first_word, rest))
                        } else {
                            Ok(format!("{}{}{}", initial, first_word, rest))
                        }
                    } else {
                        Ok(content.into_owned())
                    }
                } else {
                    Ok(content.into_owned())
                }
            }
            Token::Paragraph(ref vec) => {
                if self.first_paragraph {
                    self.first_paragraph = false;
                    if !vec.is_empty() && vec[0].is_str() {
                        // Only use initials if first element is a Token::str
                        self.first_letter = true;
                    }
                }
                Ok(format!("{}\n\n",
                           try!(self.render_vec(vec))))
            },
            Token::Header(n, ref vec) => {
                let mut content = String::new();
                if n == 1 {
                    self.first_paragraph = true;
                    if self.current_chapter == Number::Hidden {
                        if !self.is_short {
                            return Ok(r#"\chapter*{}"#.to_owned());
                        } else {
                            return Ok(r#"\section*{}"#.to_owned());
                        }
                    } else {
                        if let Number::Specified(n) = self.current_chapter {
                            content.push_str(r"\setcounter{chapter}{");
                            content.push_str(&format!("{}", n - 1));
                            content.push_str("}\n");
                        }
                    }
                }
                match n {
                    1 => {
                        if !self.is_short {
                            content.push_str(r"\chapter");
                        } else {
                            // Chapters aren't handlled for class article
                            content.push_str(r"\section");
                        }
                    },
                    2 => content.push_str(r"\section"),
                    3 => content.push_str(r"\subsection"),
                    4 => content.push_str(r"\subsubsection"),
                    _ => content.push_str(r"\paragraph"),
                }
                if self.current_chapter == Number::Unnumbered {
                    content.push_str("*");
                }
                content.push_str(r"{");
                content.push_str(&try!(self.render_vec(vec)));
                content.push_str("}\n");
                Ok(content)
            },
            Token::Emphasis(ref vec) => Ok(format!("\\emph{{{}}}", try!(self.render_vec(vec)))),
            Token::Strong(ref vec) => Ok(format!("\\textbf{{{}}}", try!(self.render_vec(vec)))),
            Token::Code(ref vec) => Ok(format!("\\texttt{{{}}}", try!(self.render_vec(vec)))),
            Token::BlockQuote(ref vec) => Ok(format!("\\begin{{quotation}}\n{}\\end{{quotation}}\n", try!(self.render_vec(vec)))),
            Token::CodeBlock(_, ref vec) => {
                self.escape = false;
                let res = try!(self.render_vec(vec));
                self.escape = true;
                Ok(format!("\\begin{{spverbatim}}{}\\end{{spverbatim}}\n\\vspace{{1em}}\n", res))
            },
            Token::Rule => Ok(String::from("\\HRule\n")),
            Token::SoftBreak => Ok(String::from(" ")),
            Token::HardBreak => Ok(String::from("\n")),
            Token::List(ref vec) => Ok(format!("\\begin{{itemize}}\n{}\\end{{itemize}}", try!(self.render_vec(vec)))),
            Token::OrderedList(_, ref vec) => Ok(format!("\\begin{{enumerate}}\n{}\\end{{enumerate}}\n", try!(self.render_vec(vec)))),
            Token::Item(ref vec) => Ok(format!("\\item {}\n", try!(self.render_vec(vec)))),
            Token::Link(ref url, _, ref vec) => {
                let content = try!(self.render_vec(vec));

                if ResourceHandler::is_local(url) {
                    Ok(format!("\\hyperref[{}]{{{}}}",
                            self.handler.get_link(url),
                            content))
                }
                else {
                    let url = escape_tex(url.as_ref());
                    if &content == &url {
                        Ok(format!("\\url{{{}}}", content))
                    } else {
                        if self.book.options.get_bool("tex.links_as_footnotes").unwrap() {
                            Ok(format!("\\href{{{}}}{{{}}}\\footnote{{\\url{{{}}}}}", url, content, url))
                        } else {
                            Ok(format!("\\href{{{}}}{{{}}}", url, content))
                        }
                    }
                }
            },
            Token::StandaloneImage(ref url, _, _) => {
                if ResourceHandler::is_local(url) {
                    let img = try!(self.handler.map_image(&self.source,
                                                          url.as_ref()));
                    Ok(format!("\\begin{{center}}
  \\includegraphics[width=0.8\\linewidth]{{{}}}
\\end{{center}}",
                            img))

                } else {
                    self.book.logger.error(&format!("LaTeX: image '{}' doesn't seem to be local; ignoring it.", url));
                    Ok(String::new())
                }
            },
            Token::Image(ref url, _, _) => {
                if ResourceHandler::is_local(url) {
                    Ok(format!("\\includegraphics{{{}}}",
                            try!(self.handler.map_image(&self.source,
                                                        url.as_ref()))))
                } else {
                    self.book.logger.error(&format!("LaTeX: image '{}' doesn't seem to be local; ignoring it.", url));
                    Ok(String::new())
                }
            },
            Token::Footnote(ref vec) => Ok(format!("\\footnote{{{}}}", try!(self.render_vec(vec)))),
            Token::Table(n, ref vec) => {
                let mut cols = String::new();
                for _ in 0..n {
                    cols.push_str("|c");
                }
                cols.push_str("|");
                Ok(format!("\\begin{{center}}
\\begin{{tabular}}{{{}}}
\\hline
{}
\\hline
\\end{{tabular}}
\\end{{center}}\n\n", cols, try!(self.render_vec(vec))))
            },
            Token::TableRow(ref vec) | Token::TableHead(ref vec) => {
                let mut res:String = try!(vec.iter()
                    .map(|v| {self.render_token(v)})
                    .collect::<Result<Vec<_>>>())
                    .join(" & ");
                res.push_str("\\\\ \n");
                if let Token::TableHead(_) = *token {
                    res.push_str("\\hline\n");
                }
                Ok(res)
            }
            Token::TableCell(ref vec) => self.render_vec(vec),
        }
    }
}
<|endoftext|>


<commit_before>use std::mem;
use utils::{GenerationalId, IdAllocator};
use shape::FeatureId;
use bounding_volume::PolyhedralCone;
use query::{Contact, ContactKinematic, TrackedContact};
use math::{Isometry, Point};

use na::{self, Real, Unit};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum CacheEntryStatus {
    Used(usize),
    Unused(usize),
}

impl CacheEntryStatus {
    fn is_used(&self) -> bool {
        match *self {
            CacheEntryStatus::Used(..) => true,
            _ => false,
        }
    }

    fn is_obsolete(&self) -> bool {
        match *self {
            CacheEntryStatus::Unused(0) => true,
            _ => false,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ContactManifold<P: Point> {
    deepest: usize,
    contacts: Vec<TrackedContact<P>>,
    // FIXME: the cache should only contain points
    // (the contact in local-space of body 1).
    cache: Vec<TrackedContact<P>>,
    cached_contact_used: Vec<CacheEntryStatus>,
    new_cached_contact_used: Vec<CacheEntryStatus>,
    max_life: usize,
}

impl<P: Point> ContactManifold<P> {
    pub fn new() -> Self {
        ContactManifold {
            deepest: 0,
            contacts: Vec::new(),
            cache: Vec::new(),
            cached_contact_used: Vec::new(),
            new_cached_contact_used: Vec::new(), // FIXME: the existence of this buffer is ugly.
            max_life: 0,                         // FIXME: don't hard-code this.
        }
    }

    pub fn len(&self) -> usize {
        self.contacts.len()
    }

    pub fn contacts(&self) -> &[TrackedContact<P>] {
        &self.contacts[..]
    }

    pub fn deepest_contact_id(&self) -> usize {
        self.deepest
    }

    pub fn deepest_contact(&self) -> Option<&TrackedContact<P>> {
        if self.len() != 0 {
            Some(&self.contacts[self.deepest])
        } else {
            None
        }
    }

    pub fn save_cache_and_clear(&mut self, gen: &mut IdAllocator) {
        for (valid, c) in self.cached_contact_used.iter_mut().zip(self.cache.iter()) {
            if valid.is_obsolete() {
                gen.free(c.id)
            }
        }

        mem::swap(&mut self.contacts, &mut self.cache);

        self.new_cached_contact_used.clear();
        self.new_cached_contact_used
            .resize(self.cache.len(), CacheEntryStatus::Unused(self.max_life));

        for (valid, c) in self.cached_contact_used
            .drain(..)
            .zip(self.contacts.drain(..))
        {
            match valid {
                CacheEntryStatus::Unused(life) if life > 0 => {
                    self.cache.push(c);
                    self.new_cached_contact_used
                        .push(CacheEntryStatus::Unused(life - 1));
                }
                _ => {}
            }
        }

        mem::swap(
            &mut self.new_cached_contact_used,
            &mut self.cached_contact_used,
        );

        self.contacts.clear();
        self.deepest = 0;
    }

    pub fn push(
        &mut self,
        contact: Contact<P>,
        kinematic: ContactKinematic<P>,
        gen: &mut IdAllocator,
    ) -> bool {
        // FIXME: all this is poorly designed and quite inefficient (but OK for a first
        // non-optimized implementation).
        let mut closest = GenerationalId::invalid();
        let mut closest_i = 0;
        let mut closest_is_contact = false;
        let mut closest_dist: P::Real = na::convert(0.02 * 0.02); // FIXME: don't hard-code this.

        for i in 0..self.cache.len() {
            let dist = na::distance_squared(&kinematic.local1(), &self.cache[i].kinematic.local1());
            if dist < closest_dist {
                closest_dist = dist;
                closest = self.cache[i].id;
                closest_i = i;
            }
        }

        for i in 0..self.contacts.len() {
            let dist =
                na::distance_squared(&kinematic.local1(), &self.contacts[i].kinematic.local1());
            if dist < closest_dist {
                closest_is_contact = true;
                closest_dist = dist;
                closest = self.contacts[i].id;
                closest_i = i;
            }
        }

        let is_deepest =
            self.contacts.len() == 0 || contact.depth > self.contacts[self.deepest].contact.depth;

        let mut matched = false;
        if closest.is_invalid() {
            closest = gen.alloc();
        } else {
            if closest_is_contact {
                self.contacts[closest_i] = TrackedContact::new(contact, kinematic, closest);
                if is_deepest {
                    self.deepest = closest_i;
                }
                return false;
            }
            if let CacheEntryStatus::Used(used_i) = self.cached_contact_used[closest_i] {
                self.contacts[used_i] = TrackedContact::new(contact, kinematic, closest);

                if is_deepest {
                    self.deepest = used_i;
                }
                return false;
            } else {
                self.cached_contact_used[closest_i] = CacheEntryStatus::Used(self.contacts.len());
                matched = true;
            }
        }

        if is_deepest {
            self.deepest = self.contacts.len();
        }
        let tracked = TrackedContact::new(contact, kinematic, closest);
        self.contacts.push(tracked);
        return matched;
    }

    /*
    pub fn push(
        &mut self,
        contact: Contact<P>,
        kinematic: ContactKinematic<P>,
        gen: &mut IdAllocator,
    ) -> bool {
        let mut id = GenerationalId::invalid();
        let mut matched = false;

        for i in 0..self.cache.len() {
            if self.cached_contact_used[i].is_obsolete()
                && self.cache[i].kinematic.feature1() == kinematic.feature1()
                && self.cache[i].kinematic.feature2() == kinematic.feature2()
            {
                self.cached_contact_used[i] = CacheEntryStatus::Used(self.contacts.len());
                id = self.cache[i].id;
                matched = true;
            }
        }

        if id.is_invalid() {
            id = gen.alloc();
        }

        let tracked = TrackedContact::new(contact, kinematic, id);
        self.contacts.push(tracked);
        matched
    }
    */
}
<commit_msg>Contact manifold: don't expand the cache with contacts order han one frame.<commit_after>use std::mem;
use utils::{GenerationalId, IdAllocator};
use shape::FeatureId;
use bounding_volume::PolyhedralCone;
use query::{Contact, ContactKinematic, TrackedContact};
use math::{Isometry, Point};

use na::{self, Real, Unit};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum CacheEntryStatus {
    Used(usize),
    Unused,
}

impl CacheEntryStatus {
    fn is_used(&self) -> bool {
        match *self {
            CacheEntryStatus::Used(..) => true,
            _ => false,
        }
    }

    fn is_obsolete(&self) -> bool {
        match *self {
            CacheEntryStatus::Unused => true,
            _ => false,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ContactManifold<P: Point> {
    deepest: usize,
    contacts: Vec<TrackedContact<P>>,
    // FIXME: the cache should only contain points
    // (the contact in local-space of body 1).
    cache: Vec<TrackedContact<P>>,
    cached_contact_used: Vec<CacheEntryStatus>,
    new_cached_contact_used: Vec<CacheEntryStatus>,
}

impl<P: Point> ContactManifold<P> {
    pub fn new() -> Self {
        ContactManifold {
            deepest: 0,
            contacts: Vec::new(),
            cache: Vec::new(),
            cached_contact_used: Vec::new(),
            new_cached_contact_used: Vec::new(), // FIXME: the existence of this buffer is ugly.
        }
    }

    pub fn len(&self) -> usize {
        self.contacts.len()
    }

    pub fn contacts(&self) -> &[TrackedContact<P>] {
        &self.contacts[..]
    }

    pub fn deepest_contact_id(&self) -> usize {
        self.deepest
    }

    pub fn deepest_contact(&self) -> Option<&TrackedContact<P>> {
        if self.len() != 0 {
            Some(&self.contacts[self.deepest])
        } else {
            None
        }
    }

    pub fn save_cache_and_clear(&mut self, gen: &mut IdAllocator) {
        for (valid, c) in self.cached_contact_used.iter_mut().zip(self.cache.iter()) {
            if valid.is_obsolete() {
                gen.free(c.id)
            }
        }

        mem::swap(&mut self.contacts, &mut self.cache);

        self.new_cached_contact_used.clear();
        self.new_cached_contact_used
            .resize(self.cache.len(), CacheEntryStatus::Unused);

        mem::swap(
            &mut self.new_cached_contact_used,
            &mut self.cached_contact_used,
        );

        self.contacts.clear();
        self.deepest = 0;
    }

    pub fn push(
        &mut self,
        contact: Contact<P>,
        kinematic: ContactKinematic<P>,
        gen: &mut IdAllocator,
    ) -> bool {
        // FIXME: all this is poorly designed and quite inefficient
        // (but OK for a first non-optimized implementation).
        let mut closest = GenerationalId::invalid();
        let mut closest_i = 0;
        let mut closest_is_contact = false;
        let mut closest_dist: P::Real = na::convert(0.02 * 0.02); // FIXME: don't hard-code this.

        for i in 0..self.cache.len() {
            let dist = na::distance_squared(&kinematic.local1(), &self.cache[i].kinematic.local1());
            if dist < closest_dist {
                closest_dist = dist;
                closest = self.cache[i].id;
                closest_i = i;
            }
        }

        for i in 0..self.contacts.len() {
            let dist =
                na::distance_squared(&kinematic.local1(), &self.contacts[i].kinematic.local1());
            if dist < closest_dist {
                closest_is_contact = true;
                closest_dist = dist;
                closest = self.contacts[i].id;
                closest_i = i;
            }
        }

        let is_deepest =
            self.contacts.len() == 0 || contact.depth > self.contacts[self.deepest].contact.depth;

        let mut matched = false;
        if closest.is_invalid() {
            closest = gen.alloc();
        } else {
            if closest_is_contact {
                self.contacts[closest_i] = TrackedContact::new(contact, kinematic, closest);
                if is_deepest {
                    self.deepest = closest_i;
                }
                return false;
            }
            if let CacheEntryStatus::Used(used_i) = self.cached_contact_used[closest_i] {
                self.contacts[used_i] = TrackedContact::new(contact, kinematic, closest);

                if is_deepest {
                    self.deepest = used_i;
                }
                return false;
            } else {
                self.cached_contact_used[closest_i] = CacheEntryStatus::Used(self.contacts.len());
                matched = true;
            }
        }

        if is_deepest {
            self.deepest = self.contacts.len();
        }
        let tracked = TrackedContact::new(contact, kinematic, closest);
        self.contacts.push(tracked);
        return matched;
    }

    /*
    pub fn push(
        &mut self,
        contact: Contact<P>,
        kinematic: ContactKinematic<P>,
        gen: &mut IdAllocator,
    ) -> bool {
        let mut id = GenerationalId::invalid();
        let mut matched = false;

        for i in 0..self.cache.len() {
            if self.cached_contact_used[i].is_obsolete()
                && self.cache[i].kinematic.feature1() == kinematic.feature1()
                && self.cache[i].kinematic.feature2() == kinematic.feature2()
            {
                self.cached_contact_used[i] = CacheEntryStatus::Used(self.contacts.len());
                id = self.cache[i].id;
                matched = true;
            }
        }

        if id.is_invalid() {
            id = gen.alloc();
        }

        let tracked = TrackedContact::new(contact, kinematic, id);
        self.contacts.push(tracked);
        matched
    }
    */
}
<|endoftext|>


<commit_before>use prelude::*;
use expression::SelectableExpression;
use expression::nullable::Nullable;
use query_builder::*;
use result::QueryResult;
use super::{QuerySource, Table};

#[derive(Debug, Clone, Copy)]
/// A query source representing the join between two tables
pub struct Join<Left, Right, Kind> {
    left: Left,
    right: Right,
    kind: Kind,
}

impl<Left, Right, Kind> Join<Left, Right, Kind> {
    pub fn new(left: Left, right: Right) -> Self where
        Kind: Default,
    {
        Join {
            left: left,
            right: right,
            kind: Kind::default(),
        }
    }
}

impl<Left, Right, Kind> AsQuery for Join<Left, Right, Kind> where
    SelectStatement<Join<Left, Right, Kind>>: Query,
{
    type SqlType = <SelectStatement<Self> as Query>::SqlType;
    type Query = SelectStatement<Self>;

    fn as_query(self) -> Self::Query {
        SelectStatement::simple(self)
    }
}

impl_query_id!(Join<Left, Right, Kind>);

impl<Left, Right> QuerySource for Join<Left, Right, Inner> where
    Left: Table + JoinTo<Right, Inner>,
    Right: Table,
    (Left::AllColumns, Right::AllColumns): SelectableExpression<
        Join<Left, Right, Inner>,
    >,
{
    type FromClause = <Left as JoinTo<Right, Inner>>::JoinClause;
    type DefaultSelection = (Left::AllColumns, Right::AllColumns);

    fn from_clause(&self) -> Self::FromClause {
        self.left.join_clause(self.kind)
    }

    fn default_selection(&self) -> Self::DefaultSelection {
        (Left::all_columns(), Right::all_columns())
    }
}

impl<Left, Right> QuerySource for Join<Left, Right, LeftOuter> where
    Left: Table + JoinTo<Right, LeftOuter>,
    Right: Table,
    (Left::AllColumns, Nullable<Right::AllColumns>): SelectableExpression<
        Join<Left, Right, LeftOuter>,
    >,
{
    type FromClause = <Left as JoinTo<Right, LeftOuter>>::JoinClause;
    type DefaultSelection = (Left::AllColumns, Nullable<Right::AllColumns>);

    fn from_clause(&self) -> Self::FromClause {
        self.left.join_clause(self.kind)
    }

    fn default_selection(&self) -> Self::DefaultSelection {
        (Left::all_columns(), Right::all_columns().nullable())
    }
}

impl<Left, Right, T> SelectableExpression<Join<Left, Right, LeftOuter>>
    for Nullable<T> where
        T: SelectableExpression<Join<Left, Right, Inner>>,
        Nullable<T>: AppearsOnTable<Join<Left, Right, LeftOuter>>,
{
}

/// Indicates that two tables can be used together in a JOIN clause.
/// Implementations of this trait will be generated for you automatically by
/// the [association annotations](FIXME: Add link) from codegen.
pub trait JoinTo<T: Table, JoinType>: Table {
    #[doc(hidden)]
    type JoinClause;
    #[doc(hidden)]
    fn join_clause(&self, join_type: JoinType) -> Self::JoinClause;
}

use backend::Backend;

#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct Inner;
impl_query_id!(Inner);

impl<DB: Backend> QueryFragment<DB> for Inner {
    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
        out.push_sql(" INNER");
        Ok(())
    }

    fn collect_binds(&self, _out: &mut DB::BindCollector) -> QueryResult<()> {
        Ok(())
    }

    fn is_safe_to_cache_prepared(&self) -> bool {
        true
    }
}

#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct LeftOuter;
impl_query_id!(LeftOuter);

impl<DB: Backend> QueryFragment<DB> for LeftOuter {
    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
        out.push_sql(" LEFT OUTER");
        Ok(())
    }

    fn collect_binds(&self, _out: &mut DB::BindCollector) -> QueryResult<()> {
        Ok(())
    }

    fn is_safe_to_cache_prepared(&self) -> bool {
        true
    }
}
<commit_msg>Add missing link to JoinTo rustdoc<commit_after>use prelude::*;
use expression::SelectableExpression;
use expression::nullable::Nullable;
use query_builder::*;
use result::QueryResult;
use super::{QuerySource, Table};

#[derive(Debug, Clone, Copy)]
/// A query source representing the join between two tables
pub struct Join<Left, Right, Kind> {
    left: Left,
    right: Right,
    kind: Kind,
}

impl<Left, Right, Kind> Join<Left, Right, Kind> {
    pub fn new(left: Left, right: Right) -> Self where
        Kind: Default,
    {
        Join {
            left: left,
            right: right,
            kind: Kind::default(),
        }
    }
}

impl<Left, Right, Kind> AsQuery for Join<Left, Right, Kind> where
    SelectStatement<Join<Left, Right, Kind>>: Query,
{
    type SqlType = <SelectStatement<Self> as Query>::SqlType;
    type Query = SelectStatement<Self>;

    fn as_query(self) -> Self::Query {
        SelectStatement::simple(self)
    }
}

impl_query_id!(Join<Left, Right, Kind>);

impl<Left, Right> QuerySource for Join<Left, Right, Inner> where
    Left: Table + JoinTo<Right, Inner>,
    Right: Table,
    (Left::AllColumns, Right::AllColumns): SelectableExpression<
        Join<Left, Right, Inner>,
    >,
{
    type FromClause = <Left as JoinTo<Right, Inner>>::JoinClause;
    type DefaultSelection = (Left::AllColumns, Right::AllColumns);

    fn from_clause(&self) -> Self::FromClause {
        self.left.join_clause(self.kind)
    }

    fn default_selection(&self) -> Self::DefaultSelection {
        (Left::all_columns(), Right::all_columns())
    }
}

impl<Left, Right> QuerySource for Join<Left, Right, LeftOuter> where
    Left: Table + JoinTo<Right, LeftOuter>,
    Right: Table,
    (Left::AllColumns, Nullable<Right::AllColumns>): SelectableExpression<
        Join<Left, Right, LeftOuter>,
    >,
{
    type FromClause = <Left as JoinTo<Right, LeftOuter>>::JoinClause;
    type DefaultSelection = (Left::AllColumns, Nullable<Right::AllColumns>);

    fn from_clause(&self) -> Self::FromClause {
        self.left.join_clause(self.kind)
    }

    fn default_selection(&self) -> Self::DefaultSelection {
        (Left::all_columns(), Right::all_columns().nullable())
    }
}

impl<Left, Right, T> SelectableExpression<Join<Left, Right, LeftOuter>>
    for Nullable<T> where
        T: SelectableExpression<Join<Left, Right, Inner>>,
        Nullable<T>: AppearsOnTable<Join<Left, Right, LeftOuter>>,
{
}

/// Indicates that two tables can be used together in a JOIN clause.
/// Implementations of this trait will be generated for you automatically by
/// the [association annotations](../associations/index.html) from codegen.
pub trait JoinTo<T: Table, JoinType>: Table {
    #[doc(hidden)]
    type JoinClause;
    #[doc(hidden)]
    fn join_clause(&self, join_type: JoinType) -> Self::JoinClause;
}

use backend::Backend;

#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct Inner;
impl_query_id!(Inner);

impl<DB: Backend> QueryFragment<DB> for Inner {
    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
        out.push_sql(" INNER");
        Ok(())
    }

    fn collect_binds(&self, _out: &mut DB::BindCollector) -> QueryResult<()> {
        Ok(())
    }

    fn is_safe_to_cache_prepared(&self) -> bool {
        true
    }
}

#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct LeftOuter;
impl_query_id!(LeftOuter);

impl<DB: Backend> QueryFragment<DB> for LeftOuter {
    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
        out.push_sql(" LEFT OUTER");
        Ok(())
    }

    fn collect_binds(&self, _out: &mut DB::BindCollector) -> QueryResult<()> {
        Ok(())
    }

    fn is_safe_to_cache_prepared(&self) -> bool {
        true
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

use support::{database, project};

#[cfg(not(windows))]
const NEW_LINE: &'static str = "\n";

#[cfg(windows)]
const NEW_LINE: &'static str = "\n";

#[test]
fn run_infer_schema_without_docs() {
    test_print_schema("print_schema_simple_without_docs", vec![]);
}

#[test]
fn run_infer_schema() {
    test_print_schema("print_schema_simple", vec!["--with-docs"]);
}

#[test]
fn run_infer_schema_include() {
    test_print_schema(
        "print_schema_only_tables",
        vec!["--with-docs", "-w", "users1"],
    );
}

#[test]
fn run_infer_schema_exclude() {
    test_print_schema(
        "print_schema_except_tables",
        vec!["--with-docs", "-b", "users1"],
    );
}

#[test]
fn run_infer_schema_order() {
    test_print_schema("print_schema_order", vec!["--with-docs"]);
}

#[test]
fn run_infer_schema_compound_primary_key() {
    test_print_schema("print_schema_compound_primary_key", vec!["--with-docs"]);
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_specifying_schema_name() {
    test_print_schema(
        "print_schema_specifying_schema_name",
        vec!["--with-docs", "--schema", "custom_schema"],
    )
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_specifying_schema_name_with_foreign_keys() {
    test_print_schema(
        "print_schema_specifying_schema_name_with_foreign_keys",
        vec!["--with-docs", "--schema", "custom_schema"],
    )
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_with_foreign_keys() {
    test_print_schema("print_schema_with_foreign_keys", vec!["--with-docs"]);
}

#[test]
fn print_schema_column_renaming() {
    test_print_schema("print_schema_column_renaming", vec!["--with-docs"]);
}

#[test]
#[cfg(feature = "mysql")]
fn print_schema_unsigned() {
    test_print_schema("print_schema_unsigned", vec!["--with-docs"]);
}

#[test]
fn print_schema_patch_file() {
    let path_to_patch_file = backend_file_path("print_schema_patch_file", "schema.patch");
    let path = path_to_patch_file.display().to_string();
    test_print_schema("print_schema_patch_file", vec!["--patch-file", &path]);
}

#[test]
fn print_schema_custom_types() {
    test_print_schema(
        "print_schema_custom_types",
        vec!["--import-types", "foo::*", "--import-types", "bar::*"],
    );
}

#[cfg(feature = "sqlite")]
const BACKEND: &str = "sqlite";
#[cfg(feature = "postgres")]
const BACKEND: &str = "postgres";
#[cfg(feature = "mysql")]
const BACKEND: &str = "mysql";

fn backend_file_path(test_name: &str, file: &str) -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("print_schema")
        .join(test_name)
        .join(BACKEND)
        .join(file)
}

fn test_print_schema(test_name: &str, args: Vec<&str>) {
    let test_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("print_schema")
        .join(test_name);
    let p = project(test_name).build();
    let db = database(&p.database_url());

    p.command("setup").run();

    let schema = read_file(&backend_file_path(test_name, "schema.sql"));
    db.execute(&schema);

    let result = p.command("print-schema").args(args).run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    let expected = read_file(&backend_file_path(test_name, "expected.rs"));

    assert_diff!(&expected, result.stdout(), NEW_LINE, 0);

    test_print_schema_config(test_name, &test_path, schema, expected);
}

fn test_print_schema_config(test_name: &str, test_path: &Path, schema: String, expected: String) {
    let config = read_file(&test_path.join("diesel.toml"));
    let mut p = project(&format!("{}_config", test_name)).file("diesel.toml", &config);

    let patch_file = backend_file_path(test_name, "schema.patch");
    if patch_file.exists() {
        let patch_contents = read_file(&patch_file);
        p = p.file("schema.patch", &patch_contents);
    }

    let p = p.build();

    p.command("setup").run();
    p.create_migration("12345_create_schema", &schema, "");

    let result = p.command("migration").arg("run").run();
    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    let schema = p.file_contents("src/schema.rs");
    assert_diff!(&expected, &schema, NEW_LINE, 0);

    let result = p.command("print-schema").run();
    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    assert_diff!(&expected, result.stdout(), NEW_LINE, 0);
}

fn read_file(path: &Path) -> String {
    let mut file = File::open(path).expect(&format!("Could not open {}", path.display()));
    let mut string = String::new();
    file.read_to_string(&mut string).unwrap();
    string
}
<commit_msg>Really fix line endings for windows tests<commit_after>use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};

use support::{database, project};

#[cfg(not(windows))]
const NEW_LINE: &'static str = "\n";

#[cfg(windows)]
const NEW_LINE: &'static str = "\r\n";

#[test]
fn run_infer_schema_without_docs() {
    test_print_schema("print_schema_simple_without_docs", vec![]);
}

#[test]
fn run_infer_schema() {
    test_print_schema("print_schema_simple", vec!["--with-docs"]);
}

#[test]
fn run_infer_schema_include() {
    test_print_schema(
        "print_schema_only_tables",
        vec!["--with-docs", "-w", "users1"],
    );
}

#[test]
fn run_infer_schema_exclude() {
    test_print_schema(
        "print_schema_except_tables",
        vec!["--with-docs", "-b", "users1"],
    );
}

#[test]
fn run_infer_schema_order() {
    test_print_schema("print_schema_order", vec!["--with-docs"]);
}

#[test]
fn run_infer_schema_compound_primary_key() {
    test_print_schema("print_schema_compound_primary_key", vec!["--with-docs"]);
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_specifying_schema_name() {
    test_print_schema(
        "print_schema_specifying_schema_name",
        vec!["--with-docs", "--schema", "custom_schema"],
    )
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_specifying_schema_name_with_foreign_keys() {
    test_print_schema(
        "print_schema_specifying_schema_name_with_foreign_keys",
        vec!["--with-docs", "--schema", "custom_schema"],
    )
}

#[test]
#[cfg(feature = "postgres")]
fn print_schema_with_foreign_keys() {
    test_print_schema("print_schema_with_foreign_keys", vec!["--with-docs"]);
}

#[test]
fn print_schema_column_renaming() {
    test_print_schema("print_schema_column_renaming", vec!["--with-docs"]);
}

#[test]
#[cfg(feature = "mysql")]
fn print_schema_unsigned() {
    test_print_schema("print_schema_unsigned", vec!["--with-docs"]);
}

#[test]
fn print_schema_patch_file() {
    let path_to_patch_file = backend_file_path("print_schema_patch_file", "schema.patch");
    let path = path_to_patch_file.display().to_string();
    test_print_schema("print_schema_patch_file", vec!["--patch-file", &path]);
}

#[test]
fn print_schema_custom_types() {
    test_print_schema(
        "print_schema_custom_types",
        vec!["--import-types", "foo::*", "--import-types", "bar::*"],
    );
}

#[cfg(feature = "sqlite")]
const BACKEND: &str = "sqlite";
#[cfg(feature = "postgres")]
const BACKEND: &str = "postgres";
#[cfg(feature = "mysql")]
const BACKEND: &str = "mysql";

fn backend_file_path(test_name: &str, file: &str) -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("print_schema")
        .join(test_name)
        .join(BACKEND)
        .join(file)
}

fn test_print_schema(test_name: &str, args: Vec<&str>) {
    let test_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("print_schema")
        .join(test_name);
    let p = project(test_name).build();
    let db = database(&p.database_url());

    p.command("setup").run();

    let schema = read_file(&backend_file_path(test_name, "schema.sql"));
    db.execute(&schema);

    let result = p.command("print-schema").args(args).run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    let expected = read_file(&backend_file_path(test_name, "expected.rs"));

    assert_diff!(&expected, result.stdout(), NEW_LINE, 0);

    test_print_schema_config(test_name, &test_path, schema, expected);
}

fn test_print_schema_config(test_name: &str, test_path: &Path, schema: String, expected: String) {
    let config = read_file(&test_path.join("diesel.toml"));
    let mut p = project(&format!("{}_config", test_name)).file("diesel.toml", &config);

    let patch_file = backend_file_path(test_name, "schema.patch");
    if patch_file.exists() {
        let patch_contents = read_file(&patch_file);
        p = p.file("schema.patch", &patch_contents);
    }

    let p = p.build();

    p.command("setup").run();
    p.create_migration("12345_create_schema", &schema, "");

    let result = p.command("migration").arg("run").run();
    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    let schema = p.file_contents("src/schema.rs");
    assert_diff!(&expected, &schema, NEW_LINE, 0);

    let result = p.command("print-schema").run();
    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    assert_diff!(&expected, result.stdout(), NEW_LINE, 0);
}

fn read_file(path: &Path) -> String {
    let mut file = File::open(path).expect(&format!("Could not open {}", path.display()));
    let mut string = String::new();
    file.read_to_string(&mut string).unwrap();
    string
}
<|endoftext|>


<commit_before>//
//
//
//! x86 ATA driver
#![feature(no_std,core)]
#![no_std]
#[macro_use] extern crate core;
#[macro_use] extern crate kernel;
use kernel::_common::*;
use kernel::lib::mem::Arc;

use kernel::device_manager;
use kernel::metadevs::storage;
use kernel::async::{ReadHandle,WriteHandle};

module_define!{ATA, [DeviceManager, Storage], init}

mod drivers;
mod io;

struct AtaVolume
{
	name: String,
	disk: u8,
	controller: Arc<io::DmaController>,

	size: u64,
}

/// Initial controller handle, owns all volumes and the first controller handle
struct ControllerRoot
{
	_controller: Arc<io::DmaController>,
	_volumes: Vec<storage::PhysicalVolumeReg>,
}

enum AtaClass
{
	Invalid,	// No valid response (timeout)
	None,	// No disk
	Unknown(u8,u8),	// Unknown type, values are regs 4 and 5
	Native,	// A standard ATA disk
	ATAPI,
}
impl Default for AtaClass { fn default() -> AtaClass { AtaClass::Invalid } }

/// ATA "IDENTIFY" packet data
#[repr(C,packed)]
//#[derive(Debug)]
struct AtaIdentifyData
{
	flags: u16,
	_unused1: [u16; 9],
	serial_number: [u8; 20],
	_unused2: [u16; 3],
	firmware_ver: [u8; 8],
	model_number: [u8; 40],
	/// NFI, TODO look up
	sect_per_int: u16,
	_unused3: u16,
	capabilities: [u16; 2],
	_unused4: [u16; 2],
	/// No idea
	valid_ext_data: u16,
	_unused5: [u16; 5],
	size_of_rw_multiple: u16,
	/// LBA 28 sector count (if zero, use 48)
	sector_count_28: u32,
	_unused6: [u16; 100-62],
	/// LBA 48 sector count
	sector_count_48: u64,
	_unused7: [u16; 256-104],
}
impl Default for AtaIdentifyData {
	fn default() -> AtaIdentifyData { unsafe { ::core::mem::zeroed() } }
}
impl ::core::fmt::Debug for AtaIdentifyData {
	fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result
	{
		try!(write!(f, "AtaIdentifyData {{"));
		try!(write!(f, " flags: {:#x}", self.flags));
		try!(write!(f, " serial_number: {:?}", ::kernel::lib::RawString(&self.serial_number)));
		try!(write!(f, " firmware_ver: {:?}", ::kernel::lib::RawString(&self.firmware_ver)));
		try!(write!(f, " model_number: {:?}", ::kernel::lib::RawString(&self.model_number)));
		try!(write!(f, " sect_per_int: {}", self.sect_per_int));
		try!(write!(f, " capabilities: [{:#x},{:#x}]", self.capabilities[0], self.capabilities[1]));
		try!(write!(f, " valid_ext_data: {}", self.valid_ext_data));
		try!(write!(f, " size_of_rw_multiple: {}", self.size_of_rw_multiple));
		try!(write!(f, " sector_count_28: {:#x}", self.sector_count_28));
		try!(write!(f, " sector_count_48: {:#x}", self.sector_count_48));
		try!(write!(f, "}}"));
		Ok( () )
	}
}

fn init()
{
	drivers::register();
}

impl AtaVolume
{
	fn new_boxed(dma_controller: Arc<io::DmaController>, disk: u8, sectors: u64) -> Box<AtaVolume>
	{
		Box::new( AtaVolume {
			name: format!("{}-{}", dma_controller.name, disk),
			disk: disk,
			controller: dma_controller,
			size: sectors,
			} )
	}
}

impl ::kernel::metadevs::storage::PhysicalVolume for AtaVolume
{
	fn name(&self) -> &str { &*self.name }
	fn blocksize(&self) -> usize { io::SECTOR_SIZE }
	fn capacity(&self) -> u64 { self.size }

	fn read<'s>(&'s self, _prio: u8, idx: u64, num: usize, dst: &'s mut [u8]) -> Result<ReadHandle<'s, 's>, ()>
	{
		assert_eq!( dst.len(), num * io::SECTOR_SIZE );
		let wh = try!( self.controller.do_dma(idx, num, dst, self.disk, false));
		Ok( ReadHandle::new(dst, wh) )
	}
	fn write<'s>(&'s self, _prio: u8, idx: u64, num: usize, src: &'s [u8]) -> Result<WriteHandle<'s, 's>, ()>
	{
		assert_eq!( src.len(), num * io::SECTOR_SIZE );
		let ctrlr = &self.controller;
		// Safe cast, as controller should not modify the buffer when write=true
		match ctrlr.do_dma(idx, num, src, self.disk, true)
		{
		Err(e) => Err(e),
		Ok(v) => Ok( WriteHandle::new(src, v) )
		}
	}

	fn wipe(&mut self, _blockidx: u64, _count: usize)
	{
		// Do nothing, no support for TRIM
	}

}

impl ControllerRoot
{
	fn new(ata_pri: u16, sts_pri: u16, irq_pri: u32,  ata_sec: u16, sts_sec: u16, irq_sec: u32,  bm: device_manager::IOBinding) -> ControllerRoot
	{

		let dma_controller = Arc::new(io::DmaController {
			name: format!("ATA[{:#x},{:#x}]", ata_pri, ata_sec),
			ata_controllers: [
				io::AtaController::new(ata_pri, sts_pri, irq_pri),
				io::AtaController::new(ata_sec, sts_sec, irq_sec),
				],
			dma_base: bm,
			});
		let mut volumes = Vec::new();

		// Send IDENTIFY to all disks
		for i in (0 .. 2)
		{
			let ctrlr_pri = &dma_controller.ata_controllers[0];
			let ctrlr_sec = &dma_controller.ata_controllers[1];

			// Create output data (defaulted, but should be written by the output function)
			let (mut identify_pri, mut type_pri) = Default::default();
			let (mut identify_sec, mut type_sec) = Default::default();

			// Perform IDENTIFY requests, both controllers in pararllel
			// TODO: Include a timeout to prevent a misbehaving controller from halting the system.
			{
				let mut wh_pri = ctrlr_pri.ata_identify(i, &mut identify_pri, &mut type_pri);
				let mut wh_sec = ctrlr_sec.ata_identify(i, &mut identify_sec, &mut type_sec);
				//let mut wh_timer = ::kernel::async::Timer::new(2*1000);

				// Wait for both complete, and obtain results
				// - Loop while the timer hasn't fired, and at least one of the waiters is still waiting
				while /* !wh_timer.is_ready() && */ !(wh_pri.is_ready() && wh_sec.is_ready())
				{
					//::kernel::async::wait_on_list(&mut [&mut wh_pri, &mut wh_sec, &mut wh_timer]);
					::kernel::async::wait_on_list(&mut [&mut wh_pri, &mut wh_sec], None);
				}
			}

			// (ugly) Handle the relevant disk types, creating devices
			for &(disk, ref class, ref ident) in [(i*2, type_pri, identify_pri), (i*2+1, type_sec, identify_sec)].iter()
			{
				match *class
				{
				AtaClass::Invalid => {
					log_log!("ATA{}: Timeout", disk);
					},
				AtaClass::None => {
					log_log!("ATA{}: No disk", disk);
					},
				AtaClass::Native => {
					let sectors = if ident.sector_count_48 == 0 { ident.sector_count_28 as u64 } else { ident.sector_count_48 };
					log_log!("ATA{}: Hard Disk, {} sectors, {}", disk, sectors, storage::SizePrinter(sectors * io::SECTOR_SIZE as u64));
					volumes.push( storage::register_pv( AtaVolume::new_boxed(dma_controller.clone(), disk, sectors) ) );
					},
				AtaClass::ATAPI => {
					log_log!("ATA{}: ATAPI", disk);
					// TODO: Support ATAPI devices with a different class
					},
				AtaClass::Unknown(r4, r5) => {
					log_warning!("ATA{}: Unknown type response ({:#x}, {:#x})", disk, r4, r5);
					},
				}
			}
		}

		// Return a controller handle, holding on to all handles
		ControllerRoot { _controller: dma_controller, _volumes: volumes, }
	}

	//fn handle_volume(volumes: &mut Vec<storage::PhysicalVolumeReg>, dma_controller: &Arc<DmaController>, disk: u8, class: AtaClass, ident: AtaIdentifyData)
}

impl device_manager::DriverInstance for ControllerRoot
{
	// Just a marker trait
}
<commit_msg>Modules/ATA - (minor) Commenting<commit_after>// "Tifflin" Kernel - ATA Driver
// - By John Hodge (thePowersGang)
//
// Modules/storage_ata/lib.rs
//! x86 ATA driver
#![feature(no_std,core)]
#![no_std]
#[macro_use] extern crate core;
#[macro_use] extern crate kernel;
use kernel::_common::*;
use kernel::lib::mem::Arc;

use kernel::device_manager;
use kernel::metadevs::storage;
use kernel::async::{ReadHandle,WriteHandle};

module_define!{ATA, [DeviceManager, Storage], init}

mod drivers;
mod io;

struct AtaVolume
{
	name: String,
	disk: u8,
	controller: Arc<io::DmaController>,

	size: u64,
}

/// Initial controller handle, owns all volumes and the first controller handle
struct ControllerRoot
{
	_controller: Arc<io::DmaController>,
	_volumes: Vec<storage::PhysicalVolumeReg>,
}

enum AtaClass
{
	Invalid,	// No valid response (timeout)
	None,	// No disk
	Unknown(u8,u8),	// Unknown type, values are regs 4 and 5
	Native,	// A standard ATA disk
	ATAPI,
}
impl Default for AtaClass { fn default() -> AtaClass { AtaClass::Invalid } }

/// ATA "IDENTIFY" packet data
#[repr(C,packed)]
struct AtaIdentifyData
{
	flags: u16,
	_unused1: [u16; 9],
	serial_number: [u8; 20],
	_unused2: [u16; 3],
	firmware_ver: [u8; 8],
	model_number: [u8; 40],
	/// NFI, TODO look up
	sect_per_int: u16,
	_unused3: u16,
	capabilities: [u16; 2],
	_unused4: [u16; 2],
	/// No idea
	valid_ext_data: u16,
	_unused5: [u16; 5],
	size_of_rw_multiple: u16,
	/// LBA 28 sector count (if zero, use 48)
	sector_count_28: u32,
	_unused6: [u16; 100-62],
	/// LBA 48 sector count
	sector_count_48: u64,
	_unused7: [u16; 256-104],
}
impl Default for AtaIdentifyData {
	fn default() -> AtaIdentifyData { unsafe { ::core::mem::zeroed() } }
}
impl ::core::fmt::Debug for AtaIdentifyData {
	fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result
	{
		try!(write!(f, "AtaIdentifyData {{"));
		try!(write!(f, " flags: {:#x}", self.flags));
		try!(write!(f, " serial_number: {:?}", ::kernel::lib::RawString(&self.serial_number)));
		try!(write!(f, " firmware_ver: {:?}", ::kernel::lib::RawString(&self.firmware_ver)));
		try!(write!(f, " model_number: {:?}", ::kernel::lib::RawString(&self.model_number)));
		try!(write!(f, " sect_per_int: {}", self.sect_per_int));
		try!(write!(f, " capabilities: [{:#x},{:#x}]", self.capabilities[0], self.capabilities[1]));
		try!(write!(f, " valid_ext_data: {}", self.valid_ext_data));
		try!(write!(f, " size_of_rw_multiple: {}", self.size_of_rw_multiple));
		try!(write!(f, " sector_count_28: {:#x}", self.sector_count_28));
		try!(write!(f, " sector_count_48: {:#x}", self.sector_count_48));
		try!(write!(f, "}}"));
		Ok( () )
	}
}

fn init()
{
	drivers::register();
}

impl AtaVolume
{
	fn new_boxed(dma_controller: Arc<io::DmaController>, disk: u8, sectors: u64) -> Box<AtaVolume>
	{
		Box::new( AtaVolume {
			name: format!("{}-{}", dma_controller.name, disk),
			disk: disk,
			controller: dma_controller,
			size: sectors,
			} )
	}
}

impl ::kernel::metadevs::storage::PhysicalVolume for AtaVolume
{
	fn name(&self) -> &str { &*self.name }
	fn blocksize(&self) -> usize { io::SECTOR_SIZE }
	fn capacity(&self) -> u64 { self.size }

	fn read<'s>(&'s self, _prio: u8, idx: u64, num: usize, dst: &'s mut [u8]) -> Result<ReadHandle<'s, 's>, ()>
	{
		assert_eq!( dst.len(), num * io::SECTOR_SIZE );
		let wh = try!( self.controller.do_dma(idx, num, dst, self.disk, false));
		Ok( ReadHandle::new(dst, wh) )
	}
	fn write<'s>(&'s self, _prio: u8, idx: u64, num: usize, src: &'s [u8]) -> Result<WriteHandle<'s, 's>, ()>
	{
		assert_eq!( src.len(), num * io::SECTOR_SIZE );
		let ctrlr = &self.controller;
		// Safe cast, as controller should not modify the buffer when write=true
		match ctrlr.do_dma(idx, num, src, self.disk, true)
		{
		Err(e) => Err(e),
		Ok(v) => Ok( WriteHandle::new(src, v) )
		}
	}

	fn wipe(&mut self, _blockidx: u64, _count: usize)
	{
		// Do nothing, no support for TRIM
	}

}

impl ControllerRoot
{
	fn new(ata_pri: u16, sts_pri: u16, irq_pri: u32,  ata_sec: u16, sts_sec: u16, irq_sec: u32,  bm: device_manager::IOBinding) -> ControllerRoot
	{

		let dma_controller = Arc::new(io::DmaController {
			name: format!("ATA[{:#x},{:#x}]", ata_pri, ata_sec),
			ata_controllers: [
				io::AtaController::new(ata_pri, sts_pri, irq_pri),
				io::AtaController::new(ata_sec, sts_sec, irq_sec),
				],
			dma_base: bm,
			});
		let mut volumes = Vec::new();

		// Send IDENTIFY to all disks
		for i in (0 .. 2)
		{
			let ctrlr_pri = &dma_controller.ata_controllers[0];
			let ctrlr_sec = &dma_controller.ata_controllers[1];

			// Create output data (defaulted, but should be written by the output function)
			let (mut identify_pri, mut type_pri) = Default::default();
			let (mut identify_sec, mut type_sec) = Default::default();

			// Perform IDENTIFY requests, both controllers in pararllel
			// TODO: Include a timeout to prevent a misbehaving controller from halting the system.
			{
				let mut wh_pri = ctrlr_pri.ata_identify(i, &mut identify_pri, &mut type_pri);
				let mut wh_sec = ctrlr_sec.ata_identify(i, &mut identify_sec, &mut type_sec);
				//let mut wh_timer = ::kernel::async::Timer::new(2*1000);

				// Wait for both complete, and obtain results
				// - Loop while the timer hasn't fired, and at least one of the waiters is still waiting
				while /* !wh_timer.is_ready() && */ !(wh_pri.is_ready() && wh_sec.is_ready())
				{
					//::kernel::async::wait_on_list(&mut [&mut wh_pri, &mut wh_sec, &mut wh_timer]);
					::kernel::async::wait_on_list(&mut [&mut wh_pri, &mut wh_sec], None);
				}
			}

			// (ugly) Handle the relevant disk types, creating devices
			for &(disk, ref class, ref ident) in [(i*2, type_pri, identify_pri), (i*2+1, type_sec, identify_sec)].iter()
			{
				match *class
				{
				AtaClass::Invalid => {
					log_log!("ATA{}: Timeout", disk);
					},
				AtaClass::None => {
					log_log!("ATA{}: No disk", disk);
					},
				AtaClass::Native => {
					let sectors = if ident.sector_count_48 == 0 { ident.sector_count_28 as u64 } else { ident.sector_count_48 };
					log_log!("ATA{}: Hard Disk, {} sectors, {}", disk, sectors, storage::SizePrinter(sectors * io::SECTOR_SIZE as u64));
					volumes.push( storage::register_pv( AtaVolume::new_boxed(dma_controller.clone(), disk, sectors) ) );
					},
				AtaClass::ATAPI => {
					log_log!("ATA{}: ATAPI", disk);
					// TODO: Support ATAPI devices with a different class
					},
				AtaClass::Unknown(r4, r5) => {
					log_warning!("ATA{}: Unknown type response ({:#x}, {:#x})", disk, r4, r5);
					},
				}
			}
		}

		// Return a controller handle, holding on to all handles
		ControllerRoot { _controller: dma_controller, _volumes: volumes, }
	}

	//fn handle_volume(volumes: &mut Vec<storage::PhysicalVolumeReg>, dma_controller: &Arc<DmaController>, disk: u8, class: AtaClass, ident: AtaIdentifyData)
}

impl device_manager::DriverInstance for ControllerRoot
{
	// Just a marker trait
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counters,
    errors::{is_internal_error, JsonRpcError},
    methods::{build_registry, JsonRpcRequest, JsonRpcService, RpcRegistry},
    response::JsonRpcResponse,
};
use futures::future::join_all;
use libra_config::config::{NodeConfig, RoleType};
use libra_logger::{info, Level, Schema};
use libra_mempool::MempoolClientSender;
use libra_types::{chain_id::ChainId, ledger_info::LedgerInfoWithSignatures};
use rand::{rngs::OsRng, RngCore};
use serde_json::{map::Map, Value};
use std::{net::SocketAddr, sync::Arc};
use storage_interface::DbReader;
use tokio::runtime::{Builder, Runtime};
use warp::{
    reject::{self, Reject},
    Filter,
};

// Counter labels for runtime metrics
const LABEL_FAIL: &str = "fail";
const LABEL_SUCCESS: &str = "success";
const LABEL_BATCH: &str = "batch";
const LABEL_SINGLE: &str = "single";

#[derive(Schema)]
struct HttpRequestLog<'a> {
    #[schema(display)]
    remote_addr: Option<std::net::SocketAddr>,
    method: String,
    path: String,
    status: u16,
    referer: Option<&'a str>,
    user_agent: Option<&'a str>,
    #[schema(debug)]
    elapsed: std::time::Duration,
}

#[derive(Schema)]
struct RpcRequestLog {
    trace_id: u64,
    request: Value,
}

#[derive(Schema)]
struct RpcResponseLog<'a> {
    trace_id: u64,
    is_batch: bool,
    response_error: bool,
    response: &'a JsonRpcResponse,
}

#[macro_export]
macro_rules! log_response {
    ($trace_id: expr, $resp: expr, $is_batch: expr) => {
        let mut level = Level::Debug;
        if let Some(ref error) = $resp.error {
            if is_internal_error(&error.code) {
                level = Level::Error
            }
        }
        libra_logger::log!(
            level,
            RpcResponseLog {
                trace_id: $trace_id,
                is_batch: $is_batch,
                response_error: $resp.error.is_some(),
                response: $resp,
            }
        );
    };
}

/// Creates HTTP server (warp-based) that serves JSON RPC requests
/// Returns handle to corresponding Tokio runtime
pub fn bootstrap(
    address: SocketAddr,
    batch_size_limit: u16,
    page_size_limit: u16,
    content_len_limit: usize,
    libra_db: Arc<dyn DbReader>,
    mp_sender: MempoolClientSender,
    role: RoleType,
    chain_id: ChainId,
) -> Runtime {
    let runtime = Builder::new()
        .thread_name("json-rpc")
        .threaded_scheduler()
        .enable_all()
        .build()
        .expect("[json-rpc] failed to create runtime");

    let registry = Arc::new(build_registry());
    let service = JsonRpcService::new(
        libra_db,
        mp_sender,
        role,
        chain_id,
        batch_size_limit,
        page_size_limit,
    );

    let base_route = warp::any()
        .and(warp::post())
        .and(warp::header::exact("content-type", "application/json"))
        .and(warp::body::content_length_limit(content_len_limit as u64))
        .and(warp::body::json())
        .and(warp::any().map(move || service.clone()))
        .and(warp::any().map(move || Arc::clone(&registry)))
        .and_then(rpc_endpoint)
        .with(warp::log::custom(|info| {
            info!(HttpRequestLog {
                remote_addr: info.remote_addr(),
                method: info.method().to_string(),
                path: info.path().to_string(),
                status: info.status().as_u16(),
                referer: info.referer(),
                user_agent: info.user_agent(),
                elapsed: info.elapsed(),
            })
        }));

    // For now we still allow user to use "/", but user should start to move to "/v1" soon
    let route_root = warp::path::end().and(base_route.clone());

    let route_v1 = warp::path::path("v1")
        .and(warp::path::end())
        .and(base_route);

    let health_route = warp::path!("-" / "healthy")
        .and(warp::path::end())
        .map(|| "libra-node:ok");

    let full_route = health_route.or(route_v1.or(route_root));

    // Ensure that we actually bind to the socket first before spawning the
    // server tasks. This helps in tests to prevent races where a client attempts
    // to make a request before the server task is actually listening on the
    // socket.
    //
    // Note: we need to enter the runtime context first to actually bind, since
    //       tokio TcpListener can only be bound inside a tokio context.
    let server = runtime.enter(move || warp::serve(full_route).bind(address));
    runtime.handle().spawn(server);
    runtime
}

/// Creates JSON RPC endpoint by given node config
pub fn bootstrap_from_config(
    config: &NodeConfig,
    chain_id: ChainId,
    libra_db: Arc<dyn DbReader>,
    mp_sender: MempoolClientSender,
) -> Runtime {
    bootstrap(
        config.rpc.address,
        config.rpc.batch_size_limit,
        config.rpc.page_size_limit,
        config.rpc.content_length_limit,
        libra_db,
        mp_sender,
        config.base.role,
        chain_id,
    )
}

/// JSON RPC entry point
/// Handles all incoming rpc requests
/// Performs routing based on methods defined in `registry`
pub(crate) async fn rpc_endpoint(
    data: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
) -> Result<Box<dyn warp::Reply>, warp::Rejection> {
    let label = match data {
        Value::Array(_) => LABEL_BATCH,
        _ => LABEL_SINGLE,
    };
    counters::RPC_REQUESTS.with_label_values(&[label]).inc();
    let timer = counters::RPC_REQUEST_LATENCY
        .with_label_values(&[label])
        .start_timer();
    let ret = rpc_endpoint_without_metrics(data, service, registry).await;
    timer.stop_and_record();
    ret
}

async fn rpc_endpoint_without_metrics(
    data: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
) -> Result<Box<dyn warp::Reply>, warp::Rejection> {
    // take snapshot of latest version of DB to be used across all requests, especially for batched requests
    let ledger_info = service
        .get_latest_ledger_info()
        .map_err(|_| reject::custom(DatabaseError))?;

    let mut rng = OsRng;
    let trace_id = rng.next_u64();
    info!(RpcRequestLog {
        trace_id,
        request: data.clone(),
    });

    let resp = Ok(if let Value::Array(requests) = data {
        match service.validate_batch_size_limit(requests.len()) {
            Ok(_) => {
                // batch API call
                let futures = requests.into_iter().map(|req| {
                    rpc_request_handler(
                        req,
                        service.clone(),
                        Arc::clone(&registry),
                        ledger_info.clone(),
                        LABEL_BATCH,
                        trace_id,
                    )
                });
                let responses = join_all(futures).await;
                for resp in &responses {
                    log_response!(trace_id, &resp, true);
                }
                warp::reply::json(&responses)
            }
            Err(err) => {
                let mut resp = JsonRpcResponse::new(
                    service.chain_id(),
                    ledger_info.ledger_info().version(),
                    ledger_info.ledger_info().timestamp_usecs(),
                );
                set_response_error(&mut resp, err, LABEL_BATCH, "unknown");
                log_response!(trace_id, &resp, true);

                warp::reply::json(&resp)
            }
        }
    } else {
        // single API call
        let resp =
            rpc_request_handler(data, service, registry, ledger_info, LABEL_SINGLE, trace_id).await;
        log_response!(trace_id, &resp, false);

        warp::reply::json(&resp)
    });

    Ok(Box::new(resp) as Box<dyn warp::Reply>)
}

/// Handler of single RPC request
/// Performs validation and executes corresponding rpc handler
async fn rpc_request_handler(
    req: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
    ledger_info: LedgerInfoWithSignatures,
    request_type_label: &str,
    trace_id: u64,
) -> JsonRpcResponse {
    let request: Map<String, Value>;
    let mut response = JsonRpcResponse::new(
        service.chain_id(),
        ledger_info.ledger_info().version(),
        ledger_info.ledger_info().timestamp_usecs(),
    );

    match req {
        Value::Object(data) => {
            request = data;
        }
        _ => {
            set_response_error(
                &mut response,
                JsonRpcError::invalid_format(),
                request_type_label,
                "unknown",
            );
            return response;
        }
    }

    // parse request id
    match parse_request_id(&request) {
        Ok(request_id) => {
            response.id = Some(request_id);
        }
        Err(err) => {
            set_response_error(&mut response, err, request_type_label, "unknown");
            return response;
        }
    };

    // verify protocol version
    if let Err(err) = verify_protocol(&request) {
        set_response_error(&mut response, err, request_type_label, "unknown");
        return response;
    }

    // parse parameters
    let params;
    match request.get("params") {
        Some(Value::Array(parameters)) => {
            params = parameters.to_vec();
        }
        _ => {
            set_response_error(
                &mut response,
                JsonRpcError::invalid_params(None),
                request_type_label,
                "unknown",
            );
            return response;
        }
    }

    let request_params = JsonRpcRequest {
        trace_id,
        ledger_info,
        params,
    };
    // get rpc handler
    match request.get("method") {
        Some(Value::String(name)) => match registry.get(name) {
            Some(handler) => {
                let timer = counters::METHOD_LATENCY
                    .with_label_values(&[request_type_label, name])
                    .start_timer();
                match handler(service, request_params).await {
                    Ok(result) => {
                        response.result = Some(result);
                        counters::REQUESTS
                            .with_label_values(&[request_type_label, name, LABEL_SUCCESS])
                            .inc();
                    }
                    Err(err) => {
                        // check for custom error
                        set_response_error(
                            &mut response,
                            err.downcast_ref::<JsonRpcError>()
                                .cloned()
                                .unwrap_or_else(|| JsonRpcError::internal_error(err.to_string())),
                            request_type_label,
                            &name,
                        );
                        counters::REQUESTS
                            .with_label_values(&[request_type_label, name, LABEL_FAIL])
                            .inc();
                    }
                }
                timer.stop_and_record();
            }
            None => set_response_error(
                &mut response,
                JsonRpcError::method_not_found(),
                request_type_label,
                "not_found",
            ),
        },
        _ => set_response_error(
            &mut response,
            JsonRpcError::method_not_found(),
            request_type_label,
            "not_found",
        ),
    }

    response
}

// Sets the JSON RPC error value for a given response.
// If a counter label is supplied, also increments the invalid request counter using the label,
fn set_response_error(
    response: &mut JsonRpcResponse,
    error: JsonRpcError,
    request_type: &str,
    method: &str,
) {
    let err_code = error.code;
    if is_internal_error(&error.code) {
        counters::INTERNAL_ERRORS
            .with_label_values(&[request_type, method, &err_code.to_string()])
            .inc();
    } else {
        let label = match err_code {
            -32600 => "invalid_request",
            -32601 => "method_not_found",
            -32602 => "invalid_params",
            -32604 => "invalid_format",
            _ => "unexpected_code",
        };
        counters::INVALID_REQUESTS
            .with_label_values(&[request_type, method, label])
            .inc();
    }

    response.error = Some(error);
}

fn parse_request_id(request: &Map<String, Value>) -> Result<Value, JsonRpcError> {
    match request.get("id") {
        Some(req_id) => {
            if req_id.is_string() || req_id.is_number() || req_id.is_null() {
                Ok(req_id.clone())
            } else {
                Err(JsonRpcError::invalid_format())
            }
        }
        None => Ok(Value::Null),
    }
}

fn verify_protocol(request: &Map<String, Value>) -> Result<(), JsonRpcError> {
    if let Some(Value::String(protocol)) = request.get("jsonrpc") {
        if protocol == "2.0" {
            return Ok(());
        }
    }
    Err(JsonRpcError::invalid_request())
}

/// Warp rejection types
#[derive(Debug)]
struct DatabaseError;

impl Reject for DatabaseError {}
<commit_msg>[json-rpc] CORS: allowing all origins for the interface<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counters,
    errors::{is_internal_error, JsonRpcError},
    methods::{build_registry, JsonRpcRequest, JsonRpcService, RpcRegistry},
    response::JsonRpcResponse,
};
use futures::future::join_all;
use libra_config::config::{NodeConfig, RoleType};
use libra_logger::{info, Level, Schema};
use libra_mempool::MempoolClientSender;
use libra_types::{chain_id::ChainId, ledger_info::LedgerInfoWithSignatures};
use rand::{rngs::OsRng, RngCore};
use serde_json::{map::Map, Value};
use std::{net::SocketAddr, sync::Arc};
use storage_interface::DbReader;
use tokio::runtime::{Builder, Runtime};
use warp::{
    reject::{self, Reject},
    Filter,
};

// Counter labels for runtime metrics
const LABEL_FAIL: &str = "fail";
const LABEL_SUCCESS: &str = "success";
const LABEL_BATCH: &str = "batch";
const LABEL_SINGLE: &str = "single";

#[derive(Schema)]
struct HttpRequestLog<'a> {
    #[schema(display)]
    remote_addr: Option<std::net::SocketAddr>,
    method: String,
    path: String,
    status: u16,
    referer: Option<&'a str>,
    user_agent: Option<&'a str>,
    #[schema(debug)]
    elapsed: std::time::Duration,
}

#[derive(Schema)]
struct RpcRequestLog {
    trace_id: u64,
    request: Value,
}

#[derive(Schema)]
struct RpcResponseLog<'a> {
    trace_id: u64,
    is_batch: bool,
    response_error: bool,
    response: &'a JsonRpcResponse,
}

#[macro_export]
macro_rules! log_response {
    ($trace_id: expr, $resp: expr, $is_batch: expr) => {
        let mut level = Level::Debug;
        if let Some(ref error) = $resp.error {
            if is_internal_error(&error.code) {
                level = Level::Error
            }
        }
        libra_logger::log!(
            level,
            RpcResponseLog {
                trace_id: $trace_id,
                is_batch: $is_batch,
                response_error: $resp.error.is_some(),
                response: $resp,
            }
        );
    };
}

/// Creates HTTP server (warp-based) that serves JSON RPC requests
/// Returns handle to corresponding Tokio runtime
pub fn bootstrap(
    address: SocketAddr,
    batch_size_limit: u16,
    page_size_limit: u16,
    content_len_limit: usize,
    libra_db: Arc<dyn DbReader>,
    mp_sender: MempoolClientSender,
    role: RoleType,
    chain_id: ChainId,
) -> Runtime {
    let runtime = Builder::new()
        .thread_name("json-rpc")
        .threaded_scheduler()
        .enable_all()
        .build()
        .expect("[json-rpc] failed to create runtime");

    let registry = Arc::new(build_registry());
    let service = JsonRpcService::new(
        libra_db,
        mp_sender,
        role,
        chain_id,
        batch_size_limit,
        page_size_limit,
    );

    let base_route = warp::any()
        .and(warp::post())
        .and(warp::header::exact("content-type", "application/json"))
        .and(warp::body::content_length_limit(content_len_limit as u64))
        .and(warp::body::json())
        .and(warp::any().map(move || service.clone()))
        .and(warp::any().map(move || Arc::clone(&registry)))
        .and_then(rpc_endpoint)
        .with(warp::log::custom(|info| {
            info!(HttpRequestLog {
                remote_addr: info.remote_addr(),
                method: info.method().to_string(),
                path: info.path().to_string(),
                status: info.status().as_u16(),
                referer: info.referer(),
                user_agent: info.user_agent(),
                elapsed: info.elapsed(),
            })
        }))
        .with(warp::cors().allow_any_origin().allow_methods(vec!["POST"]));

    // For now we still allow user to use "/", but user should start to move to "/v1" soon
    let route_root = warp::path::end().and(base_route.clone());

    let route_v1 = warp::path::path("v1")
        .and(warp::path::end())
        .and(base_route);

    let health_route = warp::path!("-" / "healthy")
        .and(warp::path::end())
        .map(|| "libra-node:ok");

    let full_route = health_route.or(route_v1.or(route_root));

    // Ensure that we actually bind to the socket first before spawning the
    // server tasks. This helps in tests to prevent races where a client attempts
    // to make a request before the server task is actually listening on the
    // socket.
    //
    // Note: we need to enter the runtime context first to actually bind, since
    //       tokio TcpListener can only be bound inside a tokio context.
    let server = runtime.enter(move || warp::serve(full_route).bind(address));
    runtime.handle().spawn(server);
    runtime
}

/// Creates JSON RPC endpoint by given node config
pub fn bootstrap_from_config(
    config: &NodeConfig,
    chain_id: ChainId,
    libra_db: Arc<dyn DbReader>,
    mp_sender: MempoolClientSender,
) -> Runtime {
    bootstrap(
        config.rpc.address,
        config.rpc.batch_size_limit,
        config.rpc.page_size_limit,
        config.rpc.content_length_limit,
        libra_db,
        mp_sender,
        config.base.role,
        chain_id,
    )
}

/// JSON RPC entry point
/// Handles all incoming rpc requests
/// Performs routing based on methods defined in `registry`
pub(crate) async fn rpc_endpoint(
    data: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
) -> Result<Box<dyn warp::Reply>, warp::Rejection> {
    let label = match data {
        Value::Array(_) => LABEL_BATCH,
        _ => LABEL_SINGLE,
    };
    counters::RPC_REQUESTS.with_label_values(&[label]).inc();
    let timer = counters::RPC_REQUEST_LATENCY
        .with_label_values(&[label])
        .start_timer();
    let ret = rpc_endpoint_without_metrics(data, service, registry).await;
    timer.stop_and_record();
    ret
}

async fn rpc_endpoint_without_metrics(
    data: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
) -> Result<Box<dyn warp::Reply>, warp::Rejection> {
    // take snapshot of latest version of DB to be used across all requests, especially for batched requests
    let ledger_info = service
        .get_latest_ledger_info()
        .map_err(|_| reject::custom(DatabaseError))?;

    let mut rng = OsRng;
    let trace_id = rng.next_u64();
    info!(RpcRequestLog {
        trace_id,
        request: data.clone(),
    });

    let resp = Ok(if let Value::Array(requests) = data {
        match service.validate_batch_size_limit(requests.len()) {
            Ok(_) => {
                // batch API call
                let futures = requests.into_iter().map(|req| {
                    rpc_request_handler(
                        req,
                        service.clone(),
                        Arc::clone(&registry),
                        ledger_info.clone(),
                        LABEL_BATCH,
                        trace_id,
                    )
                });
                let responses = join_all(futures).await;
                for resp in &responses {
                    log_response!(trace_id, &resp, true);
                }
                warp::reply::json(&responses)
            }
            Err(err) => {
                let mut resp = JsonRpcResponse::new(
                    service.chain_id(),
                    ledger_info.ledger_info().version(),
                    ledger_info.ledger_info().timestamp_usecs(),
                );
                set_response_error(&mut resp, err, LABEL_BATCH, "unknown");
                log_response!(trace_id, &resp, true);

                warp::reply::json(&resp)
            }
        }
    } else {
        // single API call
        let resp =
            rpc_request_handler(data, service, registry, ledger_info, LABEL_SINGLE, trace_id).await;
        log_response!(trace_id, &resp, false);

        warp::reply::json(&resp)
    });

    Ok(Box::new(resp) as Box<dyn warp::Reply>)
}

/// Handler of single RPC request
/// Performs validation and executes corresponding rpc handler
async fn rpc_request_handler(
    req: Value,
    service: JsonRpcService,
    registry: Arc<RpcRegistry>,
    ledger_info: LedgerInfoWithSignatures,
    request_type_label: &str,
    trace_id: u64,
) -> JsonRpcResponse {
    let request: Map<String, Value>;
    let mut response = JsonRpcResponse::new(
        service.chain_id(),
        ledger_info.ledger_info().version(),
        ledger_info.ledger_info().timestamp_usecs(),
    );

    match req {
        Value::Object(data) => {
            request = data;
        }
        _ => {
            set_response_error(
                &mut response,
                JsonRpcError::invalid_format(),
                request_type_label,
                "unknown",
            );
            return response;
        }
    }

    // parse request id
    match parse_request_id(&request) {
        Ok(request_id) => {
            response.id = Some(request_id);
        }
        Err(err) => {
            set_response_error(&mut response, err, request_type_label, "unknown");
            return response;
        }
    };

    // verify protocol version
    if let Err(err) = verify_protocol(&request) {
        set_response_error(&mut response, err, request_type_label, "unknown");
        return response;
    }

    // parse parameters
    let params;
    match request.get("params") {
        Some(Value::Array(parameters)) => {
            params = parameters.to_vec();
        }
        _ => {
            set_response_error(
                &mut response,
                JsonRpcError::invalid_params(None),
                request_type_label,
                "unknown",
            );
            return response;
        }
    }

    let request_params = JsonRpcRequest {
        trace_id,
        ledger_info,
        params,
    };
    // get rpc handler
    match request.get("method") {
        Some(Value::String(name)) => match registry.get(name) {
            Some(handler) => {
                let timer = counters::METHOD_LATENCY
                    .with_label_values(&[request_type_label, name])
                    .start_timer();
                match handler(service, request_params).await {
                    Ok(result) => {
                        response.result = Some(result);
                        counters::REQUESTS
                            .with_label_values(&[request_type_label, name, LABEL_SUCCESS])
                            .inc();
                    }
                    Err(err) => {
                        // check for custom error
                        set_response_error(
                            &mut response,
                            err.downcast_ref::<JsonRpcError>()
                                .cloned()
                                .unwrap_or_else(|| JsonRpcError::internal_error(err.to_string())),
                            request_type_label,
                            &name,
                        );
                        counters::REQUESTS
                            .with_label_values(&[request_type_label, name, LABEL_FAIL])
                            .inc();
                    }
                }
                timer.stop_and_record();
            }
            None => set_response_error(
                &mut response,
                JsonRpcError::method_not_found(),
                request_type_label,
                "not_found",
            ),
        },
        _ => set_response_error(
            &mut response,
            JsonRpcError::method_not_found(),
            request_type_label,
            "not_found",
        ),
    }

    response
}

// Sets the JSON RPC error value for a given response.
// If a counter label is supplied, also increments the invalid request counter using the label,
fn set_response_error(
    response: &mut JsonRpcResponse,
    error: JsonRpcError,
    request_type: &str,
    method: &str,
) {
    let err_code = error.code;
    if is_internal_error(&error.code) {
        counters::INTERNAL_ERRORS
            .with_label_values(&[request_type, method, &err_code.to_string()])
            .inc();
    } else {
        let label = match err_code {
            -32600 => "invalid_request",
            -32601 => "method_not_found",
            -32602 => "invalid_params",
            -32604 => "invalid_format",
            _ => "unexpected_code",
        };
        counters::INVALID_REQUESTS
            .with_label_values(&[request_type, method, label])
            .inc();
    }

    response.error = Some(error);
}

fn parse_request_id(request: &Map<String, Value>) -> Result<Value, JsonRpcError> {
    match request.get("id") {
        Some(req_id) => {
            if req_id.is_string() || req_id.is_number() || req_id.is_null() {
                Ok(req_id.clone())
            } else {
                Err(JsonRpcError::invalid_format())
            }
        }
        None => Ok(Value::Null),
    }
}

fn verify_protocol(request: &Map<String, Value>) -> Result<(), JsonRpcError> {
    if let Some(Value::String(protocol)) = request.get("jsonrpc") {
        if protocol == "2.0" {
            return Ok(());
        }
    }
    Err(JsonRpcError::invalid_request())
}

/// Warp rejection types
#[derive(Debug)]
struct DatabaseError;

impl Reject for DatabaseError {}
<|endoftext|>