generate rust code. (#86)

* WIP generate rust code.

* add test

* flake8

* Update readme

* rust template: avoid stack over flow.

* fix comment

* change rust extension

* support rust submit

* refactoring

* fix indent level

Author: fukatani

README.md

@@ -10,7 +10,7 @@ Python 3.5 以降で動作する [AtCoder](http://atcoder.jp/) からサンプ
 - AtCoderへのログイン，入出力例データなどの抽出
 - 枝刈り探索による高精度・高速な入力フォーマット解析 (ARC、ABC、AGCについては約9割ほど)
 - 問題文中に含まれるMOD値やYES/NO文字列等の定数値抽出
-- 入力フォーマット解析結果や抽出した定数値を用いたテンプレートコードの自動生成(C++, Java)
+- 入力フォーマット解析結果や抽出した定数値を用いたテンプレートコードの自動生成(C++, Java, Rust)
     - カスタムテンプレートに対応
     - 他言語対応のためのコントリビューション(≒中間形式からコードに変換する部分のPR)を募集中です!
 - コード提出機能

atcodertools/codegen/code_generators/rust.py

@@ -0,0 +1,149 @@
+from atcodertools.codegen.models.code_gen_args import CodeGenArgs
+from atcodertools.codegen.code_generators.cpp import CppCodeGenerator
+from atcodertools.codegen.template_engine import render
+from atcodertools.fmtprediction.models.format import Pattern, SingularPattern, ParallelPattern, TwoDimensionalPattern
+from atcodertools.fmtprediction.models.type import Type
+from atcodertools.fmtprediction.models.variable import Variable
+
+
+# RustCodeGenerator uses part of CppCodeGenerator just for less code clone.
+
+
+def _loop_header(var: Variable, for_second_index: bool):
+    if for_second_index:
+        index = var.second_index
+        loop_var = "j"
+    else:
+        index = var.first_index
+        loop_var = "i"
+
+    return "for {loop_var} in 0..({length}) as usize {{".format(
+        loop_var=loop_var,
+        length=index.get_length()
+    )
+
+
+class RustCodeGenerator(CppCodeGenerator):
+
+    def _input_part(self):
+        lines = ["let con = read_string();",
+                 "let mut scanner = Scanner::new(&con);"]
+        for pattern in self._format.sequence:
+            lines += self._render_pattern(pattern)
+        return "\n{indent}".format(indent=self._indent(1)).join(lines)
+
+    def _convert_type(self, type_: Type) -> str:
+        if type_ == Type.float:
+            return "f64"
+        elif type_ == Type.int:
+            return "i64"
+        elif type_ == Type.str:
+            return "String"
+        else:
+            raise NotImplementedError
+
+    def _get_declaration_type(self, var: Variable):
+        if var.dim_num() == 0:
+            template = "{type}"
+        elif var.dim_num() == 1:
+            template = "Vec<{type}>"
+        elif var.dim_num() == 2:
+            template = "Vec<Vec<{type}>>"
+        else:
+            raise NotImplementedError
+        return template.format(type=self._convert_type(var.type))
+
+    def _actual_arguments(self) -> str:
+        return ", ".join([v.name for v in self._format.all_vars()])
+
+    def _formal_arguments(self):
+        """
+            :return the string form of formal arguments e.g. "N: i64, K: i64, a: Vec<i64>"
+        """
+        return ", ".join([
+            "{name}: {decl_type}".format(
+                decl_type=self._get_declaration_type(v),
+                name=v.name)
+            for v in self._format.all_vars()
+        ])
+
+    def _generate_declaration(self, var: Variable):
+        if var.dim_num() == 0:
+            constructor = ""
+        elif var.dim_num() == 1:
+            if var.type == Type.str:
+                constructor = " = vec![String::new(); ({size}) as usize]".format(
+                    type=self._convert_type(var.type),
+                    size=var.first_index.get_length()
+                )
+            else:
+                constructor = " = vec![0{type}; ({size}) as usize]".format(
+                    type=self._convert_type(var.type),
+                    size=var.first_index.get_length()
+                )
+        elif var.dim_num() == 2:
+            if var.type == Type.str:
+                constructor = " = vec![vec![String::new(); ({col_size}) as usize]; ({row_size}) as usize]".format(
+                    type=self._convert_type(var.type),
+                    row_size=var.first_index.get_length(),
+                    col_size=var.second_index.get_length()
+                )
+            else:
+                constructor = " = vec![vec![0{type}; ({col_size}) as usize]; ({row_size}) as usize]".format(
+                    type=self._convert_type(var.type),
+                    row_size=var.first_index.get_length(),
+                    col_size=var.second_index.get_length()
+                )
+        else:
+            raise NotImplementedError
+
+        line = "let mut {name}: {decl_type}{constructor};".format(
+            name=var.name,
+            decl_type=self._get_declaration_type(var),
+            constructor=constructor
+        )
+        return line
+
+    def _input_code_for_var(self, var: Variable) -> str:
+        name = self._get_var_name(var)
+        return '{name} = scanner.next();'.format(name=name)
+
+    def _render_pattern(self, pattern: Pattern):
+        lines = []
+        for var in pattern.all_vars():
+            lines.append(self._generate_declaration(var))
+
+        representative_var = pattern.all_vars()[0]
+        if isinstance(pattern, SingularPattern):
+            lines.append(self._input_code_for_var(representative_var))
+        elif isinstance(pattern, ParallelPattern):
+            lines.append(_loop_header(representative_var, False))
+            for var in pattern.all_vars():
+                lines.append("{indent}{line}".format(indent=self._indent(1),
+                                                     line=self._input_code_for_var(var)))
+            lines.append("}")
+        elif isinstance(pattern, TwoDimensionalPattern):
+            lines.append(_loop_header(representative_var, False))
+            lines.append(
+                "{indent}{line}".format(indent=self._indent(1), line=_loop_header(representative_var, True)))
+            for var in pattern.all_vars():
+                lines.append("{indent}{line}".format(indent=self._indent(2),
+                                                     line=self._input_code_for_var(var)))
+            lines.append("{indent}}}".format(indent=self._indent(1)))
+            lines.append("}")
+        else:
+            raise NotImplementedError
+
+        return lines
+
+
+def main(args: CodeGenArgs) -> str:
+    code_parameters = RustCodeGenerator(
+        args.format, args.config).generate_parameters()
+    return render(
+        args.template,
+        mod=args.constants.mod,
+        yes_str=args.constants.yes_str,
+        no_str=args.constants.no_str,
+        **code_parameters
+    )

atcodertools/codegen/code_style_config.py

@@ -14,7 +14,7 @@ class CodeStyleConfigInitError(Exception):
 
 
 DEFAULT_WORKSPACE_DIR_PATH = os.path.join(expanduser("~"), "atcoder-workspace")
-SUPPORTED_LANGUAGES = ["cpp", "java"]
+SUPPORTED_LANGUAGES = ["cpp", "java", "rust"]
 
 
 class CodeStyleConfig:
@@ -68,9 +68,12 @@ def __init__(self,
             if lang == "cpp":
                 from atcodertools.codegen.code_generators import cpp
                 self.code_generator = cpp.main
-            else:
+            elif lang == "java":
                 from atcodertools.codegen.code_generators import java
                 self.code_generator = java.main
+            else:
+                from atcodertools.codegen.code_generators import rust
+                self.code_generator = rust.main
 
         self.template_file = normalize_path(
             template_file or get_default_template_path(lang))

atcodertools/tools/envgen.py

@@ -39,6 +39,8 @@ class BannedFileDetectedError(Exception):
 
 
 def extension(lang: str):
+    if lang == 'rust':
+        return 'rs'
     return lang
 
 
@@ -259,7 +261,9 @@ def main(prog, args):
                             "[Default (C++)] {}\n".format(
                                 get_default_template_path('cpp')),
                             "[Default (Java)] {}".format(
-                                get_default_template_path('java')))
+                                get_default_template_path('java')),
+                            "[Default (Rust)] {}".format(
+                                get_default_template_path('rust'))),
                         )
 
     # Deleted functionality

