<commit_before>use parser::{Cond, Expr, Node, Target, WS};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator<'a> {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
    next_ws: Option<&'a str>,
    skip_ws: bool,
}

impl<'a> Generator<'a> {

    fn new() -> Generator<'a> {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
            next_ws: None,
            skip_ws: false,
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &str) {
        if self.locals.contains(s) {
            self.write(s);
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_binop(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, val),
            Expr::BinOp(op, ref left, ref right) =>
                self.visit_binop(op, left, right),
        }
    }

    fn visit_target_single(&mut self, name: &str) -> Vec<String> {
        vec![name.to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn flush_ws(&mut self, ws: &WS) {
        if self.next_ws.is_some() && !ws.0 {
            let val = self.next_ws.unwrap();
            if !val.is_empty() {
                self.writeln(&format!("writer.write_str({:#?}).unwrap();",
                                      val));
            }
        }
        self.next_ws = None;
    }

    fn prepare_ws(&mut self, ws: &WS) {
        self.skip_ws = ws.1;
    }

    fn handle_ws(&mut self, ws: &WS) {
        self.flush_ws(ws);
        self.prepare_ws(ws);
    }

    fn write_lit(&mut self, lws: &'a str, val: &str, rws: &'a str) {
        assert!(self.next_ws.is_none());
        if !lws.is_empty() {
            if self.skip_ws {
                self.skip_ws = false;
            } else if val.is_empty() {
                assert!(rws.is_empty());
                self.next_ws = Some(lws);
            } else {
                self.writeln(&format!("writer.write_str({:#?}).unwrap();",
                                      lws));
            }
        }
        if !val.is_empty() {
            self.writeln(&format!("writer.write_str({:#?}).unwrap();", val));
        }
        if !rws.is_empty() {
            self.next_ws = Some(rws);
        }
    }

    fn write_expr(&mut self, ws: &WS, s: &Expr) {
        self.handle_ws(ws);
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &'a [Cond], ws: &WS) {
        for (i, &(ref cws, ref cond, ref nodes)) in conds.iter().enumerate() {
            self.handle_ws(cws);
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.handle_ws(ws);
        self.writeln("}");
    }

    fn write_loop(&mut self, ws1: &WS, var: &Target, iter: &Expr,
                  body: &'a [Node], ws2: &WS) {

        self.handle_ws(ws1);
        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(name);
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.handle_ws(ws2);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, ws1: &WS, name: &str, ws2: &WS) {
        self.flush_ws(ws1);
        self.writeln(&format!("timpl.render_block_{}_to(writer);", name));
        self.prepare_ws(ws2);
    }

    fn write_block_def(&mut self, ws1: &WS, name: &str, nodes: &'a [Node],
                       ws2: &WS) {
        self.writeln("#[allow(unused_variables)]");
        self.writeln(&format!(
            "fn impl_block_{}_to(&self, writer: &mut std::fmt::Write) {{",
            name));
        self.indent();
        self.prepare_ws(ws1);
        self.handle(nodes);
        self.flush_ws(ws2);
        self.dedent();
        self.writeln("}");
    }

    fn handle(&mut self, nodes: &'a [Node]) {
        for n in nodes {
            match *n {
                Node::Lit(lws, val, rws) => { self.write_lit(lws, val, rws); }
                Node::Expr(ref ws, ref val) => { self.write_expr(ws, val); },
                Node::Cond(ref conds, ref ws) => {
                    self.write_cond(conds, ws);
                },
                Node::Loop(ref ws1, ref var, ref iter, ref body, ref ws2) => {
                    self.write_loop(ws1, var, iter, body, ws2);
                },
                Node::Block(ref ws1, name, ref ws2) => {
                    self.write_block(ws1, name, ws2);
                },
                Node::BlockDef(ref ws1, name, ref block_nodes, ref ws2) => {
                    self.write_block_def(ws1, name, block_nodes, ws2);
                }
                Node::Extends(_) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn struct_impl(&mut self, ast: &syn::DeriveInput, path: &str,
                   blocks: &'a [Node], nodes: Option<&'a [Node]>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} {}{} {{", anno, ast.ident.as_ref(),
                              anno));
        self.indent();

        self.handle(blocks);

        if let Some(nodes) = nodes {
            let trait_name = format!("TraitFrom{}", path_as_identifier(path));
            self.writeln(&format!("fn render_trait_to(&self, timpl: &{},",
                                  trait_name));
            self.writeln("                   writer: &mut std::fmt::Write) {");
            self.indent();
            self.handle(nodes);
            self.flush_ws(&WS(false, false));
            self.dedent();
            self.writeln("}");
        }

        self.dedent();
        self.writeln("}");
    }

    fn impl_template(&mut self, ast: &syn::DeriveInput, nodes: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_to(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.flush_ws(&WS(false, false));
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn impl_trait(&mut self, ast: &syn::DeriveInput, base: &str,
                  block_names: &[&str]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();

        for bname in block_names {
            self.writeln(&format!(
                "fn render_block_{}_to(&self, writer: &mut std::fmt::Write) {{",
                bname));
            self.indent();
            self.writeln(&format!("self.impl_block_{}_to(writer);", bname));
            self.dedent();
            self.writeln("}");
        }

        self.flush_ws(&WS(false, false));
        self.dedent();
        self.writeln("}");
    }

    fn impl_template_for_trait(&mut self, ast: &syn::DeriveInput, derived: bool) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_to(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        if derived {
            self.writeln("self._parent.render_trait_to(self, writer);");
        } else {
            self.writeln("self.render_trait_to(self, writer);");
        }
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn define_trait(&mut self, path: &str, block_names: &[&str]) {
        self.writeln(&format!("trait TraitFrom{} {{",
                              path_as_identifier(path)));
        self.indent();

        for bname in block_names {
            self.writeln(&format!(
                "fn render_block_{}_to(&self, writer: &mut std::fmt::Write);",
                bname));
        }

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut block_names = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(ws1, name, _, ws2) => {
                blocks.push(n);
                block_names.push(name);
                content.push(Node::Block(ws1, name, ws2));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if !blocks.is_empty() {
        if base.is_none() {
            gen.struct_impl(ast, path, &blocks, Some(&content));
            gen.define_trait(path, &block_names);
        } else {
            gen.struct_impl(ast, path, &blocks, None);
        }
        let tmpl_path = match base {
            Some(Expr::StrLit(base_path)) => { base_path },
            _ => { path },
        };
        gen.impl_trait(ast, tmpl_path, &block_names);
        gen.impl_template_for_trait(ast, base.is_some());
    } else {
        gen.impl_template(ast, &content);
    }
    gen.result()
}
<commit_msg>Move code for whitespace handling around<commit_after>use parser::{Cond, Expr, Node, Target, WS};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator<'a> {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
    next_ws: Option<&'a str>,
    skip_ws: bool,
}

impl<'a> Generator<'a> {

    fn new() -> Generator<'a> {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
            next_ws: None,
            skip_ws: false,
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn flush_ws(&mut self, ws: &WS) {
        if self.next_ws.is_some() && !ws.0 {
            let val = self.next_ws.unwrap();
            if !val.is_empty() {
                self.writeln(&format!("writer.write_str({:#?}).unwrap();",
                                      val));
            }
        }
        self.next_ws = None;
    }

    fn prepare_ws(&mut self, ws: &WS) {
        self.skip_ws = ws.1;
    }

    fn handle_ws(&mut self, ws: &WS) {
        self.flush_ws(ws);
        self.prepare_ws(ws);
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &str) {
        if self.locals.contains(s) {
            self.write(s);
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_binop(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, val),
            Expr::BinOp(op, ref left, ref right) =>
                self.visit_binop(op, left, right),
        }
    }

    fn visit_target_single(&mut self, name: &str) -> Vec<String> {
        vec![name.to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn write_lit(&mut self, lws: &'a str, val: &str, rws: &'a str) {
        assert!(self.next_ws.is_none());
        if !lws.is_empty() {
            if self.skip_ws {
                self.skip_ws = false;
            } else if val.is_empty() {
                assert!(rws.is_empty());
                self.next_ws = Some(lws);
            } else {
                self.writeln(&format!("writer.write_str({:#?}).unwrap();",
                                      lws));
            }
        }
        if !val.is_empty() {
            self.writeln(&format!("writer.write_str({:#?}).unwrap();", val));
        }
        if !rws.is_empty() {
            self.next_ws = Some(rws);
        }
    }

    fn write_expr(&mut self, ws: &WS, s: &Expr) {
        self.handle_ws(ws);
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &'a [Cond], ws: &WS) {
        for (i, &(ref cws, ref cond, ref nodes)) in conds.iter().enumerate() {
            self.handle_ws(cws);
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.handle_ws(ws);
        self.writeln("}");
    }

    fn write_loop(&mut self, ws1: &WS, var: &Target, iter: &Expr,
                  body: &'a [Node], ws2: &WS) {

        self.handle_ws(ws1);
        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(name);
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.handle_ws(ws2);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, ws1: &WS, name: &str, ws2: &WS) {
        self.flush_ws(ws1);
        self.writeln(&format!("timpl.render_block_{}_to(writer);", name));
        self.prepare_ws(ws2);
    }

    fn write_block_def(&mut self, ws1: &WS, name: &str, nodes: &'a [Node],
                       ws2: &WS) {
        self.writeln("#[allow(unused_variables)]");
        self.writeln(&format!(
            "fn impl_block_{}_to(&self, writer: &mut std::fmt::Write) {{",
            name));
        self.indent();
        self.prepare_ws(ws1);
        self.handle(nodes);
        self.flush_ws(ws2);
        self.dedent();
        self.writeln("}");
    }

    fn handle(&mut self, nodes: &'a [Node]) {
        for n in nodes {
            match *n {
                Node::Lit(lws, val, rws) => { self.write_lit(lws, val, rws); }
                Node::Expr(ref ws, ref val) => { self.write_expr(ws, val); },
                Node::Cond(ref conds, ref ws) => {
                    self.write_cond(conds, ws);
                },
                Node::Loop(ref ws1, ref var, ref iter, ref body, ref ws2) => {
                    self.write_loop(ws1, var, iter, body, ws2);
                },
                Node::Block(ref ws1, name, ref ws2) => {
                    self.write_block(ws1, name, ws2);
                },
                Node::BlockDef(ref ws1, name, ref block_nodes, ref ws2) => {
                    self.write_block_def(ws1, name, block_nodes, ws2);
                }
                Node::Extends(_) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn struct_impl(&mut self, ast: &syn::DeriveInput, path: &str,
                   blocks: &'a [Node], nodes: Option<&'a [Node]>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} {}{} {{", anno, ast.ident.as_ref(),
                              anno));
        self.indent();

        self.handle(blocks);

        if let Some(nodes) = nodes {
            let trait_name = format!("TraitFrom{}", path_as_identifier(path));
            self.writeln(&format!("fn render_trait_to(&self, timpl: &{},",
                                  trait_name));
            self.writeln("                   writer: &mut std::fmt::Write) {");
            self.indent();
            self.handle(nodes);
            self.flush_ws(&WS(false, false));
            self.dedent();
            self.writeln("}");
        }

        self.dedent();
        self.writeln("}");
    }

    fn impl_template(&mut self, ast: &syn::DeriveInput, nodes: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_to(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.flush_ws(&WS(false, false));
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn impl_trait(&mut self, ast: &syn::DeriveInput, base: &str,
                  block_names: &[&str]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();

        for bname in block_names {
            self.writeln(&format!(
                "fn render_block_{}_to(&self, writer: &mut std::fmt::Write) {{",
                bname));
            self.indent();
            self.writeln(&format!("self.impl_block_{}_to(writer);", bname));
            self.dedent();
            self.writeln("}");
        }

        self.flush_ws(&WS(false, false));
        self.dedent();
        self.writeln("}");
    }

    fn impl_template_for_trait(&mut self, ast: &syn::DeriveInput, derived: bool) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_to(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        if derived {
            self.writeln("self._parent.render_trait_to(self, writer);");
        } else {
            self.writeln("self.render_trait_to(self, writer);");
        }
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn define_trait(&mut self, path: &str, block_names: &[&str]) {
        self.writeln(&format!("trait TraitFrom{} {{",
                              path_as_identifier(path)));
        self.indent();

        for bname in block_names {
            self.writeln(&format!(
                "fn render_block_{}_to(&self, writer: &mut std::fmt::Write);",
                bname));
        }

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut block_names = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(ws1, name, _, ws2) => {
                blocks.push(n);
                block_names.push(name);
                content.push(Node::Block(ws1, name, ws2));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if !blocks.is_empty() {
        if base.is_none() {
            gen.struct_impl(ast, path, &blocks, Some(&content));
            gen.define_trait(path, &block_names);
        } else {
            gen.struct_impl(ast, path, &blocks, None);
        }
        let tmpl_path = match base {
            Some(Expr::StrLit(base_path)) => { base_path },
            _ => { path },
        };
        gen.impl_trait(ast, tmpl_path, &block_names);
        gen.impl_template_for_trait(ast, base.is_some());
    } else {
        gen.impl_template(ast, &content);
    }
    gen.result()
}
<|endoftext|>


<commit_before>//! Script to check the validity of `href` links in our HTML documentation.
//!
//! In the past we've been quite error prone to writing in broken links as most
//! of them are manually rather than automatically added. As files move over
//! time or apis change old links become stale or broken. The purpose of this
//! script is to check all relative links in our documentation to make sure they
//! actually point to a valid place.
//!
//! Currently this doesn't actually do any HTML parsing or anything fancy like
//! that, it just has a simple "regex" to search for `href` and `id` tags.
//! These values are then translated to file URLs if possible and then the
//! destination is asserted to exist.
//!
//! A few whitelisted exceptions are allowed as there's known bugs in rustdoc,
//! but this should catch the majority of "broken link" cases.

#![deny(warnings)]

use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs;
use std::path::{Component, Path, PathBuf};
use std::rc::Rc;

use crate::Redirect::*;

macro_rules! t {
    ($e:expr) => {
        match $e {
            Ok(e) => e,
            Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
        }
    };
}

fn main() {
    let docs = env::args_os().nth(1).unwrap();
    let docs = env::current_dir().unwrap().join(docs);
    let mut errors = false;
    walk(&mut HashMap::new(), &docs, &docs, &mut errors);
    if errors {
        panic!("found some broken links");
    }
}

#[derive(Debug)]
pub enum LoadError {
    IOError(std::io::Error),
    BrokenRedirect(PathBuf, std::io::Error),
    IsRedirect,
}

enum Redirect {
    SkipRedirect,
    FromRedirect(bool),
}

struct FileEntry {
    source: Rc<String>,
    ids: HashSet<String>,
}

type Cache = HashMap<PathBuf, FileEntry>;

fn small_url_encode(s: &str) -> String {
    s.replace("<", "%3C")
        .replace(">", "%3E")
        .replace(" ", "%20")
        .replace("?", "%3F")
        .replace("'", "%27")
        .replace("&", "%26")
        .replace(",", "%2C")
        .replace(":", "%3A")
        .replace(";", "%3B")
        .replace("[", "%5B")
        .replace("]", "%5D")
        .replace("\"", "%22")
}

impl FileEntry {
    fn parse_ids(&mut self, file: &Path, contents: &str, errors: &mut bool) {
        if self.ids.is_empty() {
            with_attrs_in_source(contents, " id", |fragment, i, _| {
                let frag = fragment.trim_start_matches("#").to_owned();
                let encoded = small_url_encode(&frag);
                if !self.ids.insert(frag) {
                    *errors = true;
                    println!("{}:{}: id is not unique: `{}`", file.display(), i, fragment);
                }
                // Just in case, we also add the encoded id.
                self.ids.insert(encoded);
            });
        }
    }
}

fn walk(cache: &mut Cache, root: &Path, dir: &Path, errors: &mut bool) {
    for entry in t!(dir.read_dir()).map(|e| t!(e)) {
        let path = entry.path();
        let kind = t!(entry.file_type());
        if kind.is_dir() {
            walk(cache, root, &path, errors);
        } else {
            let pretty_path = check(cache, root, &path, errors);
            if let Some(pretty_path) = pretty_path {
                let entry = cache.get_mut(&pretty_path).unwrap();
                // we don't need the source anymore,
                // so drop to reduce memory-usage
                entry.source = Rc::new(String::new());
            }
        }
    }
}

fn check(cache: &mut Cache, root: &Path, file: &Path, errors: &mut bool) -> Option<PathBuf> {
    // Ignore none HTML files.
    if file.extension().and_then(|s| s.to_str()) != Some("html") {
        return None;
    }

    // Unfortunately we're not 100% full of valid links today to we need a few
    // whitelists to get this past `make check` today.
    // FIXME(#32129)
    if file.ends_with("std/string/struct.String.html")
        || file.ends_with("interpret/struct.ImmTy.html")
        || file.ends_with("ast/struct.ThinVec.html")
        || file.ends_with("util/struct.ThinVec.html")
        || file.ends_with("layout/struct.TyLayout.html")
        || file.ends_with("humantime/struct.Timestamp.html")
        || file.ends_with("log/index.html")
        || file.ends_with("ty/struct.Slice.html")
        || file.ends_with("ty/enum.Attributes.html")
        || file.ends_with("ty/struct.SymbolName.html")
        || file.ends_with("io/struct.IoSlice.html")
        || file.ends_with("io/struct.IoSliceMut.html")
    {
        return None;
    }
    // FIXME(#32553)
    if file.ends_with("string/struct.String.html") {
        return None;
    }
    // FIXME(#32130)
    if file.ends_with("btree_set/struct.BTreeSet.html")
        || file.ends_with("struct.BTreeSet.html")
        || file.ends_with("btree_map/struct.BTreeMap.html")
        || file.ends_with("hash_map/struct.HashMap.html")
        || file.ends_with("hash_set/struct.HashSet.html")
        || file.ends_with("sync/struct.Lrc.html")
        || file.ends_with("sync/struct.RwLock.html")
    {
        return None;
    }

    let res = load_file(cache, root, file, SkipRedirect);
    let (pretty_file, contents) = match res {
        Ok(res) => res,
        Err(_) => return None,
    };
    {
        cache.get_mut(&pretty_file).unwrap().parse_ids(&pretty_file, &contents, errors);
    }

    // Search for anything that's the regex 'href[ ]*=[ ]*".*?"'
    with_attrs_in_source(&contents, " href", |url, i, base| {
        // Ignore external URLs
        if url.starts_with("http:")
            || url.starts_with("https:")
            || url.starts_with("javascript:")
            || url.starts_with("ftp:")
            || url.starts_with("irc:")
            || url.starts_with("data:")
        {
            return;
        }
        let mut parts = url.splitn(2, "#");
        let url = parts.next().unwrap();
        let fragment = parts.next();
        let mut parts = url.splitn(2, "?");
        let url = parts.next().unwrap();

        // Once we've plucked out the URL, parse it using our base url and
        // then try to extract a file path.
        let mut path = file.to_path_buf();
        if !base.is_empty() || !url.is_empty() {
            path.pop();
            for part in Path::new(base).join(url).components() {
                match part {
                    Component::Prefix(_) | Component::RootDir => {
                        // Avoid absolute paths as they make the docs not
                        // relocatable by making assumptions on where the docs
                        // are hosted relative to the site root.
                        *errors = true;
                        println!(
                            "{}:{}: absolute path - {}",
                            pretty_file.display(),
                            i + 1,
                            Path::new(base).join(url).display()
                        );
                        return;
                    }
                    Component::CurDir => {}
                    Component::ParentDir => {
                        path.pop();
                    }
                    Component::Normal(s) => {
                        path.push(s);
                    }
                }
            }
        }

        // Alright, if we've found a file name then this file had better
        // exist! If it doesn't then we register and print an error.
        if path.exists() {
            if path.is_dir() {
                // Links to directories show as directory listings when viewing
                // the docs offline so it's best to avoid them.
                *errors = true;
                let pretty_path = path.strip_prefix(root).unwrap_or(&path);
                println!(
                    "{}:{}: directory link - {}",
                    pretty_file.display(),
                    i + 1,
                    pretty_path.display()
                );
                return;
            }
            if let Some(extension) = path.extension() {
                // Ignore none HTML files.
                if extension != "html" {
                    return;
                }
            }
            let res = load_file(cache, root, &path, FromRedirect(false));
            let (pretty_path, contents) = match res {
                Ok(res) => res,
                Err(LoadError::IOError(err)) => {
                    panic!("error loading {}: {}", path.display(), err);
                }
                Err(LoadError::BrokenRedirect(target, _)) => {
                    *errors = true;
                    println!(
                        "{}:{}: broken redirect to {}",
                        pretty_file.display(),
                        i + 1,
                        target.display()
                    );
                    return;
                }
                Err(LoadError::IsRedirect) => unreachable!(),
            };

            if let Some(ref fragment) = fragment {
                // Fragments like `#1-6` are most likely line numbers to be
                // interpreted by javascript, so we're ignoring these
                if fragment.splitn(2, '-').all(|f| f.chars().all(|c| c.is_numeric())) {
                    return;
                }

                // These appear to be broken in mdbook right now?
                if fragment.starts_with("-") {
                    return;
                }

                let entry = &mut cache.get_mut(&pretty_path).unwrap();
                entry.parse_ids(&pretty_path, &contents, errors);

                if !entry.ids.contains(*fragment) {
                    *errors = true;
                    print!("{}:{}: broken link fragment ", pretty_file.display(), i + 1);
                    println!("`#{}` pointing to `{}`", fragment, pretty_path.display());
                };
            }
        } else {
            *errors = true;
            print!("{}:{}: broken link - ", pretty_file.display(), i + 1);
            let pretty_path = path.strip_prefix(root).unwrap_or(&path);
            println!("{}", pretty_path.display());
        }
    });
    Some(pretty_file)
}

fn load_file(
    cache: &mut Cache,
    root: &Path,
    file: &Path,
    redirect: Redirect,
) -> Result<(PathBuf, Rc<String>), LoadError> {
    let pretty_file = PathBuf::from(file.strip_prefix(root).unwrap_or(&file));

    let (maybe_redirect, contents) = match cache.entry(pretty_file.clone()) {
        Entry::Occupied(entry) => (None, entry.get().source.clone()),
        Entry::Vacant(entry) => {
            let contents = match fs::read_to_string(file) {
                Ok(s) => Rc::new(s),
                Err(err) => {
                    return Err(if let FromRedirect(true) = redirect {
                        LoadError::BrokenRedirect(file.to_path_buf(), err)
                    } else {
                        LoadError::IOError(err)
                    });
                }
            };

            let maybe = maybe_redirect(&contents);
            if maybe.is_some() {
                if let SkipRedirect = redirect {
                    return Err(LoadError::IsRedirect);
                }
            } else {
                entry.insert(FileEntry { source: contents.clone(), ids: HashSet::new() });
            }
            (maybe, contents)
        }
    };
    match maybe_redirect.map(|url| file.parent().unwrap().join(url)) {
        Some(redirect_file) => load_file(cache, root, &redirect_file, FromRedirect(true)),
        None => Ok((pretty_file, contents)),
    }
}

fn maybe_redirect(source: &str) -> Option<String> {
    const REDIRECT: &'static str = "
Redirecting to <a href=";

    let mut lines = source.lines();
    let redirect_line = lines.nth(6)?;

    redirect_line.find(REDIRECT).map(|i| {
        let rest = &redirect_line[(i + REDIRECT.len() + 1)..];
        let pos_quote = rest.find('"').unwrap();
        rest[..pos_quote].to_owned()
    })
}

fn with_attrs_in_source<F: FnMut(&str, usize, &str)>(contents: &str, attr: &str, mut f: F) {
    let mut base = "";
    for (i, mut line) in contents.lines().enumerate() {
        while let Some(j) = line.find(attr) {
            let rest = &line[j + attr.len()..];
            // The base tag should always be the first link in the document so
            // we can get away with using one pass.
            let is_base = line[..j].ends_with("<base");
            line = rest;
            let pos_equals = match rest.find("=") {
                Some(i) => i,
                None => continue,
            };
            if rest[..pos_equals].trim_start_matches(" ") != "" {
                continue;
            }

            let rest = &rest[pos_equals + 1..];

            let pos_quote = match rest.find(&['"', '\''][..]) {
                Some(i) => i,
                None => continue,
            };
            let quote_delim = rest.as_bytes()[pos_quote] as char;

            if rest[..pos_quote].trim_start_matches(" ") != "" {
                continue;
            }
            let rest = &rest[pos_quote + 1..];
            let url = match rest.find(quote_delim) {
                Some(i) => &rest[..i],
                None => continue,
            };
            if is_base {
                base = url;
                continue;
            }
            f(url, i, base)
        }
    }
}
<commit_msg>Rollup merge of #67789 - ollie27:linkchecker_whitelist, r=nikomatsakis<commit_after>//! Script to check the validity of `href` links in our HTML documentation.
//!
//! In the past we've been quite error prone to writing in broken links as most
//! of them are manually rather than automatically added. As files move over
//! time or apis change old links become stale or broken. The purpose of this
//! script is to check all relative links in our documentation to make sure they
//! actually point to a valid place.
//!
//! Currently this doesn't actually do any HTML parsing or anything fancy like
//! that, it just has a simple "regex" to search for `href` and `id` tags.
//! These values are then translated to file URLs if possible and then the
//! destination is asserted to exist.
//!
//! A few whitelisted exceptions are allowed as there's known bugs in rustdoc,
//! but this should catch the majority of "broken link" cases.

#![deny(warnings)]

use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs;
use std::path::{Component, Path, PathBuf};
use std::rc::Rc;

use crate::Redirect::*;

macro_rules! t {
    ($e:expr) => {
        match $e {
            Ok(e) => e,
            Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
        }
    };
}

fn main() {
    let docs = env::args_os().nth(1).unwrap();
    let docs = env::current_dir().unwrap().join(docs);
    let mut errors = false;
    walk(&mut HashMap::new(), &docs, &docs, &mut errors);
    if errors {
        panic!("found some broken links");
    }
}

#[derive(Debug)]
pub enum LoadError {
    IOError(std::io::Error),
    BrokenRedirect(PathBuf, std::io::Error),
    IsRedirect,
}

enum Redirect {
    SkipRedirect,
    FromRedirect(bool),
}

struct FileEntry {
    source: Rc<String>,
    ids: HashSet<String>,
}

type Cache = HashMap<PathBuf, FileEntry>;

fn small_url_encode(s: &str) -> String {
    s.replace("<", "%3C")
        .replace(">", "%3E")
        .replace(" ", "%20")
        .replace("?", "%3F")
        .replace("'", "%27")
        .replace("&", "%26")
        .replace(",", "%2C")
        .replace(":", "%3A")
        .replace(";", "%3B")
        .replace("[", "%5B")
        .replace("]", "%5D")
        .replace("\"", "%22")
}

impl FileEntry {
    fn parse_ids(&mut self, file: &Path, contents: &str, errors: &mut bool) {
        if self.ids.is_empty() {
            with_attrs_in_source(contents, " id", |fragment, i, _| {
                let frag = fragment.trim_start_matches("#").to_owned();
                let encoded = small_url_encode(&frag);
                if !self.ids.insert(frag) {
                    *errors = true;
                    println!("{}:{}: id is not unique: `{}`", file.display(), i, fragment);
                }
                // Just in case, we also add the encoded id.
                self.ids.insert(encoded);
            });
        }
    }
}

fn walk(cache: &mut Cache, root: &Path, dir: &Path, errors: &mut bool) {
    for entry in t!(dir.read_dir()).map(|e| t!(e)) {
        let path = entry.path();
        let kind = t!(entry.file_type());
        if kind.is_dir() {
            walk(cache, root, &path, errors);
        } else {
            let pretty_path = check(cache, root, &path, errors);
            if let Some(pretty_path) = pretty_path {
                let entry = cache.get_mut(&pretty_path).unwrap();
                // we don't need the source anymore,
                // so drop to reduce memory-usage
                entry.source = Rc::new(String::new());
            }
        }
    }
}

fn check(cache: &mut Cache, root: &Path, file: &Path, errors: &mut bool) -> Option<PathBuf> {
    // Ignore none HTML files.
    if file.extension().and_then(|s| s.to_str()) != Some("html") {
        return None;
    }

    // Unfortunately we're not 100% full of valid links today to we need a few
    // whitelists to get this past `make check` today.
    // FIXME(#32129)
    if file.ends_with("std/io/struct.IoSlice.html")
        || file.ends_with("std/string/struct.String.html")
    {
        return None;
    }
    // FIXME(#32553)
    if file.ends_with("alloc/string/struct.String.html") {
        return None;
    }
    // FIXME(#32130)
    if file.ends_with("alloc/collections/btree_map/struct.BTreeMap.html")
        || file.ends_with("alloc/collections/btree_set/struct.BTreeSet.html")
        || file.ends_with("std/collections/btree_map/struct.BTreeMap.html")
        || file.ends_with("std/collections/btree_set/struct.BTreeSet.html")
        || file.ends_with("std/collections/hash_map/struct.HashMap.html")
        || file.ends_with("std/collections/hash_set/struct.HashSet.html")
    {
        return None;
    }

    let res = load_file(cache, root, file, SkipRedirect);
    let (pretty_file, contents) = match res {
        Ok(res) => res,
        Err(_) => return None,
    };
    {
        cache.get_mut(&pretty_file).unwrap().parse_ids(&pretty_file, &contents, errors);
    }

    // Search for anything that's the regex 'href[ ]*=[ ]*".*?"'
    with_attrs_in_source(&contents, " href", |url, i, base| {
        // Ignore external URLs
        if url.starts_with("http:")
            || url.starts_with("https:")
            || url.starts_with("javascript:")
            || url.starts_with("ftp:")
            || url.starts_with("irc:")
            || url.starts_with("data:")
        {
            return;
        }
        let mut parts = url.splitn(2, "#");
        let url = parts.next().unwrap();
        let fragment = parts.next();
        let mut parts = url.splitn(2, "?");
        let url = parts.next().unwrap();

        // Once we've plucked out the URL, parse it using our base url and
        // then try to extract a file path.
        let mut path = file.to_path_buf();
        if !base.is_empty() || !url.is_empty() {
            path.pop();
            for part in Path::new(base).join(url).components() {
                match part {
                    Component::Prefix(_) | Component::RootDir => {
                        // Avoid absolute paths as they make the docs not
                        // relocatable by making assumptions on where the docs
                        // are hosted relative to the site root.
                        *errors = true;
                        println!(
                            "{}:{}: absolute path - {}",
                            pretty_file.display(),
                            i + 1,
                            Path::new(base).join(url).display()
                        );
                        return;
                    }
                    Component::CurDir => {}
                    Component::ParentDir => {
                        path.pop();
                    }
                    Component::Normal(s) => {
                        path.push(s);
                    }
                }
            }
        }

        // Alright, if we've found a file name then this file had better
        // exist! If it doesn't then we register and print an error.
        if path.exists() {
            if path.is_dir() {
                // Links to directories show as directory listings when viewing
                // the docs offline so it's best to avoid them.
                *errors = true;
                let pretty_path = path.strip_prefix(root).unwrap_or(&path);
                println!(
                    "{}:{}: directory link - {}",
                    pretty_file.display(),
                    i + 1,
                    pretty_path.display()
                );
                return;
            }
            if let Some(extension) = path.extension() {
                // Ignore none HTML files.
                if extension != "html" {
                    return;
                }
            }
            let res = load_file(cache, root, &path, FromRedirect(false));
            let (pretty_path, contents) = match res {
                Ok(res) => res,
                Err(LoadError::IOError(err)) => {
                    panic!("error loading {}: {}", path.display(), err);
                }
                Err(LoadError::BrokenRedirect(target, _)) => {
                    *errors = true;
                    println!(
                        "{}:{}: broken redirect to {}",
                        pretty_file.display(),
                        i + 1,
                        target.display()
                    );
                    return;
                }
                Err(LoadError::IsRedirect) => unreachable!(),
            };

            if let Some(ref fragment) = fragment {
                // Fragments like `#1-6` are most likely line numbers to be
                // interpreted by javascript, so we're ignoring these
                if fragment.splitn(2, '-').all(|f| f.chars().all(|c| c.is_numeric())) {
                    return;
                }

                // These appear to be broken in mdbook right now?
                if fragment.starts_with("-") {
                    return;
                }

                let entry = &mut cache.get_mut(&pretty_path).unwrap();
                entry.parse_ids(&pretty_path, &contents, errors);

                if !entry.ids.contains(*fragment) {
                    *errors = true;
                    print!("{}:{}: broken link fragment ", pretty_file.display(), i + 1);
                    println!("`#{}` pointing to `{}`", fragment, pretty_path.display());
                };
            }
        } else {
            *errors = true;
            print!("{}:{}: broken link - ", pretty_file.display(), i + 1);
            let pretty_path = path.strip_prefix(root).unwrap_or(&path);
            println!("{}", pretty_path.display());
        }
    });
    Some(pretty_file)
}

fn load_file(
    cache: &mut Cache,
    root: &Path,
    file: &Path,
    redirect: Redirect,
) -> Result<(PathBuf, Rc<String>), LoadError> {
    let pretty_file = PathBuf::from(file.strip_prefix(root).unwrap_or(&file));

    let (maybe_redirect, contents) = match cache.entry(pretty_file.clone()) {
        Entry::Occupied(entry) => (None, entry.get().source.clone()),
        Entry::Vacant(entry) => {
            let contents = match fs::read_to_string(file) {
                Ok(s) => Rc::new(s),
                Err(err) => {
                    return Err(if let FromRedirect(true) = redirect {
                        LoadError::BrokenRedirect(file.to_path_buf(), err)
                    } else {
                        LoadError::IOError(err)
                    });
                }
            };

            let maybe = maybe_redirect(&contents);
            if maybe.is_some() {
                if let SkipRedirect = redirect {
                    return Err(LoadError::IsRedirect);
                }
            } else {
                entry.insert(FileEntry { source: contents.clone(), ids: HashSet::new() });
            }
            (maybe, contents)
        }
    };
    match maybe_redirect.map(|url| file.parent().unwrap().join(url)) {
        Some(redirect_file) => load_file(cache, root, &redirect_file, FromRedirect(true)),
        None => Ok((pretty_file, contents)),
    }
}

fn maybe_redirect(source: &str) -> Option<String> {
    const REDIRECT: &'static str = "
Redirecting to <a href=";

    let mut lines = source.lines();
    let redirect_line = lines.nth(6)?;

    redirect_line.find(REDIRECT).map(|i| {
        let rest = &redirect_line[(i + REDIRECT.len() + 1)..];
        let pos_quote = rest.find('"').unwrap();
        rest[..pos_quote].to_owned()
    })
}

fn with_attrs_in_source<F: FnMut(&str, usize, &str)>(contents: &str, attr: &str, mut f: F) {
    let mut base = "";
    for (i, mut line) in contents.lines().enumerate() {
        while let Some(j) = line.find(attr) {
            let rest = &line[j + attr.len()..];
            // The base tag should always be the first link in the document so
            // we can get away with using one pass.
            let is_base = line[..j].ends_with("<base");
            line = rest;
            let pos_equals = match rest.find("=") {
                Some(i) => i,
                None => continue,
            };
            if rest[..pos_equals].trim_start_matches(" ") != "" {
                continue;
            }

            let rest = &rest[pos_equals + 1..];

            let pos_quote = match rest.find(&['"', '\''][..]) {
                Some(i) => i,
                None => continue,
            };
            let quote_delim = rest.as_bytes()[pos_quote] as char;

            if rest[..pos_quote].trim_start_matches(" ") != "" {
                continue;
            }
            let rest = &rest[pos_quote + 1..];
            let url = match rest.find(quote_delim) {
                Some(i) => &rest[..i],
                None => continue,
            };
            if is_base {
                base = url;
                continue;
            }
            f(url, i, base)
        }
    }
}
<|endoftext|>


<commit_before>mod common;

use common::util::TestScenario;

#[test]
#[cfg(feature = "ls")]
fn execution_phrase_double() {
    use std::process::Command;

    let scenario = TestScenario::new("ls");
    let output = Command::new(&scenario.bin_path)
        .arg("ls")
        .arg("--some-invalid-arg")
        .output()
        .unwrap();
    assert!(String::from_utf8(output.stderr)
        .unwrap()
        .contains(&format!("USAGE:\n    {} ls", scenario.bin_path.display(),)));
}

#[test]
#[cfg(feature = "ls")]
fn execution_phrase_single() {
    use std::process::Command;

    let scenario = TestScenario::new("ls");
    std::fs::copy(scenario.bin_path, scenario.fixtures.plus("uu-ls")).unwrap();
    let output = Command::new(scenario.fixtures.plus("uu-ls"))
        .arg("--some-invalid-arg")
        .output()
        .unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains(&format!(
        "USAGE:\n    {}",
        scenario.fixtures.plus("uu-ls").display()
    )));
}

#[test]
#[cfg(feature = "sort")]
fn util_name_double() {
    use std::{
        io::Write,
        process::{Command, Stdio},
    };

    let scenario = TestScenario::new("sort");
    let mut child = Command::new(&scenario.bin_path)
        .arg("sort")
        .stdin(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    // input invalid utf8 to cause an error
    child.stdin.take().unwrap().write_all(&[255]).unwrap();
    let output = child.wait_with_output().unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains("sort: "));
}

#[test]
#[cfg(feature = "sort")]
fn util_name_single() {
    use std::{
        io::Write,
        process::{Command, Stdio},
    };

    let scenario = TestScenario::new("sort");
    std::fs::copy(scenario.bin_path, scenario.fixtures.plus("uu-sort")).unwrap();
    let mut child = Command::new(scenario.fixtures.plus("uu-sort"))
        .stdin(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    // input invalid utf8 to cause an error
    child.stdin.take().unwrap().write_all(&[255]).unwrap();
    let output = child.wait_with_output().unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains(&format!(
        "{}: ",
        scenario.fixtures.plus("uu-sort").display()
    )));
}
<commit_msg>test_util_name: symlink the xecutable instead of copying<commit_after>mod common;

use common::util::TestScenario;

#[cfg(unix)]
use std::os::unix::fs::symlink as symlink_file;
#[cfg(windows)]
use std::os::windows::fs::symlink_file;

#[test]
#[cfg(feature = "ls")]
fn execution_phrase_double() {
    use std::process::Command;

    let scenario = TestScenario::new("ls");
    let output = Command::new(&scenario.bin_path)
        .arg("ls")
        .arg("--some-invalid-arg")
        .output()
        .unwrap();
    assert!(String::from_utf8(output.stderr)
        .unwrap()
        .contains(&format!("USAGE:\n    {} ls", scenario.bin_path.display(),)));
}

#[test]
#[cfg(feature = "ls")]
fn execution_phrase_single() {
    use std::process::Command;

    let scenario = TestScenario::new("ls");
    symlink_file(scenario.bin_path, scenario.fixtures.plus("uu-ls")).unwrap();
    let output = Command::new(scenario.fixtures.plus("uu-ls"))
        .arg("--some-invalid-arg")
        .output()
        .unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains(&format!(
        "USAGE:\n    {}",
        scenario.fixtures.plus("uu-ls").display()
    )));
}

#[test]
#[cfg(feature = "sort")]
fn util_name_double() {
    use std::{
        io::Write,
        process::{Command, Stdio},
    };

    let scenario = TestScenario::new("sort");
    let mut child = Command::new(&scenario.bin_path)
        .arg("sort")
        .stdin(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    // input invalid utf8 to cause an error
    child.stdin.take().unwrap().write_all(&[255]).unwrap();
    let output = child.wait_with_output().unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains("sort: "));
}

#[test]
#[cfg(feature = "sort")]
fn util_name_single() {
    use std::{
        io::Write,
        process::{Command, Stdio},
    };

    let scenario = TestScenario::new("sort");
    symlink_file(scenario.bin_path, scenario.fixtures.plus("uu-sort")).unwrap();
    let mut child = Command::new(scenario.fixtures.plus("uu-sort"))
        .stdin(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    // input invalid utf8 to cause an error
    child.stdin.take().unwrap().write_all(&[255]).unwrap();
    let output = child.wait_with_output().unwrap();
    assert!(String::from_utf8(output.stderr).unwrap().contains(&format!(
        "{}: ",
        scenario.fixtures.plus("uu-sort").display()
    )));
}
<|endoftext|>


<commit_before>use std::mem;
use std::cmp::max;
use std::fmt::{Debug, Formatter};
use delete_bulk::{DeleteBulk, TraversalDriver};

pub trait Item: Sized+Clone+Debug {
    type Key: Ord+Debug;

    fn ord(&self) -> Self::Key;
}

//pub type Item = Sized+Ord;

#[derive(Debug, Clone)]
pub struct Node<T: Item> {
    pub item: Option<T>,    // TODO we can remove the option and use height==0 as null indicator
    pub height: u32,
}

#[derive(Clone)]
pub struct ImplicitTree<T: Item> {
    data: Vec<Node<T>>,
    size: usize,
}

impl<T: Item> ImplicitTree<T> {
    pub fn new(sorted: Vec<T>) -> ImplicitTree<T> {
        let size = sorted.len();

        let capacity = Self::level_from(size)*4 + 3;

        let mut data = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            data.push(Node{item: None, height: 0});
        }

        let mut sorted: Vec<Option<T>> = sorted.into_iter().map(|x| Some(x)).collect();
        Self::build(&mut sorted, 0, &mut data);
        ImplicitTree { data: data, size: size }
    }

    pub fn with_nodes(nodes: Vec<Node<T>>) -> ImplicitTree<T> {
        let size = nodes.iter().filter(|x| x.height != 0).count();
        let capacity = Self::level_from(size)*4 + 3; // allocate enough nodes that righti() is never out of bounds

        let mut data = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            data.push(Node{item: None, height: 0});
        }

        unsafe {
            use ::std::ptr;
            ptr::copy_nonoverlapping(nodes.as_ptr(), data.as_mut_ptr(), nodes.len());
        }
        ::std::mem::forget(nodes);

        ImplicitTree { data: data, size: size }
    }

    fn build(sorted: &mut [Option<T>], idx: usize, data: &mut [Node<T>]) {
        match sorted.len() {
            0 => {}
            n => {
                let mid = n/2;
                let (lefti, righti) = (Self::lefti(idx), Self::righti(idx));
                Self::build(&mut sorted[..mid], lefti, data);
                Self::build(&mut sorted[mid+1..], righti, data);

                let height = 1 + max(data[lefti].height, data[righti].height);
                data[idx] = Node { item: sorted[mid].take(), height: height };
            }
        }
    }

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

    pub fn delete_bulk<D: TraversalDriver<T>>(&mut self, drv: &mut D) -> Vec<T> {
        let output = {
            let mut d = DeleteBulk::new(self);
            d.delete_bulk(drv);
            d.output
        };
        self.size -= output.len();
        output
    }

    fn delete_idx(&mut self, mut idx: usize) -> T {
        let removed = self.delete_idx_recursive(idx);
        // update the parents
        while idx != 0 {
            idx = Self::parenti(idx);
            self.update_height(idx);
        }
        self.size -= 1;

        removed
    }

    pub fn node(&self, idx: usize) -> &Node<T> {
        &self.data[idx]
    }

    pub fn node_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[idx]
    }

    pub fn item(&self, idx: usize) -> &T {
        self.node(idx).item.as_ref().unwrap()
    }

    pub fn item_mut(&mut self, idx: usize) -> &mut T {
        self.node_mut(idx).item.as_mut().unwrap()
    }

    fn delete_idx_recursive(&mut self, idx: usize) -> T {
        assert!(!self.is_null(idx));

        if !self.has_left(idx) && !self.has_right(idx) {
            //            if idx != 0 {
            //                let parent = self.parent_mut(idx);
            //                parent.has_child[Self::branch(idx)] = false;
            //            }
            let root = self.node_mut(idx);
            root.height = 0;
            root.item.take().unwrap()
        } else {
            let removed = if self.has_left(idx) && !self.has_right(idx) {
                let left_max = self.delete_max(Self::lefti(idx));
                mem::replace(self.item_mut(idx), left_max)
            } else if !self.has_left(idx) && self.has_right(idx) {
                let right_min = self.delete_min(Self::righti(idx));
                mem::replace(self.item_mut(idx), right_min)
            } else { // self.has_left(idx) && self.has_right(idx)
                // TODO: remove from the subtree with bigger height, not always from the left
                let left_max = self.delete_max(Self::lefti(idx));
                mem::replace(self.item_mut(idx), left_max)
            };

            self.update_height(idx);
            removed
        }
    }

    #[inline]
    pub fn update_height(&mut self, idx: usize) {
        let h = max(self.left(idx).height, self.right(idx).height) + 1;
        let node = self.node_mut(idx);
        assert!(node.item.is_some());
        node.height =  h;
    }

    fn delete_max(&mut self, idx: usize) -> T {
        // TODO: rewrite with loop
        if self.has_right(idx) {
            let removed = self.delete_max(Self::righti(idx));
            self.update_height(idx);
            removed
        } else {
            // this is the max, now just need to handle the left subtree
            self.delete_idx_recursive(idx)
        }
    }

    fn delete_min(&mut self, idx: usize) -> T {
        // TODO: rewrite with loop
        if self.has_left(idx) {
            let removed = self.delete_min(Self::lefti(idx));
            self.update_height(idx);
            removed
        } else {
            // this is the min, now just need to handle the right subtree
            self.delete_idx_recursive(idx)
        }
    }

    //    #[inline]
    //    fn levels_count(&self) -> usize {
    //        if self.data.is_empty() {
    //            0
    //        } else {
    //            Self::level_of(self.data.len()-1) + 1
    //        }
    //    }

    #[inline]
    fn level_from(level: usize) -> usize {
        (1 << level) - 1
    }

    #[inline]
    fn level_of(idx: usize) -> usize {
        mem::size_of::<usize>()*8 - ((idx+1).leading_zeros() as usize) - 1
    }

    #[inline]
    fn row_start(idx: usize) -> usize {
        Self::level_from(Self::level_of(idx))
    }

    #[inline]
    pub fn parenti(idx: usize) -> usize {
        (idx-1) >> 1
    }

    #[inline]
    pub fn lefti(idx: usize) -> usize {
        (idx<<1) + 1
    }

    #[inline]
    pub fn righti(idx: usize) -> usize {
        (idx<<1) + 2
    }

    #[inline]
    pub fn parent(&self, idx: usize) -> &Node<T> {
        &self.data[Self::parenti(idx)]
    }

    #[inline]
    pub fn left(&self, idx: usize) -> &Node<T> {
        &self.data[Self::lefti(idx)]
    }

    #[inline]
    pub fn right(&self, idx: usize) -> &Node<T> {
        &self.data[Self::righti(idx)]
    }

    #[inline]
    pub fn parent_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::parenti(idx)]
    }

    #[inline]
    pub fn left_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::lefti(idx)]
    }

    #[inline]
    pub fn right_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::righti(idx)]
    }

    #[inline]
    pub fn has_left(&self, idx: usize) -> bool {
        self.left(idx).height != 0
    }

    #[inline]
    pub fn has_right(&self, idx: usize) -> bool {
        self.right(idx).height != 0
    }

    #[inline]
    pub fn is_null(&self, idx: usize) -> bool {
        self.data[idx].item.is_none()
    }

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

impl<T: Item> Debug for ImplicitTree<T> {
    fn fmt(&self, fmt: &mut Formatter) -> ::std::fmt::Result {
        let mut nz: Vec<_> = self.data.iter()
            .rev()
            .skip_while(|node| node.item.is_none() && node.height==0)
            .map(|node| match node.item {
                None => (String::from("0"), 0),
                Some(ref x) => (format!("{:?}", x.ord()), node.height)
            })
            .collect();
        nz.reverse();

        let _ = write!(fmt, "[");
        let mut sep = "";
        for &(ref key, ref height) in nz.iter() {
            let _ = write!(fmt, "{}", sep);
            sep = ", ";
            let _ = write!(fmt, "({}, {})", key, height);
        }
        let _ = write!(fmt, "]");
        Ok(())
    }
}
<commit_msg>fix capacity allocation (was exponential in input size)<commit_after>use std::mem;
use std::cmp::max;
use std::fmt::{Debug, Formatter};
use delete_bulk::{DeleteBulk, TraversalDriver};

pub trait Item: Sized+Clone+Debug {
    type Key: Ord+Debug;

    fn ord(&self) -> Self::Key;
}

//pub type Item = Sized+Ord;

#[derive(Debug, Clone)]
pub struct Node<T: Item> {
    pub item: Option<T>,    // TODO we can remove the option and use height==0 as null indicator
    pub height: u32,
}

#[derive(Clone)]
pub struct ImplicitTree<T: Item> {
    data: Vec<Node<T>>,
    size: usize,
}

impl<T: Item> ImplicitTree<T> {
    pub fn new(sorted: Vec<T>) -> ImplicitTree<T> {
        let size = sorted.len();

        let capacity = Self::row_start(size)*4 + 3;

        let mut data = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            data.push(Node{item: None, height: 0});
        }

        let mut sorted: Vec<Option<T>> = sorted.into_iter().map(|x| Some(x)).collect();
        Self::build(&mut sorted, 0, &mut data);
        ImplicitTree { data: data, size: size }
    }

    pub fn with_nodes(nodes: Vec<Node<T>>) -> ImplicitTree<T> {
        let size = nodes.iter().filter(|x| x.height != 0).count();
        let capacity = Self::row_start(nodes.len())*4 + 3; // allocate enough nodes that righti() is never out of bounds

        let mut data = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            data.push(Node{item: None, height: 0});
        }

        unsafe {
            use ::std::ptr;
            ptr::copy_nonoverlapping(nodes.as_ptr(), data.as_mut_ptr(), nodes.len());
        }
        ::std::mem::forget(nodes);

        ImplicitTree { data: data, size: size }
    }

    fn build(sorted: &mut [Option<T>], idx: usize, data: &mut [Node<T>]) {
        match sorted.len() {
            0 => {}
            n => {
                let mid = n/2;
                let (lefti, righti) = (Self::lefti(idx), Self::righti(idx));
                Self::build(&mut sorted[..mid], lefti, data);
                Self::build(&mut sorted[mid+1..], righti, data);

                let height = 1 + max(data[lefti].height, data[righti].height);
                data[idx] = Node { item: sorted[mid].take(), height: height };
            }
        }
    }

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

    pub fn delete_bulk<D: TraversalDriver<T>>(&mut self, drv: &mut D) -> Vec<T> {
        let output = {
            let mut d = DeleteBulk::new(self);
            d.delete_bulk(drv);
            d.output
        };
        self.size -= output.len();
        output
    }

    fn delete_idx(&mut self, mut idx: usize) -> T {
        let removed = self.delete_idx_recursive(idx);
        // update the parents
        while idx != 0 {
            idx = Self::parenti(idx);
            self.update_height(idx);
        }
        self.size -= 1;

        removed
    }

    pub fn node(&self, idx: usize) -> &Node<T> {
        &self.data[idx]
    }

    pub fn node_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[idx]
    }

    pub fn item(&self, idx: usize) -> &T {
        self.node(idx).item.as_ref().unwrap()
    }

    pub fn item_mut(&mut self, idx: usize) -> &mut T {
        self.node_mut(idx).item.as_mut().unwrap()
    }

    fn delete_idx_recursive(&mut self, idx: usize) -> T {
        assert!(!self.is_null(idx));

        if !self.has_left(idx) && !self.has_right(idx) {
            //            if idx != 0 {
            //                let parent = self.parent_mut(idx);
            //                parent.has_child[Self::branch(idx)] = false;
            //            }
            let root = self.node_mut(idx);
            root.height = 0;
            root.item.take().unwrap()
        } else {
            let removed = if self.has_left(idx) && !self.has_right(idx) {
                let left_max = self.delete_max(Self::lefti(idx));
                mem::replace(self.item_mut(idx), left_max)
            } else if !self.has_left(idx) && self.has_right(idx) {
                let right_min = self.delete_min(Self::righti(idx));
                mem::replace(self.item_mut(idx), right_min)
            } else { // self.has_left(idx) && self.has_right(idx)
                // TODO: remove from the subtree with bigger height, not always from the left
                let left_max = self.delete_max(Self::lefti(idx));
                mem::replace(self.item_mut(idx), left_max)
            };

            self.update_height(idx);
            removed
        }
    }

    #[inline]
    pub fn update_height(&mut self, idx: usize) {
        let h = max(self.left(idx).height, self.right(idx).height) + 1;
        let node = self.node_mut(idx);
        assert!(node.item.is_some());
        node.height =  h;
    }

    fn delete_max(&mut self, idx: usize) -> T {
        // TODO: rewrite with loop
        if self.has_right(idx) {
            let removed = self.delete_max(Self::righti(idx));
            self.update_height(idx);
            removed
        } else {
            // this is the max, now just need to handle the left subtree
            self.delete_idx_recursive(idx)
        }
    }

    fn delete_min(&mut self, idx: usize) -> T {
        // TODO: rewrite with loop
        if self.has_left(idx) {
            let removed = self.delete_min(Self::lefti(idx));
            self.update_height(idx);
            removed
        } else {
            // this is the min, now just need to handle the right subtree
            self.delete_idx_recursive(idx)
        }
    }

    //    #[inline]
    //    fn levels_count(&self) -> usize {
    //        if self.data.is_empty() {
    //            0
    //        } else {
    //            Self::level_of(self.data.len()-1) + 1
    //        }
    //    }

    #[inline]
    fn level_from(level: usize) -> usize {
        (1 << level) - 1
    }

    #[inline]
    fn level_of(idx: usize) -> usize {
        mem::size_of::<usize>()*8 - ((idx+1).leading_zeros() as usize) - 1
    }

    #[inline]
    fn row_start(idx: usize) -> usize {
        Self::level_from(Self::level_of(idx))
    }

    #[inline]
    pub fn parenti(idx: usize) -> usize {
        (idx-1) >> 1
    }

    #[inline]
    pub fn lefti(idx: usize) -> usize {
        (idx<<1) + 1
    }

    #[inline]
    pub fn righti(idx: usize) -> usize {
        (idx<<1) + 2
    }

    #[inline]
    pub fn parent(&self, idx: usize) -> &Node<T> {
        &self.data[Self::parenti(idx)]
    }

    #[inline]
    pub fn left(&self, idx: usize) -> &Node<T> {
        &self.data[Self::lefti(idx)]
    }

    #[inline]
    pub fn right(&self, idx: usize) -> &Node<T> {
        &self.data[Self::righti(idx)]
    }

    #[inline]
    pub fn parent_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::parenti(idx)]
    }

    #[inline]
    pub fn left_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::lefti(idx)]
    }

    #[inline]
    pub fn right_mut(&mut self, idx: usize) -> &mut Node<T> {
        &mut self.data[Self::righti(idx)]
    }

    #[inline]
    pub fn has_left(&self, idx: usize) -> bool {
        self.left(idx).height != 0
    }

    #[inline]
    pub fn has_right(&self, idx: usize) -> bool {
        self.right(idx).height != 0
    }

    #[inline]
    pub fn is_null(&self, idx: usize) -> bool {
        self.data[idx].item.is_none()
    }

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

impl<T: Item> Debug for ImplicitTree<T> {
    fn fmt(&self, fmt: &mut Formatter) -> ::std::fmt::Result {
        let mut nz: Vec<_> = self.data.iter()
            .rev()
            .skip_while(|node| node.item.is_none() && node.height==0)
            .map(|node| match node.item {
                None => (String::from("0"), 0),
                Some(ref x) => (format!("{:?}", x.ord()), node.height)
            })
            .collect();
        nz.reverse();

        let _ = write!(fmt, "[");
        let mut sep = "";
        for &(ref key, ref height) in nz.iter() {
            let _ = write!(fmt, "{}", sep);
            sep = ", ";
            let _ = write!(fmt, "({}, {})", key, height);
        }
        let _ = write!(fmt, "]");
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! # C API for rav1e
//!
//! [rav1e](https://github.com/xiph/rav1e/) is an [AV1](https://aomediacodec.github.io/av1-spec/)
//! encoder written in [Rust](https://rust-lang.org)
//!
//! This is the C-compatible API

extern crate rav1e;
extern crate libc;

use std::slice;
use std::sync::Arc;

use std::ffi::CStr;
use std::ffi::CString;
use std::os::raw::c_char;
use std::os::raw::c_int;

use libc::size_t;
use libc::ptrdiff_t;

/// Raw video Frame
///
/// It can be allocated throught rav1e_frame_new(), populated using rav1e_frame_fill_plane()
/// and freed using rav1e_frame_unref().
pub struct Frame(Arc<rav1e::Frame<u16>>);

#[repr(C)]
#[derive(Copy, Clone, Debug)]
pub enum EncoderStatus {
    /// Normal operation
    Success = 0,
    /// The encoder needs more data to produce an output Packet
    /// May be emitted by `Context::receive_packet`  when frame reordering is enabled.
    NeedMoreData,
    /// There are enough Frames queue
    /// May be emitted by `Context::send_frame` when the input queue is constrained
    EnoughData,
    /// The encoder already produced the number of frames requested
    /// May be emitted by `Context::receive_packet` after a flush request had been processed
    /// or the frame limit had been reached.
    LimitReached,
    /// Generic fatal error
    Failure = -1,
}

impl From<Option<rav1e::EncoderStatus>> for EncoderStatus {
    fn from(status: Option<rav1e::EncoderStatus>) -> Self {
        match status {
            None => EncoderStatus::Success,
            Some(s) => match s {
                rav1e::EncoderStatus::NeedMoreData => EncoderStatus::NeedMoreData,
                rav1e::EncoderStatus::EnoughData => EncoderStatus::EnoughData,
                rav1e::EncoderStatus::LimitReached => EncoderStatus::LimitReached,
                rav1e::EncoderStatus::Failure => EncoderStatus::Failure,
            }
        }
    }
}

/// Encoder configuration
///
/// Instantiate it using rav1e_config_default() and fine-tune it using
/// rav1e_config_parse().
///
/// Use rav1e_config_unref() to free its memory.
pub struct Config {
    cfg: rav1e::Config,
}

/// Encoder context
///
/// Contains the encoding state, it is created by rav1e_context_new() using an
/// Encoder configuration.
///
/// Use rav1e_context_unref() to free its memory.
pub struct Context {
    ctx: rav1e::Context<u16>,
    last_err: Option<rav1e::EncoderStatus>,
}

type FrameType = rav1e::FrameType;

/// Encoded Packet
///
/// The encoded packets are retrieved using rav1e_receive_packet().
///
/// Use rav1e_packet_unref() to free its memory.
#[repr(C)]
pub struct Packet {
    /// Encoded data buffer
    pub data: *const u8,
    /// Encoded data buffer size
    pub len: size_t,
    /// Frame sequence number
    pub number: u64,
    /// Frame type
    pub frame_type: FrameType,
}

type ChromaSamplePosition=rav1e::ChromaSamplePosition;
type ChromaSampling=rav1e::ChromaSampling;
type MatrixCoefficients=rav1e::MatrixCoefficients;
type ColorPrimaries=rav1e::ColorPrimaries;
type TransferCharacteristics=rav1e::TransferCharacteristics;
type Rational=rav1e::Rational;

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_default() -> *mut Config {
    let cfg = rav1e::Config {
        enc: rav1e::EncoderConfig::default(),
        threads: 0,
    };

    let c = Box::new(Config {
        cfg,
    });

    Box::into_raw(c)
}

/// Set color properties of the stream.
///
/// Supported values are defined by the enum types
/// RaMatrixCoefficients, RaColorPrimaries, and RaTransferCharacteristics
/// respectively.
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_color_description(cfg: *mut Config,
                                                            matrix: MatrixCoefficients,
                                                            primaries: ColorPrimaries,
                                                            transfer: TransferCharacteristics
) -> c_int {
    (*cfg).cfg.enc.color_description = Some(rav1e::ColorDescription {
        matrix_coefficients: matrix,
        color_primaries: primaries,
        transfer_characteristics: transfer,
    });

    if (*cfg).cfg.enc.color_description.is_some() { 0 } else { -1 }
}

/// Set the content light level information for HDR10 streams.
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_content_light(cfg: *mut Config,
                                                        max_content_light_level: u16,
                                                        max_frame_average_light_level: u16
) -> c_int {
    (*cfg).cfg.enc.content_light = Some(rav1e::ContentLight {
        max_content_light_level,
        max_frame_average_light_level,
    });

    if (*cfg).cfg.enc.content_light.is_some() { 0 } else { -1 }
}

/// Set the mastering display information for HDR10 streams.
///
/// primaries and white_point arguments are RaPoint, containing 0.16 fixed point values.
/// max_luminance is a 24.8 fixed point value.
/// min_luminance is a 18.14 fixed point value.
///
/// Returns a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_mastering_display(cfg: *mut Config,
                                                            primaries: [rav1e::Point; 3],
                                                            white_point: rav1e::Point,
                                                            max_luminance: u32,
                                                            min_luminance: u32
) -> c_int {
    (*cfg).cfg.enc.mastering_display = Some(rav1e::MasteringDisplay {
        primaries,
        white_point,
        max_luminance,
        min_luminance,
    });

    if (*cfg).cfg.enc.mastering_display.is_some() { 0 } else { -1 }
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_unref(cfg: *mut Config) {
    if !cfg.is_null() {
        let _ = Box::from_raw(cfg);
    }
}

unsafe fn option_match(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char
) -> Result<(), ()> {
    let key = CStr::from_ptr(key).to_str().map_err(|_| ())?;
    let value = CStr::from_ptr(value).to_str().map_err(|_| ())?;
    let enc = &mut(*cfg).cfg.enc;

    match key {
        "width" => enc.width = value.parse().map_err(|_| ())?,
        "height" => enc.height = value.parse().map_err(|_| ())?,
        "speed" => enc.speed_settings = rav1e::SpeedSettings::from_preset(value.parse().map_err(|_| ())?),

        "threads" => (*cfg).cfg.threads = value.parse().map_err(|_| ())?,

        "tile_rows_log2" => enc.tile_rows_log2 = value.parse().map_err(|_| ())?,
        "tile_cols_log2" => enc.tile_cols_log2 = value.parse().map_err(|_| ())?,

        "tune" => enc.tune = value.parse().map_err(|_| ())?,
        "quantizer" => enc.quantizer = value.parse().map_err(|_| ())?,
        "bitrate" => enc.bitrate = value.parse().map_err(|_| ())?,

        "key_frame_interval" => enc.max_key_frame_interval = value.parse().map_err(|_| ())?,
        "min_key_frame_interval" => enc.min_key_frame_interval = value.parse().map_err(|_| ())?,
        "low_latency" => enc.low_latency = value.parse().map_err(|_| ())?,

        _ => return Err(())
    }

    Ok(())
}

/// Set a configuration parameter using its key and value as string.
///
/// Available keys and values
/// - "quantizer": 0-255, default 100
/// - "speed": 0-10, default 3
/// - "tune": "psnr"-"psychovisual", default "psnr"
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char,
) -> c_int {
    if option_match(cfg, key, value) == Ok(()) { 0 } else { -1 }
}

/// Set a configuration parameter using its key and value as integer.
///
/// Available keys and values are the same as rav1e_config_parse()
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse_int(
    cfg: *mut Config,
    key: *const c_char,
    value: c_int,
) -> c_int {
    let val = CString::new(value.to_string()).unwrap();
    if option_match(cfg, key, val.as_ptr()) == Ok(()) { 0 } else { -1 }
}

/// Generate a new encoding context from a populated encoder configuration
///
/// Multiple contexts can be generated through it.
#[no_mangle]
pub unsafe extern "C" fn rav1e_context_new(cfg: *const Config) -> *mut Context {
    let ctx = Context {
        ctx: (*cfg).cfg.new_context(),
        last_err: None,
    };

    Box::into_raw(Box::new(ctx))
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_context_unref(ctx: *mut Context) {
    if !ctx.is_null() {
        let _ = Box::from_raw(ctx);
    }
}

/// Produce a new frame from the encoding context
///
/// It must be populated using rav1e_frame_fill_plane().
///
/// The frame is reference counted and must be released passing it to rav1e_frame_unref(),
/// see rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_new(ctx: *const Context) -> *mut Frame {
    let f = (*ctx).ctx.new_frame();
    let frame = Box::new(Frame(f));

    Box::into_raw(frame)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_unref(frame: *mut Frame) {
    if !frame.is_null() {
        let _ = Box::from_raw(frame);
    }
}

/// Send the frame for encoding
///
/// The function increases the frame internal reference count and it can be passed multiple
/// times to different rav1e_send_frame().
///
/// Returns:
/// - `0` on success,
/// - `> 0` if the input queue is full
/// - `< 0` on unrecoverable failure
#[no_mangle]
pub unsafe extern "C" fn rav1e_send_frame(ctx: *mut Context, frame: *const Frame) -> EncoderStatus {
    let frame = if frame.is_null() {
        None
    } else {
        Some((*frame).0.clone())
    };

    let ret = (*ctx)
        .ctx
        .send_frame(frame)
        .map(|_v| {
            None
        }).unwrap_or_else(|e| {
            Some(e)
        });

    (*ctx).last_err = ret;

    ret.into()
}

/// Return the last encoder status
#[no_mangle]
pub unsafe extern "C" fn rav1e_last_status(ctx: *const Context) -> EncoderStatus {
    (*ctx).last_err.into()
}

/// Return a string matching the EncooderStatus variant.
#[no_mangle]
pub unsafe extern "C" fn rav1e_status_to_str(status: EncoderStatus) -> *mut c_char {
    let status = format!("{:?}", status);
    let cptr = CString::new(status).unwrap().as_ptr();

    libc::strdup(cptr)
}

/// Receive encoded data
///
/// Returns:
/// - `0` on success
/// - `> 0` if additional frame data is required
/// - `< 0` on unrecoverable failure
#[no_mangle]
pub unsafe extern "C" fn rav1e_receive_packet(
    ctx: *mut Context,
    pkt: *mut *mut Packet,
) -> EncoderStatus {
    let ret = (*ctx)
        .ctx
        .receive_packet()
        .map(|p| {
            let rav1e::Packet { data, number, frame_type, .. } = p;
            let len  = data.len();
            let data = Box::into_raw(data.into_boxed_slice()) as *const u8;
            let packet = Packet {
                data,
                len,
                number,
                frame_type,
            };
            *pkt = Box::into_raw(Box::new(packet));
            None
        }).unwrap_or_else(|e| {
            Some(e)
        });

    (*ctx).last_err = ret;

    ret.into()
}

#[no_mangle]
pub unsafe extern fn rav1e_packet_unref(pkt: *mut Packet) {
    if !pkt.is_null() {
        let pkt = Box::from_raw(pkt);
        let _ = Box::from_raw(pkt.data as *mut u8);
    }
}

/// Produce a sequence header matching the current encoding context
///
/// Its format is compatible with the AV1 Matroska and ISOBMFF specification.
///
/// Use rav1e_container_sequence_header_unref() to free it.
#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header(ctx: *mut Context, buf_size: *mut size_t) -> *mut u8 {
    let buf = (*ctx).ctx.container_sequence_header();

    *buf_size = buf.len();
    Box::into_raw(buf.into_boxed_slice()) as *mut u8
}

#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header_unref(sequence: *mut u8) {
    if !sequence.is_null() {
        let _ = Box::from_raw(sequence);
    }
}

/// Fill a frame plane
///
/// Currently the frame contains 3 planes, the first is luminance followed by
/// chrominance.
///
/// The data is copied and this function has to be called for each plane.
///
/// frame: A frame provided by rav1e_frame_new()
/// plane: The index of the plane starting from 0
/// data: The data to be copied
/// data_len: Lenght of the buffer
/// stride: Plane line in bytes, including padding
/// bytewidth: Number of bytes per component, either 1 or 2
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_fill_plane(
    frame: *mut Frame,
    plane: c_int,
    data: *const u8,
    data_len: size_t,
    stride: ptrdiff_t,
    bytewidth: c_int,
) {
    let f = &mut (*frame).0;
    let input = Arc::make_mut(f);
    let data_slice = slice::from_raw_parts(data, data_len as usize);

    input.planes[plane as usize].copy_from_raw_u8(data_slice, stride as usize, bytewidth as usize);
}
<commit_msg>crav1e: api: Add Encoded state<commit_after>//! # C API for rav1e
//!
//! [rav1e](https://github.com/xiph/rav1e/) is an [AV1](https://aomediacodec.github.io/av1-spec/)
//! encoder written in [Rust](https://rust-lang.org)
//!
//! This is the C-compatible API

extern crate rav1e;
extern crate libc;

use std::slice;
use std::sync::Arc;

use std::ffi::CStr;
use std::ffi::CString;
use std::os::raw::c_char;
use std::os::raw::c_int;

use libc::size_t;
use libc::ptrdiff_t;

/// Raw video Frame
///
/// It can be allocated throught rav1e_frame_new(), populated using rav1e_frame_fill_plane()
/// and freed using rav1e_frame_unref().
pub struct Frame(Arc<rav1e::Frame<u16>>);

#[repr(C)]
#[derive(Copy, Clone, Debug)]
pub enum EncoderStatus {
    /// Normal operation
    Success = 0,
    /// The encoder needs more data to produce an output Packet
    /// May be emitted by `Context::receive_packet`  when frame reordering is enabled.
    NeedMoreData,
    /// There are enough Frames queue
    /// May be emitted by `Context::send_frame` when the input queue is constrained
    EnoughData,
    /// The encoder already produced the number of frames requested
    /// May be emitted by `Context::receive_packet` after a flush request had been processed
    /// or the frame limit had been reached.
    LimitReached,
    /// A Frame had been encoded but not emitted yet
    Encoded,
    /// Generic fatal error
    Failure = -1,
}

impl From<Option<rav1e::EncoderStatus>> for EncoderStatus {
    fn from(status: Option<rav1e::EncoderStatus>) -> Self {
        match status {
            None => EncoderStatus::Success,
            Some(s) => match s {
                rav1e::EncoderStatus::NeedMoreData => EncoderStatus::NeedMoreData,
                rav1e::EncoderStatus::EnoughData => EncoderStatus::EnoughData,
                rav1e::EncoderStatus::LimitReached => EncoderStatus::LimitReached,
                rav1e::EncoderStatus::Encoded => EncoderStatus::Encoded,
                rav1e::EncoderStatus::Failure => EncoderStatus::Failure,
            }
        }
    }
}

/// Encoder configuration
///
/// Instantiate it using rav1e_config_default() and fine-tune it using
/// rav1e_config_parse().
///
/// Use rav1e_config_unref() to free its memory.
pub struct Config {
    cfg: rav1e::Config,
}

/// Encoder context
///
/// Contains the encoding state, it is created by rav1e_context_new() using an
/// Encoder configuration.
///
/// Use rav1e_context_unref() to free its memory.
pub struct Context {
    ctx: rav1e::Context<u16>,
    last_err: Option<rav1e::EncoderStatus>,
}

type FrameType = rav1e::FrameType;

/// Encoded Packet
///
/// The encoded packets are retrieved using rav1e_receive_packet().
///
/// Use rav1e_packet_unref() to free its memory.
#[repr(C)]
pub struct Packet {
    /// Encoded data buffer
    pub data: *const u8,
    /// Encoded data buffer size
    pub len: size_t,
    /// Frame sequence number
    pub number: u64,
    /// Frame type
    pub frame_type: FrameType,
}

type ChromaSamplePosition=rav1e::ChromaSamplePosition;
type ChromaSampling=rav1e::ChromaSampling;
type MatrixCoefficients=rav1e::MatrixCoefficients;
type ColorPrimaries=rav1e::ColorPrimaries;
type TransferCharacteristics=rav1e::TransferCharacteristics;
type Rational=rav1e::Rational;

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_default() -> *mut Config {
    let cfg = rav1e::Config {
        enc: rav1e::EncoderConfig::default(),
        threads: 0,
    };

    let c = Box::new(Config {
        cfg,
    });

    Box::into_raw(c)
}

/// Set color properties of the stream.
///
/// Supported values are defined by the enum types
/// RaMatrixCoefficients, RaColorPrimaries, and RaTransferCharacteristics
/// respectively.
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_color_description(cfg: *mut Config,
                                                            matrix: MatrixCoefficients,
                                                            primaries: ColorPrimaries,
                                                            transfer: TransferCharacteristics
) -> c_int {
    (*cfg).cfg.enc.color_description = Some(rav1e::ColorDescription {
        matrix_coefficients: matrix,
        color_primaries: primaries,
        transfer_characteristics: transfer,
    });

    if (*cfg).cfg.enc.color_description.is_some() { 0 } else { -1 }
}

/// Set the content light level information for HDR10 streams.
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_content_light(cfg: *mut Config,
                                                        max_content_light_level: u16,
                                                        max_frame_average_light_level: u16
) -> c_int {
    (*cfg).cfg.enc.content_light = Some(rav1e::ContentLight {
        max_content_light_level,
        max_frame_average_light_level,
    });

    if (*cfg).cfg.enc.content_light.is_some() { 0 } else { -1 }
}

/// Set the mastering display information for HDR10 streams.
///
/// primaries and white_point arguments are RaPoint, containing 0.16 fixed point values.
/// max_luminance is a 24.8 fixed point value.
/// min_luminance is a 18.14 fixed point value.
///
/// Returns a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_mastering_display(cfg: *mut Config,
                                                            primaries: [rav1e::Point; 3],
                                                            white_point: rav1e::Point,
                                                            max_luminance: u32,
                                                            min_luminance: u32
) -> c_int {
    (*cfg).cfg.enc.mastering_display = Some(rav1e::MasteringDisplay {
        primaries,
        white_point,
        max_luminance,
        min_luminance,
    });

    if (*cfg).cfg.enc.mastering_display.is_some() { 0 } else { -1 }
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_unref(cfg: *mut Config) {
    if !cfg.is_null() {
        let _ = Box::from_raw(cfg);
    }
}

unsafe fn option_match(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char
) -> Result<(), ()> {
    let key = CStr::from_ptr(key).to_str().map_err(|_| ())?;
    let value = CStr::from_ptr(value).to_str().map_err(|_| ())?;
    let enc = &mut(*cfg).cfg.enc;

    match key {
        "width" => enc.width = value.parse().map_err(|_| ())?,
        "height" => enc.height = value.parse().map_err(|_| ())?,
        "speed" => enc.speed_settings = rav1e::SpeedSettings::from_preset(value.parse().map_err(|_| ())?),

        "threads" => (*cfg).cfg.threads = value.parse().map_err(|_| ())?,

        "tile_rows_log2" => enc.tile_rows_log2 = value.parse().map_err(|_| ())?,
        "tile_cols_log2" => enc.tile_cols_log2 = value.parse().map_err(|_| ())?,

        "tune" => enc.tune = value.parse().map_err(|_| ())?,
        "quantizer" => enc.quantizer = value.parse().map_err(|_| ())?,
        "bitrate" => enc.bitrate = value.parse().map_err(|_| ())?,

        "key_frame_interval" => enc.max_key_frame_interval = value.parse().map_err(|_| ())?,
        "min_key_frame_interval" => enc.min_key_frame_interval = value.parse().map_err(|_| ())?,
        "low_latency" => enc.low_latency = value.parse().map_err(|_| ())?,

        _ => return Err(())
    }

    Ok(())
}

/// Set a configuration parameter using its key and value as string.
///
/// Available keys and values
/// - "quantizer": 0-255, default 100
/// - "speed": 0-10, default 3
/// - "tune": "psnr"-"psychovisual", default "psnr"
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char,
) -> c_int {
    if option_match(cfg, key, value) == Ok(()) { 0 } else { -1 }
}

/// Set a configuration parameter using its key and value as integer.
///
/// Available keys and values are the same as rav1e_config_parse()
///
/// Return a negative value on error or 0.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse_int(
    cfg: *mut Config,
    key: *const c_char,
    value: c_int,
) -> c_int {
    let val = CString::new(value.to_string()).unwrap();
    if option_match(cfg, key, val.as_ptr()) == Ok(()) { 0 } else { -1 }
}

/// Generate a new encoding context from a populated encoder configuration
///
/// Multiple contexts can be generated through it.
#[no_mangle]
pub unsafe extern "C" fn rav1e_context_new(cfg: *const Config) -> *mut Context {
    let ctx = Context {
        ctx: (*cfg).cfg.new_context(),
        last_err: None,
    };

    Box::into_raw(Box::new(ctx))
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_context_unref(ctx: *mut Context) {
    if !ctx.is_null() {
        let _ = Box::from_raw(ctx);
    }
}

/// Produce a new frame from the encoding context
///
/// It must be populated using rav1e_frame_fill_plane().
///
/// The frame is reference counted and must be released passing it to rav1e_frame_unref(),
/// see rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_new(ctx: *const Context) -> *mut Frame {
    let f = (*ctx).ctx.new_frame();
    let frame = Box::new(Frame(f));

    Box::into_raw(frame)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_unref(frame: *mut Frame) {
    if !frame.is_null() {
        let _ = Box::from_raw(frame);
    }
}

/// Send the frame for encoding
///
/// The function increases the frame internal reference count and it can be passed multiple
/// times to different rav1e_send_frame().
///
/// Returns:
/// - `0` on success,
/// - `> 0` if the input queue is full
/// - `< 0` on unrecoverable failure
#[no_mangle]
pub unsafe extern "C" fn rav1e_send_frame(ctx: *mut Context, frame: *const Frame) -> EncoderStatus {
    let frame = if frame.is_null() {
        None
    } else {
        Some((*frame).0.clone())
    };

    let ret = (*ctx)
        .ctx
        .send_frame(frame)
        .map(|_v| {
            None
        }).unwrap_or_else(|e| {
            Some(e)
        });

    (*ctx).last_err = ret;

    ret.into()
}

/// Return the last encoder status
#[no_mangle]
pub unsafe extern "C" fn rav1e_last_status(ctx: *const Context) -> EncoderStatus {
    (*ctx).last_err.into()
}

/// Return a string matching the EncooderStatus variant.
#[no_mangle]
pub unsafe extern "C" fn rav1e_status_to_str(status: EncoderStatus) -> *mut c_char {
    let status = format!("{:?}", status);
    let cptr = CString::new(status).unwrap().as_ptr();

    libc::strdup(cptr)
}

/// Receive encoded data
///
/// Returns:
/// - `0` on success
/// - `> 0` if additional frame data is required
/// - `< 0` on unrecoverable failure
#[no_mangle]
pub unsafe extern "C" fn rav1e_receive_packet(
    ctx: *mut Context,
    pkt: *mut *mut Packet,
) -> EncoderStatus {
    let ret = (*ctx)
        .ctx
        .receive_packet()
        .map(|p| {
            let rav1e::Packet { data, number, frame_type, .. } = p;
            let len  = data.len();
            let data = Box::into_raw(data.into_boxed_slice()) as *const u8;
            let packet = Packet {
                data,
                len,
                number,
                frame_type,
            };
            *pkt = Box::into_raw(Box::new(packet));
            None
        }).unwrap_or_else(|e| {
            Some(e)
        });

    (*ctx).last_err = ret;

    ret.into()
}

#[no_mangle]
pub unsafe extern fn rav1e_packet_unref(pkt: *mut Packet) {
    if !pkt.is_null() {
        let pkt = Box::from_raw(pkt);
        let _ = Box::from_raw(pkt.data as *mut u8);
    }
}

/// Produce a sequence header matching the current encoding context
///
/// Its format is compatible with the AV1 Matroska and ISOBMFF specification.
///
/// Use rav1e_container_sequence_header_unref() to free it.
#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header(ctx: *mut Context, buf_size: *mut size_t) -> *mut u8 {
    let buf = (*ctx).ctx.container_sequence_header();

    *buf_size = buf.len();
    Box::into_raw(buf.into_boxed_slice()) as *mut u8
}

#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header_unref(sequence: *mut u8) {
    if !sequence.is_null() {
        let _ = Box::from_raw(sequence);
    }
}

/// Fill a frame plane
///
/// Currently the frame contains 3 planes, the first is luminance followed by
/// chrominance.
///
/// The data is copied and this function has to be called for each plane.
///
/// frame: A frame provided by rav1e_frame_new()
/// plane: The index of the plane starting from 0
/// data: The data to be copied
/// data_len: Lenght of the buffer
/// stride: Plane line in bytes, including padding
/// bytewidth: Number of bytes per component, either 1 or 2
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_fill_plane(
    frame: *mut Frame,
    plane: c_int,
    data: *const u8,
    data_len: size_t,
    stride: ptrdiff_t,
    bytewidth: c_int,
) {
    let f = &mut (*frame).0;
    let input = Arc::make_mut(f);
    let data_slice = slice::from_raw_parts(data, data_len as usize);

    input.planes[plane as usize].copy_from_raw_u8(data_slice, stride as usize, bytewidth as usize);
}
<|endoftext|>