Merge pull request #12 from youknowone/visitor

Visitor

Author: youknowone

ast/Cargo.toml

@@ -13,6 +13,7 @@ constant-optimization = ["fold"]
 source-code = ["fold"]
 fold = []
 unparse = ["rustpython-literal"]
+visitor = []
 
 [dependencies]
 rustpython-parser-core = { workspace = true }

ast/asdl_rs.py

@@ -12,7 +12,7 @@
 import asdl
 
 TABSIZE = 4
-AUTOGEN_MESSAGE = "// File automatically generated by {}.\n"
+AUTOGEN_MESSAGE = "// File automatically generated by {}.\n\n"
 
 builtin_type_mapping = {
     "identifier": "Ident",
@@ -68,6 +68,7 @@ class TypeInfo:
     enum_name: Optional[str]
     has_userdata: Optional[bool]
     has_attributes: bool
+    empty_field: bool
     children: set
     boxed: bool
     product: bool
@@ -78,6 +79,7 @@ def __init__(self, name):
         self.enum_name = None
         self.has_userdata = None
         self.has_attributes = False
+        self.empty_field = False
         self.children = set()
         self.boxed = False
         self.product = False
@@ -192,10 +194,9 @@ def visitSum(self, sum, name):
             info.has_userdata = False
         else:
             for t in sum.types:
-                if not t.fields:
-                    continue
                 t_info = TypeInfo(t.name)
                 t_info.enum_name = name
+                t_info.empty_field = not t.fields
                 self.typeinfo[t.name] = t_info
                 self.add_children(t.name, t.fields)
             if len(sum.types) > 1:
@@ -534,14 +535,140 @@ def gen_construction(self, header, fields, footer, depth):
 class FoldModuleVisitor(EmitVisitor):
     def visitModule(self, mod):
         depth = 0
-        self.emit('#[cfg(feature = "fold")]', depth)
-        self.emit("pub mod fold {", depth)
-        self.emit("use super::*;", depth + 1)
-        self.emit("use crate::fold_helpers::Foldable;", depth + 1)
-        FoldTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth + 1)
-        FoldImplVisitor(self.file, self.typeinfo).visit(mod, depth + 1)
+        self.emit("use crate::fold_helpers::Foldable;", depth)
+        FoldTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth)
+        FoldImplVisitor(self.file, self.typeinfo).visit(mod, depth)
+
+
+class VisitorTraitDefVisitor(StructVisitor):
+    def full_name(self, name):
+        typeinfo = self.typeinfo[name]
+        if typeinfo.enum_name:
+            return f"{typeinfo.enum_name}_{name}"
+        else:
+            return name
+
+    def node_type_name(self, name):
+        typeinfo = self.typeinfo[name]
+        if typeinfo.enum_name:
+            return f"{get_rust_type(typeinfo.enum_name)}{get_rust_type(name)}"
+        else:
+            return get_rust_type(name)
+
+    def visitModule(self, mod, depth):
+        self.emit("pub trait Visitor<U=()> {", depth)
+
+        for dfn in mod.dfns:
+            self.visit(dfn, depth + 1)
+        self.emit("}", depth)
+
+    def visitType(self, type, depth=0):
+        self.visit(type.value, type.name, depth)
+
+    def emit_visitor(self, nodename, depth, has_node=True):
+        typeinfo = self.typeinfo[nodename]
+        if has_node:
+            node_type = typeinfo.rust_sum_name
+            node_value = "node"
+        else:
+            node_type = "()"
+            node_value = "()"
+        self.emit(
+            f"fn visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{", depth
+        )
+        self.emit(f"self.generic_visit_{typeinfo.sum_name}({node_value})", depth + 1)
+        self.emit("}", depth)
+
+    def emit_generic_visitor_signature(self, nodename, depth, has_node=True):
+        typeinfo = self.typeinfo[nodename]
+        if has_node:
+            node_type = typeinfo.rust_sum_name
+        else:
+            node_type = "()"
+        self.emit(
+            f"fn generic_visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{",
+            depth,
+        )
+
+    def emit_empty_generic_visitor(self, nodename, depth):
+        self.emit_generic_visitor_signature(nodename, depth)
         self.emit("}", depth)
 
