Refactor ast to hold data as seperated type

Author: youknowone

ast/asdl_rs.py

@@ -131,6 +131,9 @@ def visitSum(self, sum, name):
         if is_simple(sum):
             info.has_userdata = False
         else:
+            for t in sum.types:
+                self.typeinfo[t.name] = TypeInfo(t.name)
+                self.add_children(t.name, t.fields)
             if len(sum.types) > 1:
                 info.boxed = True
             if sum.attributes:
@@ -205,16 +208,49 @@ def simple_sum(self, sum, name, depth):
 
     def sum_with_constructors(self, sum, name, depth):
         typeinfo = self.typeinfo[name]
-        generics, generics_applied = self.get_generics(name, "U = ()", "U")
         enumname = rustname = get_rust_type(name)
         # all the attributes right now are for location, so if it has attrs we
         # can just wrap it in Located<>
         if sum.attributes:
             enumname = rustname + "Kind"
+
+        for t in sum.types:
+            if not t.fields:
+                continue
+            self.emit_attrs(depth)
+            self.typeinfo[t] = TypeInfo(t)
+            t_generics, t_generics_applied = self.get_generics(t.name, "U = ()", "U")
+            payload_name = f"{rustname}{t.name}"
+            self.emit(f"pub struct {payload_name}{t_generics} {{", depth)
+            for f in t.fields:
+                self.visit(f, typeinfo, "pub ", depth + 1, t.name)
+            self.emit("}", depth)
+            self.emit(
+                textwrap.dedent(
+                    f"""
+            impl{t_generics_applied} From<{payload_name}{t_generics_applied}> for {enumname}{t_generics_applied} {{
+                fn from(payload: {payload_name}{t_generics_applied}) -> Self {{
+                    {enumname}::{t.name}(payload)
+                }}
+            }}
+            """
+                ),
+                depth,
+            )
+
+        generics, generics_applied = self.get_generics(name, "U = ()", "U")
         self.emit_attrs(depth)
         self.emit(f"pub enum {enumname}{generics} {{", depth)
         for t in sum.types:
-            self.visit(t, typeinfo, depth + 1)
+            if t.fields:
+                t_generics, t_generics_applied = self.get_generics(
+                    t.name, "U = ()", "U"
+                )
+                self.emit(
+                    f"{t.name}({rustname}{t.name}{t_generics_applied}),", depth + 1
+                )
+            else:
+                self.emit(f"{t.name},", depth + 1)
         self.emit("}", depth)
         if sum.attributes:
             self.emit(
@@ -238,13 +274,18 @@ def visitField(self, field, parent, vis, depth, constructor=None):
         if fieldtype and fieldtype.has_userdata:
             typ = f"{typ}<U>"
         # don't box if we're doing Vec<T>, but do box if we're doing Vec<Option<Box<T>>>
-        if fieldtype and fieldtype.boxed and (not (parent.product or field.seq) or field.opt):
+        if (
+            fieldtype
+            and fieldtype.boxed
+            and (not (parent.product or field.seq) or field.opt)
+        ):
             typ = f"Box<{typ}>"
         if field.opt or (
             # When a dictionary literal contains dictionary unpacking (e.g., `{**d}`),
             # the expression to be unpacked goes in `values` with a `None` at the corresponding
             # position in `keys`. To handle this, the type of `keys` needs to be `Option<Vec<T>>`.
-            constructor == "Dict" and field.name == "keys"
+            constructor == "Dict"
+            and field.name == "keys"
         ):
             typ = f"Option<{typ}>"
         if field.seq:
@@ -344,14 +385,21 @@ def visitSum(self, sum, name, depth):
         )
         if is_located:
             self.emit("fold_located(folder, node, |folder, node| {", depth)
-            enumname += "Kind"
+            rustname = enumname + "Kind"
+        else:
+            rustname = enumname
         self.emit("match node {", depth + 1)
         for cons in sum.types:
-            fields_pattern = self.make_pattern(cons.fields)
+            fields_pattern = self.make_pattern(
+                enumname, rustname, cons.name, cons.fields
+            )
             self.emit(
-                f"{enumname}::{cons.name} {{ {fields_pattern} }} => {{", depth + 2
+                f"{fields_pattern[0]} {{ {fields_pattern[1]} }} {fields_pattern[2]} => {{",
+                depth + 2,
+            )
+            self.gen_construction(
+                fields_pattern[0], cons.fields, fields_pattern[2], depth + 3
             )
-            self.gen_construction(f"{enumname}::{cons.name}", cons.fields, depth + 3)
             self.emit("}", depth + 2)
         self.emit("}", depth + 1)
         if is_located:
@@ -381,23 +429,33 @@ def visitProduct(self, product, name, depth):
         )
         if is_located:
             self.emit("fold_located(folder, node, |folder, node| {", depth)
-            structname += "Data"
-        fields_pattern = self.make_pattern(product.fields)
-        self.emit(f"let {structname} {{ {fields_pattern} }} = node;", depth + 1)
-        self.gen_construction(structname, product.fields, depth + 1)
+            rustname = structname + "Data"
+        else:
+            rustname = structname
+        fields_pattern = self.make_pattern(rustname, structname, None, product.fields)
+        self.emit(f"let {rustname} {{ {fields_pattern[1]} }} = node;", depth + 1)
+        self.gen_construction(rustname, product.fields, "", depth + 1)
         if is_located:
             self.emit("})", depth)
         self.emit("}", depth)
 
-    def make_pattern(self, fields):
-        return ",".join(rust_field(f.name) for f in fields)
+    def make_pattern(self, rustname, pyname, fieldname, fields):
+        if fields:
+            header = f"{pyname}::{fieldname}({rustname}{fieldname}"
+            footer = ")"
+        else:
+            header = f"{pyname}::{fieldname}"
+            footer = ""
+
+        body = ",".join(rust_field(f.name) for f in fields)
+        return header, body, footer
 
-    def gen_construction(self, cons_path, fields, depth):
-        self.emit(f"Ok({cons_path} {{", depth)
+    def gen_construction(self, header, fields, footer, depth):
+        self.emit(f"Ok({header} {{", depth)
         for field in fields:
             name = rust_field(field.name)
             self.emit(f"{name}: Foldable::fold({name}, folder)?,", depth + 1)
-        self.emit("})", depth)
+        self.emit(f"}}{footer})", depth)
 
 
 class FoldModuleVisitor(TypeInfoEmitVisitor):