Fix fold

ast/asdl_rs.py

@@ -106,7 +106,7 @@ class TypeInfo:
     def __repr__(self):
         return f"<TypeInfo: {self.name}>"
 
-    def needs_cfg(self, typeinfo):
+    def no_cfg(self, typeinfo):
         if self.product:
             return self.has_attributes
         elif self.enum_name:
@@ -157,7 +157,8 @@ class TypeInfoMixin:
         return self.type_info[typ].has_user_data
 
     def apply_generics(self, typ, *generics):
-        if self.has_user_data(typ):
+        needs_generics = not self.type_info[typ].is_simple
+        if needs_generics:
             return [f"<{g}>" for g in generics]
         else:
             return ["" for g in generics]
@@ -284,7 +285,6 @@ class StructVisitor(EmitVisitor):
             self.emit('#[cfg(not(feature = "more-attributes"))]', depth + 1)
             self.emit("pub custom: std::marker::PhantomData<U>,", depth + 1)
 
-
     def simple_sum(self, sum, name, depth):
         rust_name = rust_type_name(name)
         self.emit_attrs(depth)
@@ -314,9 +314,7 @@ class StructVisitor(EmitVisitor):
                     f'#[is(name = "{rust_field_name(t.name)}_{rust_name.lower()}")]',
                     depth + 1,
                 )
-            self.emit(
-                f"{t.name}({rust_name}{t.name}<U>),", depth + 1
-            )
+            self.emit(f"{t.name}({rust_name}{t.name}<U>),", depth + 1)
         self.emit("}", depth)
         self.emit("", depth)
 
@@ -327,9 +325,11 @@ class StructVisitor(EmitVisitor):
         for f in t.fields:
             self.visit(f, sum_type_info, "pub ", depth + 1, t.name)
 
-        assert sum_type_info.has_attributes == self.type_info[t.name].needs_cfg(self.type_info)
+        assert sum_type_info.has_attributes == self.type_info[t.name].no_cfg(
+            self.type_info
+        )
         self.emit_custom(sum_type_info.has_attributes, depth)
