<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_compare(&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::Compare(op, ref left, ref right) =>
                self.visit_compare(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.write(&format!("writer.write_str({:#?}).unwrap();",
                                    val));
            }
        } else if self.next_ws.is_some() {
        }
        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.write(&format!("writer.write_str({:#?}).unwrap();", lws));
            }
        }
        if !val.is_empty() {
            self.write(&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!("self.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, blocks: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} {}{} {{", anno, ast.ident.as_ref(),
                              anno));
        self.indent();
        self.handle(blocks);
        self.dedent();
        self.writeln("}");
    }

    fn template_impl(&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 trait_impl(&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 trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        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.writeln("self.render_trait_to(writer);");
        self.dedent();
        self.writeln("}");

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

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      block_names: &[&str], nodes: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

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

        self.writeln("fn render_trait_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 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() {
        gen.struct_impl(ast, &blocks);
        if base.is_none() {
            gen.template_trait(ast, path, &block_names, &content);
        }
        let tmpl_path = match base {
            Some(Expr::StrLit(base_path)) => { base_path },
            _ => { path },
        };
        gen.trait_impl(ast, tmpl_path, &block_names);
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<commit_msg>Improve formatting of generated code<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 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_compare(&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::Compare(op, ref left, ref right) =>
                self.visit_compare(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));
            }
        } else if self.next_ws.is_some() {
        }
        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!("self.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, blocks: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} {}{} {{", anno, ast.ident.as_ref(),
                              anno));
        self.indent();
        self.handle(blocks);
        self.dedent();
        self.writeln("}");
    }

    fn template_impl(&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 trait_impl(&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 trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        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.writeln("self.render_trait_to(writer);");
        self.dedent();
        self.writeln("}");

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

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      block_names: &[&str], nodes: &'a [Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

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

        self.writeln("fn render_trait_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 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() {
        gen.struct_impl(ast, &blocks);
        if base.is_none() {
            gen.template_trait(ast, path, &block_names, &content);
        }
        let tmpl_path = match base {
            Some(Expr::StrLit(base_path)) => { base_path },
            _ => { path },
        };
        gen.trait_impl(ast, tmpl_path, &block_names);
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<|endoftext|>


<commit_before>use field::PuyoPlainField;
use frame;
use rensa_result::RensaResult;

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
use field::BitField;

fn run_puyoplainfield_test(mut pf: PuyoPlainField, expected_result: &RensaResult) {
    let actual_result = pf.simulate();

    assert_eq!(actual_result.chain, expected_result.chain);
    assert_eq!(actual_result.score, expected_result.score);
    // Check only frames is not zero.
    if expected_result.frame > 0 {
        assert_eq!(actual_result.frame, expected_result.frame);
    }
    assert_eq!(actual_result.quick, expected_result.quick);
}

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
fn run_bitfield_test(mut bf: BitField, expected_result: &RensaResult) {
    let actual_result = bf.simulate();

    assert_eq!(actual_result.chain, expected_result.chain);
    assert_eq!(actual_result.score, expected_result.score);
    // Check only frames is not zero.
    if expected_result.frame > 0 {
        assert_eq!(actual_result.frame, expected_result.frame);
    }
    assert_eq!(actual_result.quick, expected_result.quick);
}

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
fn run_test(src: &str, expected_result: RensaResult) {
    run_puyoplainfield_test(PuyoPlainField::from_str(src), &expected_result);
    run_bitfield_test(BitField::from_str(src), &expected_result);
}

#[cfg(not(all(target_feature = "avx2", target_feature = "bmi2")))]
fn run_test(src: &str, expected_result: RensaResult) {
    run_puyoplainfield_test(PuyoPlainField::from_str(src), &expected_result);
}

#[test]
fn test_simulate_1rensa_quick() {
    let src = "..RRRR";

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case1() {
    let src = concat!(
        ".....Y",
        "..RRRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case2() {
    let src = concat!(
        ".....Y",
        ".....R",
        "...RRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case3() {
    let src = concat!(
        ".....Y",
        "....YR",
        "...RRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_quick() {
    let src = concat!(
        "Y.....",
        "RYY...",
        "RRRY..");

    let expected_result = RensaResult {
        chain: 2,
        score: 40 + 40 * 8,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_nonquick_case1() {
    let src = concat!(
        "YB....",
        "RYY...",
        "RRRY..");

    let expected_result = RensaResult {
        chain: 2,
        score: 40 + 40 * 8,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_nonquick_case2() {
    let src = concat!(
        "..B...",
        "..BBYB",
        "RRRRBB");

    let expected_result = RensaResult {
        chain: 2,
        score: 700,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_19rensa_case1() {
    let src = concat!(
        ".G.BRG",
        "GBRRYR",
        "RRYYBY",
        "RGYRBR",
        "YGYRBY",
        "YGBGYR",
        "GRBGYR",
        "BRBYBY",
        "RYYBYY",
        "BRBYBR",
        "BGBYRR",
        "YGBGBG",
        "RBGBGG");

    let expected_result = RensaResult {
        chain: 19,
        score: 175080,
        frame: 0,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate() {
    run_test(".B.GRBBGRRYRRRYYGYRBYRGRYBYRGYYBGBYRBRGBYRGRGYGYRYYGYYGRGYGRGBGYRRYBGBGBRGBGBB",
             RensaResult::new(19, 175080, 0, true));
    run_test("B..RYGGYGYGBGRRRBRGBRYBGRRGGYGYYGYRRYGRBRBRBBGYGRGGYRRRBGRGRYBYRBBRRYGGBRBBYRY",
             RensaResult::new(19, 175080, 0, true));
    run_test("...BB..B.RBB.RBRBO.RBGRB.GRGRB.GRYRB.YGYGR.YGYGR.BYBYG.BYBYGBOBOYGRRRROYBBBOBB",
             RensaResult::new(2, 38540, 0, true));
    run_test(".B.BB..R.RBB.GBRBO.GBGRB.YRGRB.YRYRB.YGYGR.BGYGRGRGBGRGGYBYGYGBOBYYRRROBBBBOBB",
             RensaResult::new(3, 43260, 0, true));
    run_test("...BB..R.RBB.GBRBOGGBGRBRYRGRBRYRYRBRYGYGRRBGYGRORGBGRGGYBYGYGBOBYYRRROBBBBOBB",
             RensaResult::new(4, 50140, 0, true));
    run_test("GRBBB.BGYRBBYYYRBOOGBGRBBYRGRBBYRGRBBYGYGROBGYGRGRGBGRGGYBYYYOBOBYYRRROBBBBOBB",
             RensaResult::new(5, 68700, 0, true));
    run_test("RRRROOOOROROROROOROOROOORORORRROOOORROROOOOORRORROOOORRROOOROOOORORRROROOOOORO",
             RensaResult::new(4, 4840, 0, true));
    run_test("BRBBRRBRRRBRRBRBRBRBRBRBRBRBRBBRBRBRRRBBRRBBRRBBRBRBRBBRBRBRBBRBRRRRBRBBRBBRRB",
             RensaResult::new(9, 49950, 0, true));
    run_test("RRRRRYBRRYOORRYOYRBYRRROBRYOYYBYBYOBRBRBBORRORRROOOOOOOOOOOOOOOOOOOOOOOOOOOOOO",
             RensaResult::new(9, 32760, 0, true));
    run_test("YYGBRGRYYBBBYYOYGGRGORGBRBORRGRYOYYYYYOBRGRBGRGGBBYRRYGGYBBBGRRYRYRGYRYYGRRBBB",
             RensaResult::new(18, 155980, 0, true));
    run_test("RRR.RRORRROROORORRROOROORORORRORORORRORORORROROROORROORROORRROROROORORORORRORR",
             RensaResult::new(11, 47080, 0, true));
    run_test("......RRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRR",
             RensaResult::new(1, 7200, 0, true));
}
<commit_msg>:white_check_mark: Add test for RealPlainField<commit_after>use field::{PuyoPlainField, RealPlainField};
use frame;
use rensa_result::RensaResult;

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
use field::BitField;

fn run_puyoplainfield_test(mut pf: PuyoPlainField, expected_result: &RensaResult) {
    let actual_result = pf.simulate();

    assert_eq!(actual_result.chain, expected_result.chain);
    assert_eq!(actual_result.score, expected_result.score);
    // Check only frames is not zero.
    if expected_result.frame > 0 {
        assert_eq!(actual_result.frame, expected_result.frame);
    }
    assert_eq!(actual_result.quick, expected_result.quick);
}

fn run_realplainfield_test(mut pf: RealPlainField, expected_result: &RensaResult) {
    let actual_result = pf.simulate();

    assert_eq!(actual_result.chain, expected_result.chain);
    assert_eq!(actual_result.score, expected_result.score);
    // Check only frames is not zero.
    if expected_result.frame > 0 {
        assert_eq!(actual_result.frame, expected_result.frame);
    }
    assert_eq!(actual_result.quick, expected_result.quick);
}

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
fn run_bitfield_test(mut bf: BitField, expected_result: &RensaResult) {
    let actual_result = bf.simulate();

    assert_eq!(actual_result.chain, expected_result.chain);
    assert_eq!(actual_result.score, expected_result.score);
    // Check only frames is not zero.
    if expected_result.frame > 0 {
        assert_eq!(actual_result.frame, expected_result.frame);
    }
    assert_eq!(actual_result.quick, expected_result.quick);
}

#[cfg(all(target_feature = "avx2", target_feature = "bmi2"))]
fn run_test(src: &str, expected_result: RensaResult) {
    run_puyoplainfield_test(PuyoPlainField::from_str(src), &expected_result);
    run_realplainfield_test(RealPlainField::from_str(src), &expected_result);
    run_bitfield_test(BitField::from_str(src), &expected_result);
}

#[cfg(not(all(target_feature = "avx2", target_feature = "bmi2")))]
fn run_test(src: &str, expected_result: RensaResult) {
    run_puyoplainfield_test(PuyoPlainField::from_str(src), &expected_result);
    run_realplainfield_test(RealPlainField::from_str(src), &expected_result);
}

#[test]
fn test_simulate_1rensa_quick() {
    let src = "..RRRR";

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case1() {
    let src = concat!(
        ".....Y",
        "..RRRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case2() {
    let src = concat!(
        ".....Y",
        ".....R",
        "...RRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_1rensa_nonquick_case3() {
    let src = concat!(
        ".....Y",
        "....YR",
        "...RRR");

    let expected_result = RensaResult {
        chain: 1,
        score: 40,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_quick() {
    let src = concat!(
        "Y.....",
        "RYY...",
        "RRRY..");

    let expected_result = RensaResult {
        chain: 2,
        score: 40 + 40 * 8,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_nonquick_case1() {
    let src = concat!(
        "YB....",
        "RYY...",
        "RRRY..");

    let expected_result = RensaResult {
        chain: 2,
        score: 40 + 40 * 8,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[2] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_2rensa_nonquick_case2() {
    let src = concat!(
        "..B...",
        "..BBYB",
        "RRRRBB");

    let expected_result = RensaResult {
        chain: 2,
        score: 700,
        frame: frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING +
            frame::FRAMES_VANISH_ANIMATION + frame::FRAMES_TO_DROP_FAST[1] + frame::FRAMES_GROUNDING,
        quick: false
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate_19rensa_case1() {
    let src = concat!(
        ".G.BRG",
        "GBRRYR",
        "RRYYBY",
        "RGYRBR",
        "YGYRBY",
        "YGBGYR",
        "GRBGYR",
        "BRBYBY",
        "RYYBYY",
        "BRBYBR",
        "BGBYRR",
        "YGBGBG",
        "RBGBGG");

    let expected_result = RensaResult {
        chain: 19,
        score: 175080,
        frame: 0,
        quick: true,
    };

    run_test(src, expected_result);
}

#[test]
fn test_simulate() {
    run_test(".B.GRBBGRRYRRRYYGYRBYRGRYBYRGYYBGBYRBRGBYRGRGYGYRYYGYYGRGYGRGBGYRRYBGBGBRGBGBB",
             RensaResult::new(19, 175080, 0, true));
    run_test("B..RYGGYGYGBGRRRBRGBRYBGRRGGYGYYGYRRYGRBRBRBBGYGRGGYRRRBGRGRYBYRBBRRYGGBRBBYRY",
             RensaResult::new(19, 175080, 0, true));
    run_test("...BB..B.RBB.RBRBO.RBGRB.GRGRB.GRYRB.YGYGR.YGYGR.BYBYG.BYBYGBOBOYGRRRROYBBBOBB",
             RensaResult::new(2, 38540, 0, true));
    run_test(".B.BB..R.RBB.GBRBO.GBGRB.YRGRB.YRYRB.YGYGR.BGYGRGRGBGRGGYBYGYGBOBYYRRROBBBBOBB",
             RensaResult::new(3, 43260, 0, true));
    run_test("...BB..R.RBB.GBRBOGGBGRBRYRGRBRYRYRBRYGYGRRBGYGRORGBGRGGYBYGYGBOBYYRRROBBBBOBB",
             RensaResult::new(4, 50140, 0, true));
    run_test("GRBBB.BGYRBBYYYRBOOGBGRBBYRGRBBYRGRBBYGYGROBGYGRGRGBGRGGYBYYYOBOBYYRRROBBBBOBB",
             RensaResult::new(5, 68700, 0, true));
    run_test("RRRROOOOROROROROOROOROOORORORRROOOORROROOOOORRORROOOORRROOOROOOORORRROROOOOORO",
             RensaResult::new(4, 4840, 0, true));
    run_test("BRBBRRBRRRBRRBRBRBRBRBRBRBRBRBBRBRBRRRBBRRBBRRBBRBRBRBBRBRBRBBRBRRRRBRBBRBBRRB",
             RensaResult::new(9, 49950, 0, true));
    run_test("RRRRRYBRRYOORRYOYRBYRRROBRYOYYBYBYOBRBRBBORRORRROOOOOOOOOOOOOOOOOOOOOOOOOOOOOO",
             RensaResult::new(9, 32760, 0, true));
    run_test("YYGBRGRYYBBBYYOYGGRGORGBRBORRGRYOYYYYYOBRGRBGRGGBBYRRYGGYBBBGRRYRYRGYRYYGRRBBB",
             RensaResult::new(18, 155980, 0, true));
    run_test("RRR.RRORRROROORORRROOROORORORRORORORRORORORROROROORROORROORRROROROORORORORRORR",
             RensaResult::new(11, 47080, 0, true));
    run_test("......RRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRR",
             RensaResult::new(1, 7200, 0, true));
}
<|endoftext|>


<commit_before>use core::hint;

/// Tock programs' entry point. Called by the kernel at program start. Sets up
/// the stack then calls rust_start() for the remainder of setup.
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub unsafe extern "C" fn _start(
    app_start: usize,
    mem_start: usize,
    _memory_len: usize,
    app_heap_break: usize,
) -> ! {
    asm!("
        // Because ROPI-RWPI support in LLVM/rustc is incomplete, Rust
        // applications must be statically linked. An offset between the
        // location the program is linked at and its actual location in flash
        // would cause references in .data and .rodata to point to the wrong
        // data. To mitigate this, this section checks that .text (and .start)
        // are loaded at the correct location. If the application was linked and
        // loaded correctly, the location of the first instruction (read using
        // the Program Counter) will match the intended location of .start. We
        // don't have an easy way to signal an error, so for now we just yield
        // if the location is wrong.
        mov r4, pc
        sub r4, r4, #4    // r4 = pc
        ldr r5, =.start   // r5 = address of .start
        cmp r4, r5
        beq .Lstack_init  // Jump to stack initialization if pc was correct
        movw r0, #8       // LowLevelDebug driver number
        movw r1, #1       // LowLevelDebug 'print status code' command
        movw r2, #2       // LowLevelDebug relocation failed status code
        svc 2             // command() syscall
        .Lyield_loop:
        svc 0             // yield() syscall (in infinite loop)
        b .Lyield_loop

        .Lstack_init:
        // Compute the stacktop (stack_start). The stacktop is computed as
        // stack_size + mem_start plus padding to align the stack to a multiple
        // of 8 bytes. The 8 byte alignment is to follow ARM AAPCS:
        // http://infocenter.arm.com/help/index.jsp?topic=/com.arm.doc.faqs/ka4127.html
        ldr r4, [r0, #36]  // r4 = app_start->stack_size
        add r4, r4, r1     // r4 = app_start->stack_size + mem_start
        add r4, #7         // r4 = app_start->stack_size + mem_start + 7
        bic r4, r4, #7     // r4 = (app_start->stack_size + mem_start + 7) & ~0x7
        mov sp, r4         // sp = r4

        // We need to pass app_start, stacktop and app_heap_break to rust_start.
        // Temporarily store them in r6, r7 and r8
        mov r6, r0
        mov r7, sp

        // Debug support, tell the kernel the stack location
        //
        // memop(10, stacktop)
        // r7 contains stacktop
        mov r0, #10
        mov r1, r7
        svc 4

        // Debug support, tell the kernel the heap_start location
        mov r0, r6
        ldr r4, [r0, #24] // r4 = app_start->bss_start
        ldr r5, [r0, #28] // r5 = app_start->bss_size
        add r4, r4, r5    // r4 = bss_start + bss_size
        //
        // memop(11, r4)
        mov r0, #11
        mov r1, r4
        svc 4

        // Store heap_start (and soon to be app_heap_break) in r8
        mov r8, r4

        // There is a possibility that stack + .data + .bss is greater than
        // 3072. Therefore setup the initial app_heap_break to heap_start (that
        // is zero initial heap) and let rust_start determine where the actual
        // app_heap_break should go.
        //
        // Also, because app_heap_break is where the unprivileged MPU region
        // ends, in case mem_start + stack + .data + .bss is greater than
        // initial app_heap_break (mem_start + 3072), we will get a memory fault
        // in rust_start when initializing .data and .bss. Setting
        // app_heap_break to heap_start avoids that.

        // memop(0, r8)
        mov r0, #0
        mov r1, r8
        svc 4

        // NOTE: If there is a hard-fault before this point, then
        //       process_detail_fmt in kernel/src/process.rs panics which
        //       will result in us losing the PC of the instruction
        //       generating the hard-fault. Therefore any code before
        //       this point is critical code

        // Setup parameters needed by rust_start
        // r6 (app_start), r7 (stacktop), r8 (app_heap_break)
        mov r0, r6
        mov r1, r7
        mov r2, r8

        // Call rust_start
        bl rust_start"
        :                                                              // No output operands
        : "{r0}"(app_start), "{r1}"(mem_start), "{r3}"(app_heap_break) // Input operands
        : "r0", "r1", "r2", "r3", "r4", "r5", "r6", "r8", "r12",
          "cc", "memory"                                               // Clobbers
        : "volatile"                                                   // Options
    );
    hint::unreachable_unchecked()
}
<commit_msg>imxrt1050 without start_item_arm fix<commit_after>use core::hint;

/// Tock programs' entry point. Called by the kernel at program start. Sets up
/// the stack then calls rust_start() for the remainder of setup.
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub unsafe extern "C" fn _start(
    app_start: usize,
    mem_start: usize,
    _memory_len: usize,
    app_heap_break: usize,
) -> ! {
    asm!("
        // Because ROPI-RWPI support in LLVM/rustc is incomplete, Rust
        // applications must be statically linked. An offset between the
        // location the program is linked at and its actual location in flash
        // would cause references in .data and .rodata to point to the wrong
        // data. To mitigate this, this section checks that .text (and .start)
        // are loaded at the correct location. If the application was linked and
        // loaded correctly, the location of the first instruction (read using
        // the Program Counter) will match the intended location of .start. We
        // don't have an easy way to signal an error, so for now we just yield
        // if the location is wrong.
        sub r4, pc, #4    // r4 = pc
        ldr r5, =.start   // r5 = address of .start
        cmp r4, r5
        beq .Lstack_init  // Jump to stack initialization if pc was correct
        movw r0, #8       // LowLevelDebug driver number
        movw r1, #1       // LowLevelDebug 'print status code' command
        movw r2, #2       // LowLevelDebug relocation failed status code
        svc 2             // command() syscall
        .Lyield_loop:
        svc 0             // yield() syscall (in infinite loop)
        b .Lyield_loop

        .Lstack_init:
        // Compute the stacktop (stack_start). The stacktop is computed as
        // stack_size + mem_start plus padding to align the stack to a multiple
        // of 8 bytes. The 8 byte alignment is to follow ARM AAPCS:
        // http://infocenter.arm.com/help/index.jsp?topic=/com.arm.doc.faqs/ka4127.html
        ldr r4, [r0, #36]  // r4 = app_start->stack_size
        add r4, r4, r1     // r4 = app_start->stack_size + mem_start
        add r4, #7         // r4 = app_start->stack_size + mem_start + 7
        bic r4, r4, #7     // r4 = (app_start->stack_size + mem_start + 7) & ~0x7
        mov sp, r4         // sp = r4

        // We need to pass app_start, stacktop and app_heap_break to rust_start.
        // Temporarily store them in r6, r7 and r8
        mov r6, r0
        mov r7, sp

        // Debug support, tell the kernel the stack location
        //
        // memop(10, stacktop)
        // r7 contains stacktop
        mov r0, #10
        mov r1, r7
        svc 4

        // Debug support, tell the kernel the heap_start location
        mov r0, r6
        ldr r4, [r0, #24] // r4 = app_start->bss_start
        ldr r5, [r0, #28] // r5 = app_start->bss_size
        add r4, r4, r5    // r4 = bss_start + bss_size
        //
        // memop(11, r4)
        mov r0, #11
        mov r1, r4
        svc 4

        // Store heap_start (and soon to be app_heap_break) in r8
        mov r8, r4

        // There is a possibility that stack + .data + .bss is greater than
        // 3072. Therefore setup the initial app_heap_break to heap_start (that
        // is zero initial heap) and let rust_start determine where the actual
        // app_heap_break should go.
        //
        // Also, because app_heap_break is where the unprivileged MPU region
        // ends, in case mem_start + stack + .data + .bss is greater than
        // initial app_heap_break (mem_start + 3072), we will get a memory fault
        // in rust_start when initializing .data and .bss. Setting
        // app_heap_break to heap_start avoids that.

        // memop(0, r8)
        mov r0, #0
        mov r1, r8
        svc 4

        // NOTE: If there is a hard-fault before this point, then
        //       process_detail_fmt in kernel/src/process.rs panics which
        //       will result in us losing the PC of the instruction
        //       generating the hard-fault. Therefore any code before
        //       this point is critical code

        // Setup parameters needed by rust_start
        // r6 (app_start), r7 (stacktop), r8 (app_heap_break)
        mov r0, r6
        mov r1, r7
        mov r2, r8

        // Call rust_start
        bl rust_start"
        :                                                              // No output operands
        : "{r0}"(app_start), "{r1}"(mem_start), "{r3}"(app_heap_break) // Input operands
        : "r0", "r1", "r2", "r3", "r4", "r5", "r6", "r8", "r12",
          "cc", "memory"                                               // Clobbers
        : "volatile"                                                   // Options
    );
    hint::unreachable_unchecked()
}
<|endoftext|>


<commit_before>//! This modules contains various "assits": suggestions for source code edits
//! which are likely to occur at a given cursor positon. For example, if the
//! cursor is on the `,`, a possible assist is swapping the elments around the
//! comma.

mod flip_comma;
mod add_derive;
mod add_impl;
mod introduce_variable;
mod change_visibility;
mod split_import;
mod replace_if_let_with_match;

use ra_text_edit::{TextEdit, TextEditBuilder};
use ra_syntax::{
    Direction, SyntaxNode, TextUnit, TextRange, SourceFile, AstNode,
    algo::{find_leaf_at_offset, find_node_at_offset, find_covering_node, LeafAtOffset},
};
use itertools::Itertools;

use crate::formatting::leading_indent;

pub use self::{
    flip_comma::flip_comma,
    add_derive::add_derive,
    add_impl::add_impl,
    introduce_variable::introduce_variable,
    change_visibility::change_visibility,
    split_import::split_import,
    replace_if_let_with_match::replace_if_let_with_match,
};

/// Return all the assists applicable at the given position.
pub fn assists(file: &SourceFile, range: TextRange) -> Vec<LocalEdit> {
    let ctx = AssistCtx::new(file, range);
    [
        flip_comma,
        add_derive,
        add_impl,
        introduce_variable,
        change_visibility,
        split_import,
        replace_if_let_with_match,
    ]
    .iter()
    .filter_map(|&assist| ctx.clone().apply(assist))
    .collect()
}

#[derive(Debug)]
pub struct LocalEdit {
    pub label: String,
    pub edit: TextEdit,
    pub cursor_position: Option<TextUnit>,
}

fn non_trivia_sibling(node: &SyntaxNode, direction: Direction) -> Option<&SyntaxNode> {
    node.siblings(direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

/// `AssistCtx` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat overengeneered approach, given the current needs. The
/// assists workflow consists of two phases. In the first phase, a user asks for
/// the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then neccessary to answer
///   "is assist applicable" in the first phase.
/// * second, when we are appling assist, we don't have a gurantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach remeniscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `should_compute_edit`. It first check if an edit is applicable (potentially
/// computing info required to compute the actual edit). If it is applicable,
/// and `should_compute_edit` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `should_compute_edit = false`, and then applying the selected edit
/// again, with `should_compute_edit = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagarly :-)
#[derive(Debug, Clone)]
pub struct AssistCtx<'a> {
    source_file: &'a SourceFile,
    range: TextRange,
    should_compute_edit: bool,
}

#[derive(Debug)]
pub enum Assist {
    Applicable,
    Edit(LocalEdit),
}

#[derive(Default)]
struct AssistBuilder {
    edit: TextEditBuilder,
    cursor_position: Option<TextUnit>,
}

impl<'a> AssistCtx<'a> {
    pub fn new(source_file: &'a SourceFile, range: TextRange) -> AssistCtx {
        AssistCtx {
            source_file,
            range,
            should_compute_edit: false,
        }
    }

    pub fn apply(mut self, assist: fn(AssistCtx) -> Option<Assist>) -> Option<LocalEdit> {
        self.should_compute_edit = true;
        match assist(self) {
            None => None,
            Some(Assist::Edit(e)) => Some(e),
            Some(Assist::Applicable) => unreachable!(),
        }
    }

    pub fn check(mut self, assist: fn(AssistCtx) -> Option<Assist>) -> bool {
        self.should_compute_edit = false;
        match assist(self) {
            None => false,
            Some(Assist::Edit(_)) => unreachable!(),
            Some(Assist::Applicable) => true,
        }
    }

    fn build(self, label: impl Into<String>, f: impl FnOnce(&mut AssistBuilder)) -> Option<Assist> {
        if !self.should_compute_edit {
            return Some(Assist::Applicable);
        }
        let mut edit = AssistBuilder::default();
        f(&mut edit);
        Some(Assist::Edit(LocalEdit {
            label: label.into(),
            edit: edit.edit.finish(),
            cursor_position: edit.cursor_position,
        }))
    }

    pub(crate) fn leaf_at_offset(&self) -> LeafAtOffset<&'a SyntaxNode> {
        find_leaf_at_offset(self.source_file.syntax(), self.range.start())
    }
    pub(crate) fn node_at_offset<N: AstNode>(&self) -> Option<&'a N> {
        find_node_at_offset(self.source_file.syntax(), self.range.start())
    }
    pub(crate) fn covering_node(&self) -> &'a SyntaxNode {
        find_covering_node(self.source_file.syntax(), self.range)
    }
}

impl AssistBuilder {
    fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    fn replace_node_and_indent(&mut self, node: &SyntaxNode, replace_with: impl Into<String>) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, indent)
        }
        self.replace(node.range(), replace_with)
    }
    #[allow(unused)]
    fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    fn insert(&mut self, offset: TextUnit, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    fn set_cursor(&mut self, offset: TextUnit) {
        self.cursor_position = Some(offset)
    }
}

fn reindent(text: &str, indent: &str) -> String {
    let indent = format!("\n{}", indent);
    text.lines().intersperse(&indent).collect()
}

#[cfg(test)]
fn check_assist(assist: fn(AssistCtx) -> Option<Assist>, before: &str, after: &str) {
    crate::test_utils::check_action(before, after, |file, off| {
        let range = TextRange::offset_len(off, 0.into());
        AssistCtx::new(file, range).apply(assist)
    })
}

#[cfg(test)]
fn check_assist_range(assist: fn(AssistCtx) -> Option<Assist>, before: &str, after: &str) {
    crate::test_utils::check_action_range(before, after, |file, range| {
        AssistCtx::new(file, range).apply(assist)
    })
}
<commit_msg>Fix typos<commit_after>//! This modules contains various "assists": suggestions for source code edits
//! which are likely to occur at a given cursor positon. For example, if the
//! cursor is on the `,`, a possible assist is swapping the elements around the
//! comma.

mod flip_comma;
mod add_derive;
mod add_impl;
mod introduce_variable;
mod change_visibility;
mod split_import;
mod replace_if_let_with_match;

use ra_text_edit::{TextEdit, TextEditBuilder};
use ra_syntax::{
    Direction, SyntaxNode, TextUnit, TextRange, SourceFile, AstNode,
    algo::{find_leaf_at_offset, find_node_at_offset, find_covering_node, LeafAtOffset},
};
use itertools::Itertools;

use crate::formatting::leading_indent;

pub use self::{
    flip_comma::flip_comma,
    add_derive::add_derive,
    add_impl::add_impl,
    introduce_variable::introduce_variable,
    change_visibility::change_visibility,
    split_import::split_import,
    replace_if_let_with_match::replace_if_let_with_match,
};

/// Return all the assists applicable at the given position.
pub fn assists(file: &SourceFile, range: TextRange) -> Vec<LocalEdit> {
    let ctx = AssistCtx::new(file, range);
    [
        flip_comma,
        add_derive,
        add_impl,
        introduce_variable,
        change_visibility,
        split_import,
        replace_if_let_with_match,
    ]
    .iter()
    .filter_map(|&assist| ctx.clone().apply(assist))
    .collect()
}

#[derive(Debug)]
pub struct LocalEdit {
    pub label: String,
    pub edit: TextEdit,
    pub cursor_position: Option<TextUnit>,
}

fn non_trivia_sibling(node: &SyntaxNode, direction: Direction) -> Option<&SyntaxNode> {
    node.siblings(direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

/// `AssistCtx` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat overengeneered approach, given the current needs. The
/// assists workflow consists of two phases. In the first phase, a user asks for
/// the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then neccessary to answer
///   "is assist applicable" in the first phase.
/// * second, when we are appling assist, we don't have a gurantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach remeniscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `should_compute_edit`. It first check if an edit is applicable (potentially
/// computing info required to compute the actual edit). If it is applicable,
/// and `should_compute_edit` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `should_compute_edit = false`, and then applying the selected edit
/// again, with `should_compute_edit = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagarly :-)
#[derive(Debug, Clone)]
pub struct AssistCtx<'a> {
    source_file: &'a SourceFile,
    range: TextRange,
    should_compute_edit: bool,
}

#[derive(Debug)]
pub enum Assist {
    Applicable,
    Edit(LocalEdit),
}

#[derive(Default)]
struct AssistBuilder {
    edit: TextEditBuilder,
    cursor_position: Option<TextUnit>,
}

impl<'a> AssistCtx<'a> {
    pub fn new(source_file: &'a SourceFile, range: TextRange) -> AssistCtx {
        AssistCtx {
            source_file,
            range,
            should_compute_edit: false,
        }
    }

    pub fn apply(mut self, assist: fn(AssistCtx) -> Option<Assist>) -> Option<LocalEdit> {
        self.should_compute_edit = true;
        match assist(self) {
            None => None,
            Some(Assist::Edit(e)) => Some(e),
            Some(Assist::Applicable) => unreachable!(),
        }
    }

    pub fn check(mut self, assist: fn(AssistCtx) -> Option<Assist>) -> bool {
        self.should_compute_edit = false;
        match assist(self) {
            None => false,
            Some(Assist::Edit(_)) => unreachable!(),
            Some(Assist::Applicable) => true,
        }
    }

    fn build(self, label: impl Into<String>, f: impl FnOnce(&mut AssistBuilder)) -> Option<Assist> {
        if !self.should_compute_edit {
            return Some(Assist::Applicable);
        }
        let mut edit = AssistBuilder::default();
        f(&mut edit);
        Some(Assist::Edit(LocalEdit {
            label: label.into(),
            edit: edit.edit.finish(),
            cursor_position: edit.cursor_position,
        }))
    }

    pub(crate) fn leaf_at_offset(&self) -> LeafAtOffset<&'a SyntaxNode> {
        find_leaf_at_offset(self.source_file.syntax(), self.range.start())
    }
    pub(crate) fn node_at_offset<N: AstNode>(&self) -> Option<&'a N> {
        find_node_at_offset(self.source_file.syntax(), self.range.start())
    }
    pub(crate) fn covering_node(&self) -> &'a SyntaxNode {
        find_covering_node(self.source_file.syntax(), self.range)
    }
}

impl AssistBuilder {
    fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    fn replace_node_and_indent(&mut self, node: &SyntaxNode, replace_with: impl Into<String>) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, indent)
        }
        self.replace(node.range(), replace_with)
    }
    #[allow(unused)]
    fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    fn insert(&mut self, offset: TextUnit, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    fn set_cursor(&mut self, offset: TextUnit) {
        self.cursor_position = Some(offset)
    }
}

fn reindent(text: &str, indent: &str) -> String {
    let indent = format!("\n{}", indent);
    text.lines().intersperse(&indent).collect()
}

#[cfg(test)]
fn check_assist(assist: fn(AssistCtx) -> Option<Assist>, before: &str, after: &str) {
    crate::test_utils::check_action(before, after, |file, off| {
        let range = TextRange::offset_len(off, 0.into());
        AssistCtx::new(file, range).apply(assist)
    })
}

#[cfg(test)]
fn check_assist_range(assist: fn(AssistCtx) -> Option<Assist>, before: &str, after: &str) {
    crate::test_utils::check_action_range(before, after, |file, range| {
        AssistCtx::new(file, range).apply(assist)
    })
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use anyhow::{anyhow, bail, Context, Result};
use std::{
    env,
    path::{Path, PathBuf},
    process::Command,
};

use ra_arena::{impl_arena_id, Arena, RawId};

#[derive(Default, Debug, Clone)]
pub struct Sysroot {
    crates: Arena<SysrootCrate, SysrootCrateData>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SysrootCrate(RawId);
impl_arena_id!(SysrootCrate);

#[derive(Debug, Clone)]
struct SysrootCrateData {
    name: String,
    root: PathBuf,
    deps: Vec<SysrootCrate>,
}

impl Sysroot {
    pub fn core(&self) -> Option<SysrootCrate> {
        self.by_name("core")
    }

    pub fn alloc(&self) -> Option<SysrootCrate> {
        self.by_name("alloc")
    }

    pub fn std(&self) -> Option<SysrootCrate> {
        self.by_name("std")
    }

    pub fn proc_macro(&self) -> Option<SysrootCrate> {
        self.by_name("proc_macro")
    }

    pub fn crates<'a>(&'a self) -> impl Iterator<Item = SysrootCrate> + ExactSizeIterator + 'a {
        self.crates.iter().map(|(id, _data)| id)
    }

    pub fn discover(cargo_toml: &Path) -> Result<Sysroot> {
        let src = get_or_install_rust_src(cargo_toml)?;

        if !src.exists() {
            Err(anyhow!(
                "can't load standard library from sysroot\n\
                {}\n\
                (discovered via `rustc --print sysroot`)\n\
                try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
                src.display(),
            ))?;
        }

        let mut sysroot = Sysroot { crates: Arena::default() };
        for name in SYSROOT_CRATES.trim().lines() {
            let root = src.join(format!("lib{}", name)).join("lib.rs");
            if root.exists() {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }
        if let Some(std) = sysroot.std() {
            for dep in STD_DEPS.trim().lines() {
                if let Some(dep) = sysroot.by_name(dep) {
                    sysroot.crates[std].deps.push(dep)
                }
            }
        }
        if let Some(alloc) = sysroot.by_name("alloc") {
            if let Some(core) = sysroot.core() {
                sysroot.crates[alloc].deps.push(core);
            }
        }
        Ok(sysroot)
    }

    fn by_name(&self, name: &str) -> Option<SysrootCrate> {
        self.crates.iter().find(|(_id, data)| data.name == name).map(|(id, _data)| id)
    }
}

fn get_or_install_rust_src(cargo_toml: &Path) -> Result<PathBuf> {
    fn try_find_src_path(cargo_toml: &Path) -> Result<PathBuf> {
        if let Ok(path) = env::var("RUST_SRC_PATH") {
            return Ok(path.into());
        }

        let rustc_output = Command::new("rustc")
            .current_dir(cargo_toml.parent().unwrap())
            .args(&["--print", "sysroot"])
            .output()
            .context("rustc --print sysroot failed")?;
        if !rustc_output.status.success() {
            match rustc_output.status.code() {
                Some(code) => bail!(
                    "failed to locate sysroot: rustc --print sysroot exited with code {}",
                    code
                ),
                None => {
                    bail!("failed to locate sysroot: rustc --print sysroot terminated by signal")
                }
            };
        }
        let stdout = String::from_utf8(rustc_output.stdout)?;
        let sysroot_path = Path::new(stdout.trim());
        Ok(sysroot_path.join("lib/rustlib/src/rust/src"))
    }

    fn try_install_rust_src(cargo_toml: &Path) -> Result<PathBuf> {
        let rustup_output = Command::new("rustup")
            .current_dir(cargo_toml.parent().unwrap())
            .args(&["component", "add", "rust-src"])
            .output()
            .context("rustup component add rust-src failed")?;
        if !rustup_output.status.success() {
            match rustup_output.status.code() {
                Some(code) => bail!(
                    "failed to install rust-src: rustup component add rust-src exited with code {}",
                    code
                ),
                None => bail!(
                    "failed to install rust-src: rustup component add rust-src terminated by signal"
                ),
            };
        }
        try_find_src_path(cargo_toml)
    }

    let src = try_find_src_path(cargo_toml)?;
    if !src.exists() {
        try_install_rust_src(cargo_toml)
    } else {
        Ok(src)
    }
}

impl SysrootCrate {
    pub fn name(self, sysroot: &Sysroot) -> &str {
        &sysroot.crates[self].name
    }
    pub fn root(self, sysroot: &Sysroot) -> &Path {
        sysroot.crates[self].root.as_path()
    }
    pub fn root_dir(self, sysroot: &Sysroot) -> &Path {
        self.root(sysroot).parent().unwrap()
    }
    pub fn deps<'a>(self, sysroot: &'a Sysroot) -> impl Iterator<Item = SysrootCrate> + 'a {
        sysroot.crates[self].deps.iter().copied()
    }
}

const SYSROOT_CRATES: &str = "
std
core
alloc
collections
libc
panic_unwind
proc_macro
rustc_unicode
std_unicode
test
alloc_jemalloc
alloc_system
compiler_builtins
getopts
panic_unwind
panic_abort
rand
term
unwind
build_helper
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
syntax";

const STD_DEPS: &str = "
alloc
alloc_jemalloc
alloc_system
core
panic_abort
rand
compiler_builtins
unwind
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
build_helper";
<commit_msg>Simplify the code further<commit_after>//! FIXME: write short doc here

use anyhow::{bail, Context, Result};
use std::{
    env,
    path::{Path, PathBuf},
    process::Command,
};

use ra_arena::{impl_arena_id, Arena, RawId};

#[derive(Default, Debug, Clone)]
pub struct Sysroot {
    crates: Arena<SysrootCrate, SysrootCrateData>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SysrootCrate(RawId);
impl_arena_id!(SysrootCrate);

#[derive(Debug, Clone)]
struct SysrootCrateData {
    name: String,
    root: PathBuf,
    deps: Vec<SysrootCrate>,
}

impl Sysroot {
    pub fn core(&self) -> Option<SysrootCrate> {
        self.by_name("core")
    }

    pub fn alloc(&self) -> Option<SysrootCrate> {
        self.by_name("alloc")
    }

    pub fn std(&self) -> Option<SysrootCrate> {
        self.by_name("std")
    }

    pub fn proc_macro(&self) -> Option<SysrootCrate> {
        self.by_name("proc_macro")
    }

    pub fn crates<'a>(&'a self) -> impl Iterator<Item = SysrootCrate> + ExactSizeIterator + 'a {
        self.crates.iter().map(|(id, _data)| id)
    }

    pub fn discover(cargo_toml: &Path) -> Result<Sysroot> {
        let src = get_or_install_rust_src(cargo_toml)?;
        let mut sysroot = Sysroot { crates: Arena::default() };
        for name in SYSROOT_CRATES.trim().lines() {
            let root = src.join(format!("lib{}", name)).join("lib.rs");
            if root.exists() {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }
        if let Some(std) = sysroot.std() {
            for dep in STD_DEPS.trim().lines() {
                if let Some(dep) = sysroot.by_name(dep) {
                    sysroot.crates[std].deps.push(dep)
                }
            }
        }
        if let Some(alloc) = sysroot.by_name("alloc") {
            if let Some(core) = sysroot.core() {
                sysroot.crates[alloc].deps.push(core);
            }
        }
        Ok(sysroot)
    }

    fn by_name(&self, name: &str) -> Option<SysrootCrate> {
        self.crates.iter().find(|(_id, data)| data.name == name).map(|(id, _data)| id)
    }
}

fn get_or_install_rust_src(cargo_toml: &Path) -> Result<PathBuf> {
    fn try_install_rust_src(cargo_toml: &Path) -> Result<()> {
        let rustup_output = Command::new("rustup")
            .current_dir(cargo_toml.parent().unwrap())
            .args(&["component", "add", "rust-src"])
            .output()
            .context("rustup component add rust-src failed")?;
        if !rustup_output.status.success() {
            match rustup_output.status.code() {
                Some(code) => bail!(
                    "failed to install rust-src: rustup component add rust-src exited with code {}",
                    code
                ),
                None => bail!(
                    "failed to install rust-src: rustup component add rust-src terminated by signal"
                ),
            };
        }
        Ok(())
    }

    if let Ok(path) = env::var("RUST_SRC_PATH") {
        return Ok(path.into());
    }
    let rustc_output = Command::new("rustc")
        .current_dir(cargo_toml.parent().unwrap())
        .args(&["--print", "sysroot"])
        .output()
        .context("rustc --print sysroot failed")?;
    if !rustc_output.status.success() {
        match rustc_output.status.code() {
            Some(code) => {
                bail!("failed to locate sysroot: rustc --print sysroot exited with code {}", code)
            }
            None => bail!("failed to locate sysroot: rustc --print sysroot terminated by signal"),
        };
    }
    let stdout = String::from_utf8(rustc_output.stdout)?;
    let sysroot_path = Path::new(stdout.trim());
    let src_path = sysroot_path.join("lib/rustlib/src/rust/src");

    if !src_path.exists() {
        try_install_rust_src(cargo_toml)?
    }
    if !src_path.exists() {
        bail!(
            "can't load standard library from sysroot\n\
            {}\n\
            (discovered via `rustc --print sysroot`)\n\
            try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
            src_path.display(),
        )
    }
    Ok(src_path)
}

impl SysrootCrate {
    pub fn name(self, sysroot: &Sysroot) -> &str {
        &sysroot.crates[self].name
    }
    pub fn root(self, sysroot: &Sysroot) -> &Path {
        sysroot.crates[self].root.as_path()
    }
    pub fn root_dir(self, sysroot: &Sysroot) -> &Path {
        self.root(sysroot).parent().unwrap()
    }
    pub fn deps<'a>(self, sysroot: &'a Sysroot) -> impl Iterator<Item = SysrootCrate> + 'a {
        sysroot.crates[self].deps.iter().copied()
    }
}

const SYSROOT_CRATES: &str = "
std
core
alloc
collections
libc
panic_unwind
proc_macro
rustc_unicode
std_unicode
test
alloc_jemalloc
alloc_system
compiler_builtins
getopts
panic_unwind
panic_abort
rand
term
unwind
build_helper
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
syntax";

const STD_DEPS: &str = "
alloc
alloc_jemalloc
alloc_system
core
panic_abort
rand
compiler_builtins
unwind
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
build_helper";
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use crate::TextRange;

use ra_syntax::{
    ast::{self, AttrsOwner, NameOwner, TypeAscriptionOwner, TypeParamsOwner},
    match_ast, AstNode, SourceFile, SyntaxKind, SyntaxNode, WalkEvent,
};

#[derive(Debug, Clone)]
pub struct StructureNode {
    pub parent: Option<usize>,
    pub label: String,
    pub navigation_range: TextRange,
    pub node_range: TextRange,
    pub kind: SyntaxKind,
    pub detail: Option<String>,
    pub deprecated: bool,
}

pub fn file_structure(file: &SourceFile) -> Vec<StructureNode> {
    let mut res = Vec::new();
    let mut stack = Vec::new();

    for event in file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = structure_node(&node) {
                    symbol.parent = stack.last().copied();
                    stack.push(res.len());
                    res.push(symbol);
                }
            }
            WalkEvent::Leave(node) => {
                if structure_node(&node).is_some() {
                    stack.pop().unwrap();
                }
            }
        }
    }
    res
}

fn structure_node(node: &SyntaxNode) -> Option<StructureNode> {
    fn decl<N: NameOwner + AttrsOwner>(node: N) -> Option<StructureNode> {
        decl_with_detail(node, None)
    }

    fn decl_with_ascription<N: NameOwner + AttrsOwner + TypeAscriptionOwner>(
        node: N,
    ) -> Option<StructureNode> {
        let ty = node.ascribed_type();
        decl_with_type_ref(node, ty)
    }

    fn decl_with_type_ref<N: NameOwner + AttrsOwner>(
        node: N,
        type_ref: Option<ast::TypeRef>,
    ) -> Option<StructureNode> {
        let detail = type_ref.map(|type_ref| {
            let mut detail = String::new();
            collapse_ws(type_ref.syntax(), &mut detail);
            detail
        });
        decl_with_detail(node, detail)
    }

    fn decl_with_detail<N: NameOwner + AttrsOwner>(
        node: N,
        detail: Option<String>,
    ) -> Option<StructureNode> {
        let name = node.name()?;

        Some(StructureNode {
            parent: None,
            label: name.text().to_string(),
            navigation_range: name.syntax().text_range(),
            node_range: node.syntax().text_range(),
            kind: node.syntax().kind(),
            detail,
            deprecated: node.attrs().filter_map(|x| x.simple_name()).any(|x| x == "deprecated"),
        })
    }

    fn collapse_ws(node: &SyntaxNode, output: &mut String) {
        let mut can_insert_ws = false;
        node.text().for_each_chunk(|chunk| {
            for line in chunk.lines() {
                let line = line.trim();
                if line.is_empty() {
                    if can_insert_ws {
                        output.push(' ');
                        can_insert_ws = false;
                    }
                } else {
                    output.push_str(line);
                    can_insert_ws = true;
                }
            }
        })
    }

    match_ast! {
        match node {
            ast::FnDef(it) => {
                let mut detail = String::from("fn");
                if let Some(type_param_list) = it.type_param_list() {
                    collapse_ws(type_param_list.syntax(), &mut detail);
                }
                if let Some(param_list) = it.param_list() {
                    collapse_ws(param_list.syntax(), &mut detail);
                }
                if let Some(ret_type) = it.ret_type() {
                    detail.push_str(" ");
                    collapse_ws(ret_type.syntax(), &mut detail);
                }

                decl_with_detail(it, Some(detail))
            },
            ast::StructDef(it) => { decl(it) },
            ast::EnumDef(it) => { decl(it) },
            ast::EnumVariant(it) => { decl(it) },
            ast::TraitDef(it) => { decl(it) },
            ast::Module(it) => { decl(it) },
            ast::TypeAliasDef(it) => {
                let ty = it.type_ref();
                decl_with_type_ref(it, ty)
            },
            ast::RecordFieldDef(it) => { decl_with_ascription(it) },
            ast::ConstDef(it) => { decl_with_ascription(it) },
            ast::StaticDef(it) => { decl_with_ascription(it) },
            ast::ImplBlock(it) => {
                let target_type = it.target_type()?;
                let target_trait = it.target_trait();
                let label = match target_trait {
                    None => format!("impl {}", target_type.syntax().text()),
                    Some(t) => {
                        format!("impl {} for {}", t.syntax().text(), target_type.syntax().text(),)
                    }
                };

                let node = StructureNode {
                    parent: None,
                    label,
                    navigation_range: target_type.syntax().text_range(),
                    node_range: it.syntax().text_range(),
                    kind: it.syntax().kind(),
                    detail: None,
                    deprecated: false,
                };
                Some(node)
            },
            ast::MacroCall(it) => {
                let macro_name = it.syntax()
                    .children()
                    .find(|c|
                        ![
                            SyntaxKind::COMMENT,
                            SyntaxKind::WHITESPACE,
                            SyntaxKind::ATTR
                        ].iter()
                        .any(|&k| k == c.kind())
                    );

                match macro_name {
                    None => return None,
                    Some(n) => if n.first_token().unwrap().text().as_str() != "macro_rules" {
                        return None;
                    }
                }

                decl(it)
            },
            _ => None,
        }
    }
}

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

    #[test]
    fn test_file_structure() {
        let file = SourceFile::parse(
            r#"
struct Foo {
    x: i32
}

mod m {
    fn bar1() {}
    fn bar2<T>(t: T) -> T {}
    fn bar3<A,
        B>(a: A,
        b: B) -> Vec<
        u32
    > {}
}

enum E { X, Y(i32) }
type T = ();
static S: i32 = 92;
const C: i32 = 92;

impl E {}

impl fmt::Debug for E {}

macro_rules! mc {
    () => {}
}

#[deprecated]
fn obsolete() {}

#[deprecated(note = "for awhile")]
fn very_obsolete() {}
"#,
        )
        .ok()
        .unwrap();
        let structure = file_structure(&file);
        assert_debug_snapshot!(structure,
        @r###"
        [
            StructureNode {
                parent: None,
                label: "Foo",
                navigation_range: [8; 11),
                node_range: [1; 26),
                kind: STRUCT_DEF,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    0,
                ),
                label: "x",
                navigation_range: [18; 19),
                node_range: [18; 24),
                kind: RECORD_FIELD_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "m",
                navigation_range: [32; 33),
                node_range: [28; 158),
                kind: MODULE,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar1",
                navigation_range: [43; 47),
                node_range: [40; 52),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar2",
                navigation_range: [60; 64),
                node_range: [57; 81),
                kind: FN_DEF,
                detail: Some(
                    "fn<T>(t: T) -> T",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar3",
                navigation_range: [89; 93),
                node_range: [86; 156),
                kind: FN_DEF,
                detail: Some(
                    "fn<A, B>(a: A, b: B) -> Vec< u32 >",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "E",
                navigation_range: [165; 166),
                node_range: [160; 180),
                kind: ENUM_DEF,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    6,
                ),
                label: "X",
                navigation_range: [169; 170),
                node_range: [169; 170),
                kind: ENUM_VARIANT,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    6,
                ),
                label: "Y",
                navigation_range: [172; 173),
                node_range: [172; 178),
                kind: ENUM_VARIANT,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "T",
                navigation_range: [186; 187),
                node_range: [181; 193),
                kind: TYPE_ALIAS_DEF,
                detail: Some(
                    "()",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "S",
                navigation_range: [201; 202),
                node_range: [194; 213),
                kind: STATIC_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "C",
                navigation_range: [220; 221),
                node_range: [214; 232),
                kind: CONST_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "impl E",
                navigation_range: [239; 240),
                node_range: [234; 243),
                kind: IMPL_BLOCK,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "impl fmt::Debug for E",
                navigation_range: [265; 266),
                node_range: [245; 269),
                kind: IMPL_BLOCK,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "mc",
                navigation_range: [284; 286),
                node_range: [271; 303),
                kind: MACRO_CALL,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "obsolete",
                navigation_range: [322; 330),
                node_range: [305; 335),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: true,
            },
            StructureNode {
                parent: None,
                label: "very_obsolete",
                navigation_range: [375; 388),
                node_range: [337; 393),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: true,
            },
        ]
        "###
                );
    }
}
<commit_msg>Added test for file_structure unidentified macro definition<commit_after>//! FIXME: write short doc here

use crate::TextRange;

use ra_syntax::{
    ast::{self, AttrsOwner, NameOwner, TypeAscriptionOwner, TypeParamsOwner},
    match_ast, AstNode, SourceFile, SyntaxKind, SyntaxNode, WalkEvent,
};

#[derive(Debug, Clone)]
pub struct StructureNode {
    pub parent: Option<usize>,
    pub label: String,
    pub navigation_range: TextRange,
    pub node_range: TextRange,
    pub kind: SyntaxKind,
    pub detail: Option<String>,
    pub deprecated: bool,
}

pub fn file_structure(file: &SourceFile) -> Vec<StructureNode> {
    let mut res = Vec::new();
    let mut stack = Vec::new();

    for event in file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = structure_node(&node) {
                    symbol.parent = stack.last().copied();
                    stack.push(res.len());
                    res.push(symbol);
                }
            }
            WalkEvent::Leave(node) => {
                if structure_node(&node).is_some() {
                    stack.pop().unwrap();
                }
            }
        }
    }
    res
}

fn structure_node(node: &SyntaxNode) -> Option<StructureNode> {
    fn decl<N: NameOwner + AttrsOwner>(node: N) -> Option<StructureNode> {
        decl_with_detail(node, None)
    }

    fn decl_with_ascription<N: NameOwner + AttrsOwner + TypeAscriptionOwner>(
        node: N,
    ) -> Option<StructureNode> {
        let ty = node.ascribed_type();
        decl_with_type_ref(node, ty)
    }

    fn decl_with_type_ref<N: NameOwner + AttrsOwner>(
        node: N,
        type_ref: Option<ast::TypeRef>,
    ) -> Option<StructureNode> {
        let detail = type_ref.map(|type_ref| {
            let mut detail = String::new();
            collapse_ws(type_ref.syntax(), &mut detail);
            detail
        });
        decl_with_detail(node, detail)
    }

    fn decl_with_detail<N: NameOwner + AttrsOwner>(
        node: N,
        detail: Option<String>,
    ) -> Option<StructureNode> {
        let name = node.name()?;

        Some(StructureNode {
            parent: None,
            label: name.text().to_string(),
            navigation_range: name.syntax().text_range(),
            node_range: node.syntax().text_range(),
            kind: node.syntax().kind(),
            detail,
            deprecated: node.attrs().filter_map(|x| x.simple_name()).any(|x| x == "deprecated"),
        })
    }

    fn collapse_ws(node: &SyntaxNode, output: &mut String) {
        let mut can_insert_ws = false;
        node.text().for_each_chunk(|chunk| {
            for line in chunk.lines() {
                let line = line.trim();
                if line.is_empty() {
                    if can_insert_ws {
                        output.push(' ');
                        can_insert_ws = false;
                    }
                } else {
                    output.push_str(line);
                    can_insert_ws = true;
                }
            }
        })
    }

    match_ast! {
        match node {
            ast::FnDef(it) => {
                let mut detail = String::from("fn");
                if let Some(type_param_list) = it.type_param_list() {
                    collapse_ws(type_param_list.syntax(), &mut detail);
                }
                if let Some(param_list) = it.param_list() {
                    collapse_ws(param_list.syntax(), &mut detail);
                }
                if let Some(ret_type) = it.ret_type() {
                    detail.push_str(" ");
                    collapse_ws(ret_type.syntax(), &mut detail);
                }

                decl_with_detail(it, Some(detail))
            },
            ast::StructDef(it) => { decl(it) },
            ast::EnumDef(it) => { decl(it) },
            ast::EnumVariant(it) => { decl(it) },
            ast::TraitDef(it) => { decl(it) },
            ast::Module(it) => { decl(it) },
            ast::TypeAliasDef(it) => {
                let ty = it.type_ref();
                decl_with_type_ref(it, ty)
            },
            ast::RecordFieldDef(it) => { decl_with_ascription(it) },
            ast::ConstDef(it) => { decl_with_ascription(it) },
            ast::StaticDef(it) => { decl_with_ascription(it) },
            ast::ImplBlock(it) => {
                let target_type = it.target_type()?;
                let target_trait = it.target_trait();
                let label = match target_trait {
                    None => format!("impl {}", target_type.syntax().text()),
                    Some(t) => {
                        format!("impl {} for {}", t.syntax().text(), target_type.syntax().text(),)
                    }
                };

                let node = StructureNode {
                    parent: None,
                    label,
                    navigation_range: target_type.syntax().text_range(),
                    node_range: it.syntax().text_range(),
                    kind: it.syntax().kind(),
                    detail: None,
                    deprecated: false,
                };
                Some(node)
            },
            ast::MacroCall(it) => {
                let macro_name = it.syntax()
                    .children()
                    .find(|c|
                        ![
                            SyntaxKind::COMMENT,
                            SyntaxKind::WHITESPACE,
                            SyntaxKind::ATTR
                        ].iter()
                        .any(|&k| k == c.kind())
                    );

                match macro_name {
                    None => return None,
                    Some(n) => if n.first_token().unwrap().text().as_str() != "macro_rules" {
                        return None;
                    }
                }

                decl(it)
            },
            _ => None,
        }
    }
}

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

    #[test]
    fn test_file_structure() {
        let file = SourceFile::parse(
            r#"
struct Foo {
    x: i32
}

mod m {
    fn bar1() {}
    fn bar2<T>(t: T) -> T {}
    fn bar3<A,
        B>(a: A,
        b: B) -> Vec<
        u32
    > {}
}

enum E { X, Y(i32) }
type T = ();
static S: i32 = 92;
const C: i32 = 92;

impl E {}

impl fmt::Debug for E {}

macro_rules! mc {
    () => {}
}

#[macro_export]
macro_rules! mcexp {
    () => {}
}

#[deprecated]
fn obsolete() {}

#[deprecated(note = "for awhile")]
fn very_obsolete() {}
"#,
        )
        .ok()
        .unwrap();
        let structure = file_structure(&file);
        assert_debug_snapshot!(structure,
        @r###"
        [
            StructureNode {
                parent: None,
                label: "Foo",
                navigation_range: [8; 11),
                node_range: [1; 26),
                kind: STRUCT_DEF,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    0,
                ),
                label: "x",
                navigation_range: [18; 19),
                node_range: [18; 24),
                kind: RECORD_FIELD_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "m",
                navigation_range: [32; 33),
                node_range: [28; 158),
                kind: MODULE,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar1",
                navigation_range: [43; 47),
                node_range: [40; 52),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar2",
                navigation_range: [60; 64),
                node_range: [57; 81),
                kind: FN_DEF,
                detail: Some(
                    "fn<T>(t: T) -> T",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    2,
                ),
                label: "bar3",
                navigation_range: [89; 93),
                node_range: [86; 156),
                kind: FN_DEF,
                detail: Some(
                    "fn<A, B>(a: A, b: B) -> Vec< u32 >",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "E",
                navigation_range: [165; 166),
                node_range: [160; 180),
                kind: ENUM_DEF,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    6,
                ),
                label: "X",
                navigation_range: [169; 170),
                node_range: [169; 170),
                kind: ENUM_VARIANT,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: Some(
                    6,
                ),
                label: "Y",
                navigation_range: [172; 173),
                node_range: [172; 178),
                kind: ENUM_VARIANT,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "T",
                navigation_range: [186; 187),
                node_range: [181; 193),
                kind: TYPE_ALIAS_DEF,
                detail: Some(
                    "()",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "S",
                navigation_range: [201; 202),
                node_range: [194; 213),
                kind: STATIC_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "C",
                navigation_range: [220; 221),
                node_range: [214; 232),
                kind: CONST_DEF,
                detail: Some(
                    "i32",
                ),
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "impl E",
                navigation_range: [239; 240),
                node_range: [234; 243),
                kind: IMPL_BLOCK,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "impl fmt::Debug for E",
                navigation_range: [265; 266),
                node_range: [245; 269),
                kind: IMPL_BLOCK,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "mc",
                navigation_range: [284; 286),
                node_range: [271; 303),
                kind: MACRO_CALL,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "mcexp",
                navigation_range: [334; 339),
                node_range: [305; 356),
                kind: MACRO_CALL,
                detail: None,
                deprecated: false,
            },
            StructureNode {
                parent: None,
                label: "obsolete",
                navigation_range: [375; 383),
                node_range: [358; 388),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: true,
            },
            StructureNode {
                parent: None,
                label: "very_obsolete",
                navigation_range: [428; 441),
                node_range: [390; 446),
                kind: FN_DEF,
                detail: Some(
                    "fn()",
                ),
                deprecated: true,
            },
        ]
        "###
                );
    }
}
<|endoftext|>