atcodertools/tools/submit.py

@@ -19,6 +19,8 @@ def infer_detailed_lang(lang: str):
         return "Java8 (OpenJDK 1.8.0)"
     if lang == "cpp":
         return "C++14 (GCC 5.4.1)"
+    if lang == "rust":
+        return "Rust (1.15.1)"
     raise NotImplementedError
 
 

atcodertools/tools/templates/default_template.rust

@@ -0,0 +1,78 @@
+use io::*;
+use std::*;
+
+{% if mod %}
+const MOD: i64 = {{ mod }};
+{% endif %}
+{% if yes_str %}
+const YES: String = "{{ yes_str }}";
+{% endif %}
+{% if no_str %}
+const NO: String = "{{ no_str }}";
+{% endif %}
+{% if prediction_success %}
+fn solve({{ formal_arguments }}) {
+
+}
+{% endif %}
+
+fn main() {
+    {% if prediction_success %}
+    {{input_part}}
+    // In order to avoid potential stack overflow, spawn a new thread.
+    let stack_size = 104_857_600; // 100 MB
+    let thd = std::thread::Builder::new().stack_size(stack_size);
+    thd.spawn(move || solve({{ actual_arguments }})).unwrap().join().unwrap();
+    {% else %}
+    // Failed to predict input format
+    {% endif %}
+}
+
+pub mod io {
+    use std;
+    use std::str::FromStr;
+
+    pub struct Scanner<'a> {
+        iter: std::str::SplitWhitespace<'a>,
+    }
+
+    impl<'a> Scanner<'a> {
+        pub fn new(s: &'a str) -> Scanner<'a> {
+            Scanner {
+                iter: s.split_whitespace(),
+            }
+        }
+
+        pub fn next<T: FromStr>(&mut self) -> T {
+            let s = self.iter.next().unwrap();
+            if let Ok(v) = s.parse::<T>() {
+                v
+            } else {
+                panic!("Parse error")
+            }
+        }
+
+        pub fn next_vec_len<T: FromStr>(&mut self) -> Vec<T> {
+            let n: usize = self.next();
+            self.next_vec(n)
+        }
+
+        pub fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T> {
+            (0..n).map(|_| self.next()).collect()
+        }
+    }
+
+    pub fn read_string() -> String {
+        use std::io::Read;
+
+        let mut s = String::new();
+        std::io::stdin().read_to_string(&mut s).unwrap();
+        s
+    }
+
+    pub fn read_line() -> String {
+        let mut s = String::new();
+        std::io::stdin().read_line(&mut s).unwrap();
+        s.trim_right().to_owned()
+    }
+}