+    def simple_sum(self, sum, name, depth):
+        self.emit_visitor(name, depth)
+        self.emit_empty_generic_visitor(name, depth)
+
+    def visit_match_for_type(self, nodename, rustname, type_, depth):
+        self.emit(f"{rustname}::{type_.name}", depth)
+        if type_.fields:
+            self.emit("(data)", depth)
+            data = "data"
+        else:
+            data = "()"
+        self.emit(f"=> self.visit_{nodename}_{type_.name}({data}),", depth)
+
+    def visit_sumtype(self, name, type_, depth):
+        self.emit_visitor(type_.name, depth, has_node=type_.fields)
+        self.emit_generic_visitor_signature(type_.name, depth, has_node=type_.fields)
+        for f in type_.fields:
+            fieldname = rust_field(f.name)
+            fieldtype = self.typeinfo.get(f.type)
+            if not (fieldtype and fieldtype.has_userdata):
+                continue
+
+            if f.opt:
+                self.emit(f"if let Some(value) = node.{fieldname} {{", depth + 1)
+            elif f.seq:
+                iterable = f"node.{fieldname}"
+                if type_.name == "Dict" and f.name == "keys":
+                    iterable = f"{iterable}.into_iter().flatten()"
+                self.emit(f"for value in {iterable} {{", depth + 1)
+            else:
+                self.emit("{", depth + 1)
+                self.emit(f"let value = node.{fieldname};", depth + 2)
+
+            variable = "value"
+            if fieldtype.boxed and (not f.seq or f.opt):
+                variable = "*" + variable
+            typeinfo = self.typeinfo[fieldtype.name]
+            self.emit(f"self.visit_{typeinfo.sum_name}({variable});", depth + 2)
+
+            self.emit("}", depth + 1)
+
+        self.emit("}", depth)
+
+    def sum_with_constructors(self, sum, name, depth):
+        if not sum.attributes:
+            return
+
+        rustname = enumname = get_rust_type(name)
+        if sum.attributes:
+            rustname = enumname + "Kind"
+        self.emit_visitor(name, depth)
+        self.emit_generic_visitor_signature(name, depth)
+        depth += 1
+        self.emit("match node.node {", depth)
+        for t in sum.types:
+            self.visit_match_for_type(name, rustname, t, depth + 1)
+        self.emit("}", depth)
+        depth -= 1
+        self.emit("}", depth)
+
+        # Now for the visitors for the types
+        for t in sum.types:
+            self.visit_sumtype(name, t, depth)
+
+    def visitProduct(self, product, name, depth):
+        self.emit_visitor(name, depth)
+        self.emit_empty_generic_visitor(name, depth)
+
+
+class VisitorModuleVisitor(EmitVisitor):
+    def visitModule(self, mod):
+        depth = 0
+        self.emit("#[allow(unused_variables, non_snake_case)]", depth)
+        VisitorTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth)
+
 
 class ClassDefVisitor(EmitVisitor):
     def visitModule(self, mod):
@@ -799,23 +926,19 @@ def visit(self, object):
             v.emit("", 0)
 
 
-def write_generic_def(mod, typeinfo, f):
-    f.write(
-        textwrap.dedent(
-            """
-        pub use crate::{Attributed, constant::*};
+def write_ast_def(mod, typeinfo, f):
+    StructVisitor(f, typeinfo).visit(mod)
 
-        type Ident = String;
-        \n
-        """
-        )
-    )
 
-    c = ChainOfVisitors(StructVisitor(f, typeinfo), FoldModuleVisitor(f, typeinfo))
-    c.visit(mod)
+def write_fold_def(mod, typeinfo, f):
+    FoldModuleVisitor(f, typeinfo).visit(mod)
+
+
+def write_visitor_def(mod, typeinfo, f):
+    VisitorModuleVisitor(f, typeinfo).visit(mod)
 
 
-def write_located_def(typeinfo, f):
+def write_located_def(mod, typeinfo, f):
     f.write(
         textwrap.dedent(
             """
@@ -826,6 +949,8 @@ def write_located_def(typeinfo, f):
         )
     )
     for info in typeinfo.values():
+        if info.empty_field:
+            continue
         if info.has_userdata:
             generics = "::<SourceRange>"
         else:
@@ -863,8 +988,7 @@ def write_ast_mod(mod, typeinfo, f):
 
 def main(
     input_filename,
-    generic_filename,
-    located_filename,
+    ast_dir,
     module_filename,
     dump_module=False,
 ):
@@ -879,34 +1003,34 @@ def main(
     typeinfo = {}
     FindUserdataTypesVisitor(typeinfo).visit(mod)
 
-    with generic_filename.open("w") as generic_file, located_filename.open(
-        "w"
-    ) as located_file:
-        generic_file.write(auto_gen_msg)
-        write_generic_def(mod, typeinfo, generic_file)
-        located_file.write(auto_gen_msg)
-        write_located_def(typeinfo, located_file)
+    for filename, write in [
+        ("generic", write_ast_def),
+        ("fold", write_fold_def),
+        ("located", write_located_def),
+        ("visitor", write_visitor_def),
+    ]:
+        with (ast_dir / f"{filename}.rs").open("w") as f:
+            f.write(auto_gen_msg)
+            write(mod, typeinfo, f)
 
     with module_filename.open("w") as module_file:
         module_file.write(auto_gen_msg)
         write_ast_mod(mod, typeinfo, module_file)
 
-    print(f"{generic_filename}, {located_filename}, {module_filename} regenerated.")
+    print(f"{ast_dir}, {module_filename} regenerated.")
 
 
 if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("input_file", type=Path)
-    parser.add_argument("-G", "--generic-file", type=Path, required=True)
-    parser.add_argument("-L", "--located-file", type=Path, required=True)
+    parser.add_argument("-A", "--ast-dir", type=Path, required=True)
     parser.add_argument("-M", "--module-file", type=Path, required=True)
     parser.add_argument("-d", "--dump-module", action="store_true")
 
     args = parser.parse_args()
     main(
         args.input_file,
-        args.generic_file,
-        args.located_file,
+        args.ast_dir,
         args.module_file,
         args.dump_module,
     )