-            
+
         self.emit("}", depth)
         self.emit(
             textwrap.dedent(
@@ -398,7 +398,7 @@ class StructVisitor(EmitVisitor):
         self.emit(f"pub struct {product_name}<U> {{", depth)
         for f in product.fields:
             self.visit(f, type_info, "pub ", depth + 1)
-        assert bool(product.attributes) == type_info.needs_cfg(self.type_info)
+        assert bool(product.attributes) == type_info.no_cfg(self.type_info)
         self.emit_custom(product.attributes, depth + 1)
         self.emit("}", depth)
 
@@ -418,7 +418,17 @@ class FoldTraitDefVisitor(EmitVisitor):
         self.emit("type TargetU;", depth + 1)
         self.emit("type Error;", depth + 1)
         self.emit(
-            "fn map_user(&mut self, user: U) -> Result<Self::TargetU, Self::Error>;",
+            """
+            fn map_user(&mut self, user: U) -> Result<Self::TargetU, Self::Error>;
+            #[cfg(feature = "more-attributes")]
+            fn map_user_cfg(&mut self, user: U) -> Result<Self::TargetU, Self::Error> {
+                self.map_user(user)
+            }
+            #[cfg(not(feature = "more-attributes"))]
+            fn map_user_cfg(&mut self, _user: U) -> Result<std::marker::PhantomData<Self::TargetU>, Self::Error> {
+                Ok(std::marker::PhantomData)
+            }
+            """,
             depth + 1,
         )
         self.emit(
@@ -446,15 +456,6 @@ class FoldTraitDefVisitor(EmitVisitor):
 
 class FoldImplVisitor(EmitVisitor):
     def visitModule(self, mod, depth):
-        self.emit(
-            "fn fold_attributed<U, F: Fold<U> + ?Sized, T, MT>(folder: &mut F, node: Attributed<T, U>, f: impl FnOnce(&mut F, T) -> Result<MT, F::Error>) -> Result<Attributed<MT, F::TargetU>, F::Error> where T: Ranged{",
-            depth,
-        )
-        self.emit(
-            "let node = folder.map_attributed(node)?; Ok(Attributed { custom: node.custom, node: f(folder, node.node)? })",
-            depth + 1,
-        )
-        self.emit("}", depth)
         for dfn in mod.dfns:
             self.visit(dfn, depth)
 
@@ -483,25 +484,28 @@ class FoldImplVisitor(EmitVisitor):
             f"pub fn fold_{name}<U, F: Fold<U> + ?Sized>(#[allow(unused)] folder: &mut F, node: {enum_name}{apply_u}) -> Result<{enum_name}{apply_target_u}, F::Error> {{",
             depth,
         )
-        if type_info.has_attributes:
-            self.emit("fold_attributed(folder, node, |folder, node| {", depth)
+
+        if simple:
+            self.emit("Ok(node) }", depth + 1)
+            return
 
         self.emit("match node {", depth + 1)
         for cons in sum.types:
-            fields_pattern = self.make_pattern(
-                enum_name, cons.name, cons.fields, simple
-            )
+            fields_pattern = self.make_pattern(enum_name, cons.name, cons.fields)
             self.emit(
                 f"{fields_pattern[0]} {{ {fields_pattern[1]}}} {fields_pattern[2]} => {{",
                 depth + 2,
             )
+            if not type_info.has_attributes:
+                self.emit('#[cfg(not(feature = "more-attributes"))]', depth + 3)
+                self.emit("let custom = std::marker::PhantomData;", depth + 3)
+                self.emit('#[cfg(feature = "more-attributes")]', depth + 3)
+            self.emit("let custom = folder.map_user(custom)?;", depth + 3)
             self.gen_construction(
-                fields_pattern[0], cons.fields, fields_pattern[2], depth + 3, simple
+                fields_pattern[0], cons.fields, fields_pattern[2], depth + 3
             )
             self.emit("}", depth + 2)
         self.emit("}", depth + 1)
-        if type_info.has_attributes:
-            self.emit("})", depth)
         self.emit("}", depth)
 
     def visitProduct(self, product, name, depth):
@@ -525,41 +529,38 @@ class FoldImplVisitor(EmitVisitor):
             f"pub fn fold_{name}<U, F: Fold<U> + ?Sized>(#[allow(unused)] folder: &mut F, node: {struct_name}{apply_u}) -> Result<{struct_name}{apply_target_u}, F::Error> {{",
             depth,
         )
-        if has_attributes:
-            self.emit("fold_attributed(folder, node, |folder, node| {", depth)
-            rust_name = struct_name + "Data"
-        else:
-            rust_name = struct_name
-        fields_pattern = self.make_pattern(rust_name, struct_name, product.fields, False)
-        self.emit(f"let {rust_name} {{ {fields_pattern[1]} }} = node;", depth + 1)
-        self.gen_construction(rust_name, product.fields, "", depth + 1, False)
-        if has_attributes:
-            self.emit("})", depth)
+
+        fields_pattern = self.make_pattern(struct_name, struct_name, product.fields)
+        self.emit(f"let {struct_name} {{ {fields_pattern[1]} }} = node;", depth + 1)
+
+        if not has_attributes:
+            self.emit('#[cfg(not(feature = "more-attributes"))]', depth + 3)
+            self.emit("let custom = std::marker::PhantomData;", depth + 3)
+            self.emit('#[cfg(feature = "more-attributes")]', depth + 3)
+        self.emit("let custom = folder.map_user(custom)?;", depth + 3)
+
+        self.gen_construction(struct_name, product.fields, "", depth + 1)
+
         self.emit("}", depth)
 
-    def make_pattern(self, rust_name, fieldname: str, fields, simple_sum: bool):
-        if fields or not simple_sum:
-            header = f"{rust_name}::{fieldname}({rust_name}{fieldname}"
-            footer = ")"
-        else:
-            header = f"{rust_name}::{fieldname}"
-            footer = ""
+    def make_pattern(self, rust_name, fieldname: str, fields):
+        header = f"{rust_name}::{fieldname}({rust_name}{fieldname}"
+        footer = ")"
 
         body = ",".join(rust_field(f.name) for f in fields)
-
-        if not simple_sum:
-            body = f"range, {body}"
+        if body:
+            body += ","
+        body += "custom"
 
         return header, body, footer
 
-    def gen_construction(self, header, fields, footer, depth, simple_sum: bool):
+    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("custom,", depth + 1)
 
-        if not simple_sum:
-            self.emit("range", depth + 1)
         self.emit(f"}}{footer})", depth)
 
 
@@ -668,15 +669,13 @@ class VisitorTraitDefVisitor(StructVisitor):
         if not sum.attributes:
             return
 
-        rust_name = enum_name = rust_type_name(name)
-        if sum.attributes:
-            rust_name = enum_name + "Kind"
+        enum_name = rust_type_name(name)
         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, rust_name, t, depth + 1)