tests/resources/test_codegen/template.rust

@@ -0,0 +1,59 @@
+use std::*;
+
+fn solve(${formal_arguments}) {
+
+}
+
+fn main() {
+    ${input_part}
+    solve(${actual_arguments});
+}
+
+pub mod io {
+    use std;
+    use std::str::FromStr;
+
+    pub struct Scanner<'a> {
+        iter: std::str::SplitWhitespace<'a>,
+    }
+
+    impl<'a> Scanner<'a> {
+        pub fn new(s: &'a str) -> Scanner<'a> {
+            Scanner {
+                iter: s.split_whitespace(),
+            }
+        }
+
+        pub fn next<T: FromStr>(&mut self) -> T {
+            let s = self.iter.next().unwrap();
+            if let Ok(v) = s.parse::<T>() {
+                v
+            } else {
+                panic!("Parse error")
+            }
+        }
+
+        pub fn next_vec_len<T: FromStr>(&mut self) -> Vec<T> {
+            let n: usize = self.next();
+            self.next_vec(n)
+        }
+
+        pub fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T> {
+            (0..n).map(|_| self.next()).collect()
+        }
+    }
+
+    pub fn read_string() -> String {
+        use std::io::Read;
+
+        let mut s = String::new();
+        std::io::stdin().read_to_string(&mut s).unwrap();
+        s
+    }
+
+    pub fn read_line() -> String {
+        let mut s = String::new();
+        std::io::stdin().read_line(&mut s).unwrap();
+        s.trim_right().to_owned()
+    }
+}