+            self.visit_match_for_type(name, enum_name, t, depth + 1)
         self.emit("}", depth)
         depth -= 1
         self.emit("}", depth)
@@ -837,8 +836,6 @@ class TraitImplVisitor(EmitVisitor):
 
     def visitProduct(self, product, name, depth):
         struct_name = rust_type_name(name)
-        if product.attributes:
-            struct_name += "Data"
 
         self.emit(f"impl NamedNode for ast::located::{struct_name} {{", depth)
         self.emit(f"const NAME: &'static str = {json.dumps(name)};", depth + 1)
@@ -959,9 +956,11 @@ def write_ast_def(mod, type_info, f):
 
 
 def write_fold_def(mod, type_info, f):
-    f.write("""
+    f.write(
+        """
 use crate::generic::Custom;
-""")
+"""
+    )
     FoldModuleVisitor(f, type_info).visit(mod)
 
 
@@ -972,15 +971,17 @@ def write_visitor_def(mod, type_info, f):
 def write_ranged_def(mod, type_info, f):
     for info in type_info.values():
         if not info.is_simple:
-            if info.needs_cfg:
+            if info.no_cfg:
                 f.write('#[cfg(feature = "more-attributes")]')
-            f.write(f"""
+            f.write(
+                f"""
             impl Ranged for {info.rust_sum_name} {{
                 fn range(&self) -> TextRange {{
                     self.custom
                 }}
             }}
-            """)
+            """
+            )
             generics = "::<TextRange>"
         else:
             generics = ""
@@ -988,18 +989,21 @@ def write_ranged_def(mod, type_info, f):
             f"pub type {info.rust_sum_name} = crate::generic::{info.rust_sum_name}{generics};\n"
         )
 
+
 def write_located_def(mod, type_info, f):
     for info in type_info.values():
         if not info.is_simple:
-            if info.needs_cfg:
+            if info.no_cfg:
                 f.write('#[cfg(feature = "more-attributes")]')
-            f.write(f"""
+            f.write(
+                f"""
             impl Located for {info.rust_sum_name} {{
                 fn range(&self) -> SourceRange {{
                     self.custom
                 }}
             }}
-            """)
+            """
+            )
             generics = "::<SourceRange>"
         else:
             generics = ""
@@ -1007,6 +1011,7 @@ def write_located_def(mod, type_info, f):
             f"pub type {info.rust_sum_name} = crate::generic::{info.rust_sum_name}{generics};\n"
         )
 
+
 def write_ast_mod(mod, type_info, f):
     f.write(
         textwrap.dedent(

ast/src/gen/visitor.rs

@@ -7,33 +7,33 @@ pub trait Visitor<U = ()> {
     }
     fn generic_visit_stmt(&mut self, node: Stmt) {
         match node.node {
-            StmtKind::FunctionDef(data) => self.visit_stmt_FunctionDef(data),
-            StmtKind::AsyncFunctionDef(data) => self.visit_stmt_AsyncFunctionDef(data),
-            StmtKind::ClassDef(data) => self.visit_stmt_ClassDef(data),
-            StmtKind::Return(data) => self.visit_stmt_Return(data),
-            StmtKind::Delete(data) => self.visit_stmt_Delete(data),
-            StmtKind::Assign(data) => self.visit_stmt_Assign(data),
-            StmtKind::AugAssign(data) => self.visit_stmt_AugAssign(data),
-            StmtKind::AnnAssign(data) => self.visit_stmt_AnnAssign(data),
-            StmtKind::For(data) => self.visit_stmt_For(data),
-            StmtKind::AsyncFor(data) => self.visit_stmt_AsyncFor(data),
-            StmtKind::While(data) => self.visit_stmt_While(data),
-            StmtKind::If(data) => self.visit_stmt_If(data),
-            StmtKind::With(data) => self.visit_stmt_With(data),
-            StmtKind::AsyncWith(data) => self.visit_stmt_AsyncWith(data),
-            StmtKind::Match(data) => self.visit_stmt_Match(data),
-            StmtKind::Raise(data) => self.visit_stmt_Raise(data),
-            StmtKind::Try(data) => self.visit_stmt_Try(data),
-            StmtKind::TryStar(data) => self.visit_stmt_TryStar(data),
-            StmtKind::Assert(data) => self.visit_stmt_Assert(data),
-            StmtKind::Import(data) => self.visit_stmt_Import(data),
-            StmtKind::ImportFrom(data) => self.visit_stmt_ImportFrom(data),
-            StmtKind::Global(data) => self.visit_stmt_Global(data),
-            StmtKind::Nonlocal(data) => self.visit_stmt_Nonlocal(data),
-            StmtKind::Expr(data) => self.visit_stmt_Expr(data),
-            StmtKind::Pass(data) => self.visit_stmt_Pass(data),
-            StmtKind::Break(data) => self.visit_stmt_Break(data),
-            StmtKind::Continue(data) => self.visit_stmt_Continue(data),
+            Stmt::FunctionDef(data) => self.visit_stmt_FunctionDef(data),
+            Stmt::AsyncFunctionDef(data) => self.visit_stmt_AsyncFunctionDef(data),
+            Stmt::ClassDef(data) => self.visit_stmt_ClassDef(data),
+            Stmt::Return(data) => self.visit_stmt_Return(data),
+            Stmt::Delete(data) => self.visit_stmt_Delete(data),
+            Stmt::Assign(data) => self.visit_stmt_Assign(data),
+            Stmt::AugAssign(data) => self.visit_stmt_AugAssign(data),
+            Stmt::AnnAssign(data) => self.visit_stmt_AnnAssign(data),
+            Stmt::For(data) => self.visit_stmt_For(data),
+            Stmt::AsyncFor(data) => self.visit_stmt_AsyncFor(data),
+            Stmt::While(data) => self.visit_stmt_While(data),
+            Stmt::If(data) => self.visit_stmt_If(data),
+            Stmt::With(data) => self.visit_stmt_With(data),
+            Stmt::AsyncWith(data) => self.visit_stmt_AsyncWith(data),
+            Stmt::Match(data) => self.visit_stmt_Match(data),
+            Stmt::Raise(data) => self.visit_stmt_Raise(data),
+            Stmt::Try(data) => self.visit_stmt_Try(data),
+            Stmt::TryStar(data) => self.visit_stmt_TryStar(data),
+            Stmt::Assert(data) => self.visit_stmt_Assert(data),
+            Stmt::Import(data) => self.visit_stmt_Import(data),
+            Stmt::ImportFrom(data) => self.visit_stmt_ImportFrom(data),
+            Stmt::Global(data) => self.visit_stmt_Global(data),
+            Stmt::Nonlocal(data) => self.visit_stmt_Nonlocal(data),
+            Stmt::Expr(data) => self.visit_stmt_Expr(data),
+            Stmt::Pass(data) => self.visit_stmt_Pass(data),
+            Stmt::Break(data) => self.visit_stmt_Break(data),
+            Stmt::Continue(data) => self.visit_stmt_Continue(data),
         }
     }
     fn visit_stmt_FunctionDef(&mut self, node: StmtFunctionDef) {
@@ -346,33 +346,33 @@ pub trait Visitor<U = ()> {
     }
     fn generic_visit_expr(&mut self, node: Expr) {
         match node.node {
-            ExprKind::BoolOp(data) => self.visit_expr_BoolOp(data),
-            ExprKind::NamedExpr(data) => self.visit_expr_NamedExpr(data),
-            ExprKind::BinOp(data) => self.visit_expr_BinOp(data),
-            ExprKind::UnaryOp(data) => self.visit_expr_UnaryOp(data),
-            ExprKind::Lambda(data) => self.visit_expr_Lambda(data),
-            ExprKind::IfExp(data) => self.visit_expr_IfExp(data),
-            ExprKind::Dict(data) => self.visit_expr_Dict(data),
-            ExprKind::Set(data) => self.visit_expr_Set(data),
-            ExprKind::ListComp(data) => self.visit_expr_ListComp(data),
-            ExprKind::SetComp(data) => self.visit_expr_SetComp(data),
-            ExprKind::DictComp(data) => self.visit_expr_DictComp(data),
-            ExprKind::GeneratorExp(data) => self.visit_expr_GeneratorExp(data),
-            ExprKind::Await(data) => self.visit_expr_Await(data),
-            ExprKind::Yield(data) => self.visit_expr_Yield(data),
-            ExprKind::YieldFrom(data) => self.visit_expr_YieldFrom(data),
-            ExprKind::Compare(data) => self.visit_expr_Compare(data),
-            ExprKind::Call(data) => self.visit_expr_Call(data),
-            ExprKind::FormattedValue(data) => self.visit_expr_FormattedValue(data),
-            ExprKind::JoinedStr(data) => self.visit_expr_JoinedStr(data),
-            ExprKind::Constant(data) => self.visit_expr_Constant(data),
-            ExprKind::Attribute(data) => self.visit_expr_Attribute(data),
-            ExprKind::Subscript(data) => self.visit_expr_Subscript(data),
-            ExprKind::Starred(data) => self.visit_expr_Starred(data),
-            ExprKind::Name(data) => self.visit_expr_Name(data),
-            ExprKind::List(data) => self.visit_expr_List(data),
-            ExprKind::Tuple(data) => self.visit_expr_Tuple(data),
-            ExprKind::Slice(data) => self.visit_expr_Slice(data),
+            Expr::BoolOp(data) => self.visit_expr_BoolOp(data),
+            Expr::NamedExpr(data) => self.visit_expr_NamedExpr(data),
+            Expr::BinOp(data) => self.visit_expr_BinOp(data),
+            Expr::UnaryOp(data) => self.visit_expr_UnaryOp(data),
+            Expr::Lambda(data) => self.visit_expr_Lambda(data),
+            Expr::IfExp(data) => self.visit_expr_IfExp(data),
+            Expr::Dict(data) => self.visit_expr_Dict(data),
+            Expr::Set(data) => self.visit_expr_Set(data),
+            Expr::ListComp(data) => self.visit_expr_ListComp(data),
+            Expr::SetComp(data) => self.visit_expr_SetComp(data),
+            Expr::DictComp(data) => self.visit_expr_DictComp(data),
+            Expr::GeneratorExp(data) => self.visit_expr_GeneratorExp(data),
+            Expr::Await(data) => self.visit_expr_Await(data),
+            Expr::Yield(data) => self.visit_expr_Yield(data),
+            Expr::YieldFrom(data) => self.visit_expr_YieldFrom(data),
+            Expr::Compare(data) => self.visit_expr_Compare(data),
+            Expr::Call(data) => self.visit_expr_Call(data),
+            Expr::FormattedValue(data) => self.visit_expr_FormattedValue(data),
+            Expr::JoinedStr(data) => self.visit_expr_JoinedStr(data),
+            Expr::Constant(data) => self.visit_expr_Constant(data),
+            Expr::Attribute(data) => self.visit_expr_Attribute(data),
+            Expr::Subscript(data) => self.visit_expr_Subscript(data),
+            Expr::Starred(data) => self.visit_expr_Starred(data),
+            Expr::Name(data) => self.visit_expr_Name(data),
+            Expr::List(data) => self.visit_expr_List(data),
+            Expr::Tuple(data) => self.visit_expr_Tuple(data),
+            Expr::Slice(data) => self.visit_expr_Slice(data),
         }
     }
     fn visit_expr_BoolOp(&mut self, node: ExprBoolOp) {
@@ -690,7 +690,7 @@ pub trait Visitor<U = ()> {
     }
     fn generic_visit_excepthandler(&mut self, node: Excepthandler) {
         match node.node {
-            ExcepthandlerKind::ExceptHandler(data) => self.visit_excepthandler_ExceptHandler(data),
+            Excepthandler::ExceptHandler(data) => self.visit_excepthandler_ExceptHandler(data),
         }
     }
     fn visit_excepthandler_ExceptHandler(&mut self, node: ExcepthandlerExceptHandler) {
@@ -733,14 +733,14 @@ pub trait Visitor<U = ()> {
     }
     fn generic_visit_pattern(&mut self, node: Pattern) {
         match node.node {
-            PatternKind::MatchValue(data) => self.visit_pattern_MatchValue(data),
-            PatternKind::MatchSingleton(data) => self.visit_pattern_MatchSingleton(data),
-            PatternKind::MatchSequence(data) => self.visit_pattern_MatchSequence(data),
-            PatternKind::MatchMapping(data) => self.visit_pattern_MatchMapping(data),
-            PatternKind::MatchClass(data) => self.visit_pattern_MatchClass(data),
-            PatternKind::MatchStar(data) => self.visit_pattern_MatchStar(data),
-            PatternKind::MatchAs(data) => self.visit_pattern_MatchAs(data),
-            PatternKind::MatchOr(data) => self.visit_pattern_MatchOr(data),
+            Pattern::MatchValue(data) => self.visit_pattern_MatchValue(data),
+            Pattern::MatchSingleton(data) => self.visit_pattern_MatchSingleton(data),
+            Pattern::MatchSequence(data) => self.visit_pattern_MatchSequence(data),
+            Pattern::MatchMapping(data) => self.visit_pattern_MatchMapping(data),
+            Pattern::MatchClass(data) => self.visit_pattern_MatchClass(data),
+            Pattern::MatchStar(data) => self.visit_pattern_MatchStar(data),
+            Pattern::MatchAs(data) => self.visit_pattern_MatchAs(data),
+            Pattern::MatchOr(data) => self.visit_pattern_MatchOr(data),
         }
     }
     fn visit_pattern_MatchValue(&mut self, node: PatternMatchValue) {

ast/src/located.rs

@@ -5,11 +5,11 @@ pub trait Located {
     fn range(&self) -> SourceRange;
 
     fn location(&self) -> SourceLocation {
-        self.range().start()
+        self.range().start
     }
 
-    fn end_location(&self) -> SourceLocation {
-        self.range().end()
+    fn end_location(&self) -> Option<SourceLocation> {
+        self.range().end
     }
 }
 

ast/src/source_locator.rs

@@ -1,42 +1,15 @@
-use crate::builtin::Attributed;
-use crate::Ranged;
-use rustpython_parser_core::source_code::{SourceLocation, SourceLocator, SourceRange};
+use rustpython_parser_core::{
+    source_code::{SourceLocator, SourceRange},
+    text_size::TextRange,
+};
 
-impl crate::fold::Fold<()> for SourceLocator<'_> {
+impl crate::fold::Fold<TextRange> for SourceLocator<'_> {
     type TargetU = SourceRange;
     type Error = std::convert::Infallible;
 
-    #[cold]
-    fn map_user(&mut self, _user: ()) -> Result<Self::TargetU, Self::Error> {
-        unreachable!("implemented map_attributed");
-    }
-
-    fn map_attributed<T>(
-        &mut self,
-        node: Attributed<T, ()>,
-    ) -> Result<Attributed<T, Self::TargetU>, Self::Error>
-    where
-        T: Ranged,
-    {
-        let start = self.locate(node.start());
-        let end = self.locate(node.end());
-        Ok(Attributed {
-            custom: (start..end).into(),
-            node: node.node,
-        })
-    }
-}
-
-impl<T> Attributed<T, SourceRange> {
-    /// Returns the absolute start position of the node from the beginning of the document.
-    #[inline]
-    pub const fn location(&self) -> SourceLocation {
-        self.custom.start
-    }
-
-    /// Returns the absolute position at which the node ends in the source document.
-    #[inline]
-    pub const fn end_location(&self) -> Option<SourceLocation> {
-        self.custom.end
+    fn map_user(&mut self, user: TextRange) -> Result<Self::TargetU, Self::Error> {
+        let start = self.locate(user.start());
+        let end = self.locate(user.end());
+        Ok((start..end).into())
     }
 }