From d07dca789f7cb66e59d9d803a15f937ebbea7334 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Sun, 4 Apr 2021 17:52:19 -0500 Subject: [PATCH 1/3] Add the ast::fold::Fold trait, replacing MapAst --- ast/Cargo.toml | 5 + ast/asdl_rs.py | 154 +++++---- ast/src/ast_gen.rs | 738 ++++++++++++++++++++++++++++------------ ast/src/fold_helpers.rs | 72 ++++ ast/src/lib.rs | 6 +- ast/src/map_ast.rs | 66 ---- 6 files changed, 699 insertions(+), 342 deletions(-) create mode 100644 ast/src/fold_helpers.rs delete mode 100644 ast/src/map_ast.rs diff --git a/ast/Cargo.toml b/ast/Cargo.toml index 74dfbded5..24aa0e124 100644 --- a/ast/Cargo.toml +++ b/ast/Cargo.toml @@ -4,5 +4,10 @@ version = "0.1.0" authors = ["RustPython Team"] edition = "2018" +[features] +default = ["constant-optimization", "fold"] +constant-optimization = ["fold"] +fold = [] + [dependencies] num-bigint = "0.4.0" diff --git a/ast/asdl_rs.py b/ast/asdl_rs.py index d1f6dd7a0..95a4ad8b2 100755 --- a/ast/asdl_rs.py +++ b/ast/asdl_rs.py @@ -154,13 +154,23 @@ def rust_field(field_name): else: return field_name -class StructVisitor(EmitVisitor): - """Visitor to generate typedefs for AST.""" - +class TypeInfoEmitVisitor(EmitVisitor): def __init__(self, file, typeinfo): self.typeinfo = typeinfo super().__init__(file) + def has_userdata(self, typ): + return self.typeinfo[typ].has_userdata + + def get_generics(self, typ, *generics): + if self.has_userdata(typ): + return [f"<{g}>" for g in generics] + else: + return ["" for g in generics] + +class StructVisitor(TypeInfoEmitVisitor): + """Visitor to generate typedefs for AST.""" + def visitModule(self, mod): for dfn in mod.dfns: self.visit(dfn) @@ -188,7 +198,7 @@ class StructVisitor(EmitVisitor): def sum_with_constructors(self, sum, name, depth): typeinfo = self.typeinfo[name] - generics, generics_applied = self.get_generics(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<> @@ -229,7 +239,7 @@ class StructVisitor(EmitVisitor): def visitProduct(self, product, name, depth): typeinfo = self.typeinfo[name] - generics, generics_applied = self.get_generics(name) + generics, generics_applied = self.get_generics(name, "U = ()", "U") dataname = rustname = get_rust_type(name) if product.attributes: dataname = rustname + "Data" @@ -243,71 +253,86 @@ class StructVisitor(EmitVisitor): self.emit(f"pub type {rustname} = Located<{dataname}{generics_applied}, U>;", depth); self.emit("", depth) - def get_generics(self, typ): - if self.typeinfo[typ].has_userdata: - return "", "" - else: - return "", "" -class MapAstImplVisitor(EmitVisitor): - def __init__(self, file, typeinfo): - self.typeinfo = typeinfo - super().__init__(file) - - def get_generics(self, typ): - if self.typeinfo[typ].has_userdata: - return "", "" - else: - return "", "" - - def visitModule(self, mod): +class FoldTraitDefVisitor(TypeInfoEmitVisitor): + def visitModule(self, mod, depth): + self.emit("pub trait Fold {", depth) + self.emit("type TargetU;", depth + 1) + self.emit("type Error;", depth + 1) + self.emit("fn map_user(&mut self, user: U) -> Result;", depth + 2) for dfn in mod.dfns: - self.visit(dfn) + self.visit(dfn, depth + 2) + self.emit("}", depth) + + def visitType(self, type, depth): + name = type.name + apply_u, apply_target_u = self.get_generics(name, "U", "Self::TargetU") + enumname = get_rust_type(name) + self.emit(f"fn fold_{name}(&mut self, node: {enumname}{apply_u}) -> Result<{enumname}{apply_target_u}, Self::Error> {{", depth) + self.emit(f"fold_{name}(self, node)", depth + 1) + self.emit("}", depth) + + +class FoldImplVisitor(TypeInfoEmitVisitor): + def visitModule(self, mod, depth): + self.emit("fn fold_located + ?Sized, T, MT>(folder: &mut F, node: Located, f: impl FnOnce(&mut F, T) -> Result) -> Result, F::Error> {", depth) + self.emit("Ok(Located { custom: folder.map_user(node.custom)?, location: node.location, node: f(folder, node.node)? })", depth + 1) + self.emit("}", depth) + for dfn in mod.dfns: + self.visit(dfn, depth) def visitType(self, type, depth=0): self.visit(type.value, type.name, depth) def visitSum(self, sum, name, depth): - apply_t, apply_u = self.get_generics(name) + apply_t, apply_u, apply_target_u = self.get_generics(name, "T", "U", "F::TargetU") enumname = get_rust_type(name) - if sum.attributes: - enumname += "Kind" + is_located = bool(sum.attributes) - if not apply_t: - self.emit(f"no_user!({enumname});", depth) - return - - - self.emit(f"impl MapAst for {enumname}{apply_t} {{", depth) + self.emit(f"impl Foldable for {enumname}{apply_t} {{", depth) self.emit(f"type Mapped = {enumname}{apply_u};", depth + 1) - self.emit("fn try_map_ast Result>(self, f: &mut F) -> Result {", depth + 1) - self.emit("match self {", depth + 2) - for cons in sum.types: - fields_pattern = self.make_pattern(cons.fields) - self.emit(f"{enumname}::{cons.name} {{ {fields_pattern} }} => {{", depth + 3) - self.gen_construction(f"{enumname}::{cons.name}", cons.fields, depth + 4) - self.emit("}", depth + 3) - self.emit("}", depth + 2) + self.emit("fn fold + ?Sized>(self, folder: &mut F) -> Result {", depth + 1) + self.emit(f"fold_{name}(folder, self)", depth + 2) self.emit("}", depth + 1) self.emit("}", depth) - def visitProduct(self, product, name, depth): - apply_t, apply_u = self.get_generics(name) - structname = get_rust_type(name) - if product.attributes: - structname += "Data" - - if not apply_t: - self.emit(f"no_user!({structname});", depth) - return - - self.emit(f"impl MapAst for {structname}{apply_t} {{", depth) - self.emit(f"type Mapped = {structname}{apply_u};", depth + 1) - self.emit("fn try_map_ast Result>(self, f: &mut F) -> Result {", depth + 1) - fields_pattern = self.make_pattern(product.fields) - self.emit(f"let {structname} {{ {fields_pattern} }} = self;", depth + 2) - self.gen_construction(structname, product.fields, depth + 2) + self.emit(f"pub fn fold_{name} + ?Sized>(#[allow(unused)] folder: &mut F, node: {enumname}{apply_u}) -> Result<{enumname}{apply_target_u}, F::Error> {{", depth) + if is_located: + self.emit("fold_located(folder, node, |folder, node| {", depth) + enumname += "Kind" + self.emit("match node {", depth + 1) + for cons in sum.types: + fields_pattern = self.make_pattern(cons.fields) + self.emit(f"{enumname}::{cons.name} {{ {fields_pattern} }} => {{", depth + 2) + self.gen_construction(f"{enumname}::{cons.name}", cons.fields, depth + 3) + self.emit("}", depth + 2) self.emit("}", depth + 1) + if is_located: + self.emit("})", depth) + self.emit("}", depth) + + + def visitProduct(self, product, name, depth): + apply_t, apply_u, apply_target_u = self.get_generics(name, "T", "U", "F::TargetU") + structname = get_rust_type(name) + is_located = bool(product.attributes) + + self.emit(f"impl Foldable for {structname}{apply_t} {{", depth) + self.emit(f"type Mapped = {structname}{apply_u};", depth + 1) + self.emit("fn fold + ?Sized>(self, folder: &mut F) -> Result {", depth + 1) + self.emit(f"fold_{name}(folder, self)", depth + 2) + self.emit("}", depth + 1) + self.emit("}", depth) + + self.emit(f"pub fn fold_{name} + ?Sized>(#[allow(unused)] folder: &mut F, node: {structname}{apply_u}) -> Result<{structname}{apply_target_u}, F::Error> {{", 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) + if is_located: + self.emit("})", depth) self.emit("}", depth) def make_pattern(self, fields): @@ -317,10 +342,22 @@ class MapAstImplVisitor(EmitVisitor): self.emit(f"Ok({cons_path} {{", depth) for field in fields: name = rust_field(field.name) - self.emit(f"{name}: {name}.try_map_ast(f)?,", depth + 1) + self.emit(f"{name}: Foldable::fold({name}, folder)?,", depth + 1) self.emit("})", depth) +class FoldModuleVisitor(TypeInfoEmitVisitor): + 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("}", depth) + + class ClassDefVisitor(EmitVisitor): def visitModule(self, mod): @@ -508,7 +545,6 @@ class ChainOfVisitors: def write_ast_def(mod, typeinfo, f): f.write('pub use crate::location::Location;\n') f.write('pub use crate::constant::*;\n') - f.write('use crate::map_ast::MapAst;\n') f.write('\n') f.write('type Ident = String;\n') f.write('\n') @@ -527,7 +563,7 @@ def write_ast_def(mod, typeinfo, f): f.write('\n') c = ChainOfVisitors(StructVisitor(f, typeinfo), - MapAstImplVisitor(f, typeinfo)) + FoldModuleVisitor(f, typeinfo)) c.visit(mod) diff --git a/ast/src/ast_gen.rs b/ast/src/ast_gen.rs index 6ceb8414a..5d2aa6704 100644 --- a/ast/src/ast_gen.rs +++ b/ast/src/ast_gen.rs @@ -2,7 +2,6 @@ pub use crate::constant::*; pub use crate::location::Location; -use crate::map_ast::MapAst; type Ident = String; @@ -389,31 +388,133 @@ pub enum TypeIgnore { TypeIgnore { lineno: usize, tag: String }, } -impl MapAst for Mod { - type Mapped = Mod; - fn try_map_ast Result>(self, f: &mut F) -> Result { - match self { +#[cfg(feature = "fold")] +pub mod fold { + use super::*; + use crate::fold_helpers::Foldable; + pub trait Fold { + type TargetU; + type Error; + fn map_user(&mut self, user: U) -> Result; + fn fold_mod(&mut self, node: Mod) -> Result, Self::Error> { + fold_mod(self, node) + } + fn fold_stmt(&mut self, node: Stmt) -> Result, Self::Error> { + fold_stmt(self, node) + } + fn fold_expr(&mut self, node: Expr) -> Result, Self::Error> { + fold_expr(self, node) + } + fn fold_expr_context(&mut self, node: ExprContext) -> Result { + fold_expr_context(self, node) + } + fn fold_boolop(&mut self, node: Boolop) -> Result { + fold_boolop(self, node) + } + fn fold_operator(&mut self, node: Operator) -> Result { + fold_operator(self, node) + } + fn fold_unaryop(&mut self, node: Unaryop) -> Result { + fold_unaryop(self, node) + } + fn fold_cmpop(&mut self, node: Cmpop) -> Result { + fold_cmpop(self, node) + } + fn fold_comprehension( + &mut self, + node: Comprehension, + ) -> Result, Self::Error> { + fold_comprehension(self, node) + } + fn fold_excepthandler( + &mut self, + node: Excepthandler, + ) -> Result, Self::Error> { + fold_excepthandler(self, node) + } + fn fold_arguments( + &mut self, + node: Arguments, + ) -> Result, Self::Error> { + fold_arguments(self, node) + } + fn fold_arg(&mut self, node: Arg) -> Result, Self::Error> { + fold_arg(self, node) + } + fn fold_keyword( + &mut self, + node: Keyword, + ) -> Result, Self::Error> { + fold_keyword(self, node) + } + fn fold_alias(&mut self, node: Alias) -> Result { + fold_alias(self, node) + } + fn fold_withitem( + &mut self, + node: Withitem, + ) -> Result, Self::Error> { + fold_withitem(self, node) + } + fn fold_type_ignore(&mut self, node: TypeIgnore) -> Result { + fold_type_ignore(self, node) + } + } + fn fold_located + ?Sized, T, MT>( + folder: &mut F, + node: Located, + f: impl FnOnce(&mut F, T) -> Result, + ) -> Result, F::Error> { + Ok(Located { + custom: folder.map_user(node.custom)?, + location: node.location, + node: f(folder, node.node)?, + }) + } + impl Foldable for Mod { + type Mapped = Mod; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_mod(folder, self) + } + } + pub fn fold_mod + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Mod, + ) -> Result, F::Error> { + match node { Mod::Module { body, type_ignores } => Ok(Mod::Module { - body: body.try_map_ast(f)?, - type_ignores: type_ignores.try_map_ast(f)?, + body: Foldable::fold(body, folder)?, + type_ignores: Foldable::fold(type_ignores, folder)?, }), Mod::Interactive { body } => Ok(Mod::Interactive { - body: body.try_map_ast(f)?, + body: Foldable::fold(body, folder)?, }), Mod::Expression { body } => Ok(Mod::Expression { - body: body.try_map_ast(f)?, + body: Foldable::fold(body, folder)?, }), Mod::FunctionType { argtypes, returns } => Ok(Mod::FunctionType { - argtypes: argtypes.try_map_ast(f)?, - returns: returns.try_map_ast(f)?, + argtypes: Foldable::fold(argtypes, folder)?, + returns: Foldable::fold(returns, folder)?, }), } } -} -impl MapAst for StmtKind { - type Mapped = StmtKind; - fn try_map_ast Result>(self, f: &mut F) -> Result { - match self { + impl Foldable for Stmt { + type Mapped = Stmt; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_stmt(folder, self) + } + } + pub fn fold_stmt + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Stmt, + ) -> Result, F::Error> { + fold_located(folder, node, |folder, node| match node { StmtKind::FunctionDef { name, args, @@ -422,12 +523,12 @@ impl MapAst for StmtKind { returns, type_comment, } => Ok(StmtKind::FunctionDef { - name: name.try_map_ast(f)?, - args: args.try_map_ast(f)?, - body: body.try_map_ast(f)?, - decorator_list: decorator_list.try_map_ast(f)?, - returns: returns.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + name: Foldable::fold(name, folder)?, + args: Foldable::fold(args, folder)?, + body: Foldable::fold(body, folder)?, + decorator_list: Foldable::fold(decorator_list, folder)?, + returns: Foldable::fold(returns, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::AsyncFunctionDef { name, @@ -437,12 +538,12 @@ impl MapAst for StmtKind { returns, type_comment, } => Ok(StmtKind::AsyncFunctionDef { - name: name.try_map_ast(f)?, - args: args.try_map_ast(f)?, - body: body.try_map_ast(f)?, - decorator_list: decorator_list.try_map_ast(f)?, - returns: returns.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + name: Foldable::fold(name, folder)?, + args: Foldable::fold(args, folder)?, + body: Foldable::fold(body, folder)?, + decorator_list: Foldable::fold(decorator_list, folder)?, + returns: Foldable::fold(returns, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::ClassDef { name, @@ -451,31 +552,31 @@ impl MapAst for StmtKind { body, decorator_list, } => Ok(StmtKind::ClassDef { - name: name.try_map_ast(f)?, - bases: bases.try_map_ast(f)?, - keywords: keywords.try_map_ast(f)?, - body: body.try_map_ast(f)?, - decorator_list: decorator_list.try_map_ast(f)?, + name: Foldable::fold(name, folder)?, + bases: Foldable::fold(bases, folder)?, + keywords: Foldable::fold(keywords, folder)?, + body: Foldable::fold(body, folder)?, + decorator_list: Foldable::fold(decorator_list, folder)?, }), StmtKind::Return { value } => Ok(StmtKind::Return { - value: value.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, }), StmtKind::Delete { targets } => Ok(StmtKind::Delete { - targets: targets.try_map_ast(f)?, + targets: Foldable::fold(targets, folder)?, }), StmtKind::Assign { targets, value, type_comment, } => Ok(StmtKind::Assign { - targets: targets.try_map_ast(f)?, - value: value.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + targets: Foldable::fold(targets, folder)?, + value: Foldable::fold(value, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::AugAssign { target, op, value } => Ok(StmtKind::AugAssign { - target: target.try_map_ast(f)?, - op: op.try_map_ast(f)?, - value: value.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + op: Foldable::fold(op, folder)?, + value: Foldable::fold(value, folder)?, }), StmtKind::AnnAssign { target, @@ -483,10 +584,10 @@ impl MapAst for StmtKind { value, simple, } => Ok(StmtKind::AnnAssign { - target: target.try_map_ast(f)?, - annotation: annotation.try_map_ast(f)?, - value: value.try_map_ast(f)?, - simple: simple.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + annotation: Foldable::fold(annotation, folder)?, + value: Foldable::fold(value, folder)?, + simple: Foldable::fold(simple, folder)?, }), StmtKind::For { target, @@ -495,11 +596,11 @@ impl MapAst for StmtKind { orelse, type_comment, } => Ok(StmtKind::For { - target: target.try_map_ast(f)?, - iter: iter.try_map_ast(f)?, - body: body.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + iter: Foldable::fold(iter, folder)?, + body: Foldable::fold(body, folder)?, + orelse: Foldable::fold(orelse, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::AsyncFor { target, @@ -508,43 +609,43 @@ impl MapAst for StmtKind { orelse, type_comment, } => Ok(StmtKind::AsyncFor { - target: target.try_map_ast(f)?, - iter: iter.try_map_ast(f)?, - body: body.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + iter: Foldable::fold(iter, folder)?, + body: Foldable::fold(body, folder)?, + orelse: Foldable::fold(orelse, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::While { test, body, orelse } => Ok(StmtKind::While { - test: test.try_map_ast(f)?, - body: body.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, + test: Foldable::fold(test, folder)?, + body: Foldable::fold(body, folder)?, + orelse: Foldable::fold(orelse, folder)?, }), StmtKind::If { test, body, orelse } => Ok(StmtKind::If { - test: test.try_map_ast(f)?, - body: body.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, + test: Foldable::fold(test, folder)?, + body: Foldable::fold(body, folder)?, + orelse: Foldable::fold(orelse, folder)?, }), StmtKind::With { items, body, type_comment, } => Ok(StmtKind::With { - items: items.try_map_ast(f)?, - body: body.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + items: Foldable::fold(items, folder)?, + body: Foldable::fold(body, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::AsyncWith { items, body, type_comment, } => Ok(StmtKind::AsyncWith { - items: items.try_map_ast(f)?, - body: body.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + items: Foldable::fold(items, folder)?, + body: Foldable::fold(body, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, }), StmtKind::Raise { exc, cause } => Ok(StmtKind::Raise { - exc: exc.try_map_ast(f)?, - cause: cause.try_map_ast(f)?, + exc: Foldable::fold(exc, folder)?, + cause: Foldable::fold(cause, folder)?, }), StmtKind::Try { body, @@ -552,216 +653,359 @@ impl MapAst for StmtKind { orelse, finalbody, } => Ok(StmtKind::Try { - body: body.try_map_ast(f)?, - handlers: handlers.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, - finalbody: finalbody.try_map_ast(f)?, + body: Foldable::fold(body, folder)?, + handlers: Foldable::fold(handlers, folder)?, + orelse: Foldable::fold(orelse, folder)?, + finalbody: Foldable::fold(finalbody, folder)?, }), StmtKind::Assert { test, msg } => Ok(StmtKind::Assert { - test: test.try_map_ast(f)?, - msg: msg.try_map_ast(f)?, + test: Foldable::fold(test, folder)?, + msg: Foldable::fold(msg, folder)?, }), StmtKind::Import { names } => Ok(StmtKind::Import { - names: names.try_map_ast(f)?, + names: Foldable::fold(names, folder)?, }), StmtKind::ImportFrom { module, names, level, } => Ok(StmtKind::ImportFrom { - module: module.try_map_ast(f)?, - names: names.try_map_ast(f)?, - level: level.try_map_ast(f)?, + module: Foldable::fold(module, folder)?, + names: Foldable::fold(names, folder)?, + level: Foldable::fold(level, folder)?, }), StmtKind::Global { names } => Ok(StmtKind::Global { - names: names.try_map_ast(f)?, + names: Foldable::fold(names, folder)?, }), StmtKind::Nonlocal { names } => Ok(StmtKind::Nonlocal { - names: names.try_map_ast(f)?, + names: Foldable::fold(names, folder)?, }), StmtKind::Expr { value } => Ok(StmtKind::Expr { - value: value.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, }), StmtKind::Pass {} => Ok(StmtKind::Pass {}), StmtKind::Break {} => Ok(StmtKind::Break {}), StmtKind::Continue {} => Ok(StmtKind::Continue {}), + }) + } + impl Foldable for Expr { + type Mapped = Expr; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_expr(folder, self) } } -} -impl MapAst for ExprKind { - type Mapped = ExprKind; - fn try_map_ast Result>(self, f: &mut F) -> Result { - match self { + pub fn fold_expr + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Expr, + ) -> Result, F::Error> { + fold_located(folder, node, |folder, node| match node { ExprKind::BoolOp { op, values } => Ok(ExprKind::BoolOp { - op: op.try_map_ast(f)?, - values: values.try_map_ast(f)?, + op: Foldable::fold(op, folder)?, + values: Foldable::fold(values, folder)?, }), ExprKind::NamedExpr { target, value } => Ok(ExprKind::NamedExpr { - target: target.try_map_ast(f)?, - value: value.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + value: Foldable::fold(value, folder)?, }), ExprKind::BinOp { left, op, right } => Ok(ExprKind::BinOp { - left: left.try_map_ast(f)?, - op: op.try_map_ast(f)?, - right: right.try_map_ast(f)?, + left: Foldable::fold(left, folder)?, + op: Foldable::fold(op, folder)?, + right: Foldable::fold(right, folder)?, }), ExprKind::UnaryOp { op, operand } => Ok(ExprKind::UnaryOp { - op: op.try_map_ast(f)?, - operand: operand.try_map_ast(f)?, + op: Foldable::fold(op, folder)?, + operand: Foldable::fold(operand, folder)?, }), ExprKind::Lambda { args, body } => Ok(ExprKind::Lambda { - args: args.try_map_ast(f)?, - body: body.try_map_ast(f)?, + args: Foldable::fold(args, folder)?, + body: Foldable::fold(body, folder)?, }), ExprKind::IfExp { test, body, orelse } => Ok(ExprKind::IfExp { - test: test.try_map_ast(f)?, - body: body.try_map_ast(f)?, - orelse: orelse.try_map_ast(f)?, + test: Foldable::fold(test, folder)?, + body: Foldable::fold(body, folder)?, + orelse: Foldable::fold(orelse, folder)?, }), ExprKind::Dict { keys, values } => Ok(ExprKind::Dict { - keys: keys.try_map_ast(f)?, - values: values.try_map_ast(f)?, + keys: Foldable::fold(keys, folder)?, + values: Foldable::fold(values, folder)?, }), ExprKind::Set { elts } => Ok(ExprKind::Set { - elts: elts.try_map_ast(f)?, + elts: Foldable::fold(elts, folder)?, }), ExprKind::ListComp { elt, generators } => Ok(ExprKind::ListComp { - elt: elt.try_map_ast(f)?, - generators: generators.try_map_ast(f)?, + elt: Foldable::fold(elt, folder)?, + generators: Foldable::fold(generators, folder)?, }), ExprKind::SetComp { elt, generators } => Ok(ExprKind::SetComp { - elt: elt.try_map_ast(f)?, - generators: generators.try_map_ast(f)?, + elt: Foldable::fold(elt, folder)?, + generators: Foldable::fold(generators, folder)?, }), ExprKind::DictComp { key, value, generators, } => Ok(ExprKind::DictComp { - key: key.try_map_ast(f)?, - value: value.try_map_ast(f)?, - generators: generators.try_map_ast(f)?, + key: Foldable::fold(key, folder)?, + value: Foldable::fold(value, folder)?, + generators: Foldable::fold(generators, folder)?, }), ExprKind::GeneratorExp { elt, generators } => Ok(ExprKind::GeneratorExp { - elt: elt.try_map_ast(f)?, - generators: generators.try_map_ast(f)?, + elt: Foldable::fold(elt, folder)?, + generators: Foldable::fold(generators, folder)?, }), ExprKind::Await { value } => Ok(ExprKind::Await { - value: value.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, }), ExprKind::Yield { value } => Ok(ExprKind::Yield { - value: value.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, }), ExprKind::YieldFrom { value } => Ok(ExprKind::YieldFrom { - value: value.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, }), ExprKind::Compare { left, ops, comparators, } => Ok(ExprKind::Compare { - left: left.try_map_ast(f)?, - ops: ops.try_map_ast(f)?, - comparators: comparators.try_map_ast(f)?, + left: Foldable::fold(left, folder)?, + ops: Foldable::fold(ops, folder)?, + comparators: Foldable::fold(comparators, folder)?, }), ExprKind::Call { func, args, keywords, } => Ok(ExprKind::Call { - func: func.try_map_ast(f)?, - args: args.try_map_ast(f)?, - keywords: keywords.try_map_ast(f)?, + func: Foldable::fold(func, folder)?, + args: Foldable::fold(args, folder)?, + keywords: Foldable::fold(keywords, folder)?, }), ExprKind::FormattedValue { value, conversion, format_spec, } => Ok(ExprKind::FormattedValue { - value: value.try_map_ast(f)?, - conversion: conversion.try_map_ast(f)?, - format_spec: format_spec.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, + conversion: Foldable::fold(conversion, folder)?, + format_spec: Foldable::fold(format_spec, folder)?, }), ExprKind::JoinedStr { values } => Ok(ExprKind::JoinedStr { - values: values.try_map_ast(f)?, + values: Foldable::fold(values, folder)?, }), ExprKind::Constant { value, kind } => Ok(ExprKind::Constant { - value: value.try_map_ast(f)?, - kind: kind.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, + kind: Foldable::fold(kind, folder)?, }), ExprKind::Attribute { value, attr, ctx } => Ok(ExprKind::Attribute { - value: value.try_map_ast(f)?, - attr: attr.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, + attr: Foldable::fold(attr, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::Subscript { value, slice, ctx } => Ok(ExprKind::Subscript { - value: value.try_map_ast(f)?, - slice: slice.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, + slice: Foldable::fold(slice, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::Starred { value, ctx } => Ok(ExprKind::Starred { - value: value.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + value: Foldable::fold(value, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::Name { id, ctx } => Ok(ExprKind::Name { - id: id.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + id: Foldable::fold(id, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::List { elts, ctx } => Ok(ExprKind::List { - elts: elts.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + elts: Foldable::fold(elts, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::Tuple { elts, ctx } => Ok(ExprKind::Tuple { - elts: elts.try_map_ast(f)?, - ctx: ctx.try_map_ast(f)?, + elts: Foldable::fold(elts, folder)?, + ctx: Foldable::fold(ctx, folder)?, }), ExprKind::Slice { lower, upper, step } => Ok(ExprKind::Slice { - lower: lower.try_map_ast(f)?, - upper: upper.try_map_ast(f)?, - step: step.try_map_ast(f)?, + lower: Foldable::fold(lower, folder)?, + upper: Foldable::fold(upper, folder)?, + step: Foldable::fold(step, folder)?, }), + }) + } + impl Foldable for ExprContext { + type Mapped = ExprContext; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_expr_context(folder, self) } } -} -no_user!(ExprContext); -no_user!(Boolop); -no_user!(Operator); -no_user!(Unaryop); -no_user!(Cmpop); -impl MapAst for Comprehension { - type Mapped = Comprehension; - fn try_map_ast Result>(self, f: &mut F) -> Result { + pub fn fold_expr_context + ?Sized>( + #[allow(unused)] folder: &mut F, + node: ExprContext, + ) -> Result { + match node { + ExprContext::Load {} => Ok(ExprContext::Load {}), + ExprContext::Store {} => Ok(ExprContext::Store {}), + ExprContext::Del {} => Ok(ExprContext::Del {}), + } + } + impl Foldable for Boolop { + type Mapped = Boolop; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_boolop(folder, self) + } + } + pub fn fold_boolop + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Boolop, + ) -> Result { + match node { + Boolop::And {} => Ok(Boolop::And {}), + Boolop::Or {} => Ok(Boolop::Or {}), + } + } + impl Foldable for Operator { + type Mapped = Operator; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_operator(folder, self) + } + } + pub fn fold_operator + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Operator, + ) -> Result { + match node { + Operator::Add {} => Ok(Operator::Add {}), + Operator::Sub {} => Ok(Operator::Sub {}), + Operator::Mult {} => Ok(Operator::Mult {}), + Operator::MatMult {} => Ok(Operator::MatMult {}), + Operator::Div {} => Ok(Operator::Div {}), + Operator::Mod {} => Ok(Operator::Mod {}), + Operator::Pow {} => Ok(Operator::Pow {}), + Operator::LShift {} => Ok(Operator::LShift {}), + Operator::RShift {} => Ok(Operator::RShift {}), + Operator::BitOr {} => Ok(Operator::BitOr {}), + Operator::BitXor {} => Ok(Operator::BitXor {}), + Operator::BitAnd {} => Ok(Operator::BitAnd {}), + Operator::FloorDiv {} => Ok(Operator::FloorDiv {}), + } + } + impl Foldable for Unaryop { + type Mapped = Unaryop; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_unaryop(folder, self) + } + } + pub fn fold_unaryop + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Unaryop, + ) -> Result { + match node { + Unaryop::Invert {} => Ok(Unaryop::Invert {}), + Unaryop::Not {} => Ok(Unaryop::Not {}), + Unaryop::UAdd {} => Ok(Unaryop::UAdd {}), + Unaryop::USub {} => Ok(Unaryop::USub {}), + } + } + impl Foldable for Cmpop { + type Mapped = Cmpop; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_cmpop(folder, self) + } + } + pub fn fold_cmpop + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Cmpop, + ) -> Result { + match node { + Cmpop::Eq {} => Ok(Cmpop::Eq {}), + Cmpop::NotEq {} => Ok(Cmpop::NotEq {}), + Cmpop::Lt {} => Ok(Cmpop::Lt {}), + Cmpop::LtE {} => Ok(Cmpop::LtE {}), + Cmpop::Gt {} => Ok(Cmpop::Gt {}), + Cmpop::GtE {} => Ok(Cmpop::GtE {}), + Cmpop::Is {} => Ok(Cmpop::Is {}), + Cmpop::IsNot {} => Ok(Cmpop::IsNot {}), + Cmpop::In {} => Ok(Cmpop::In {}), + Cmpop::NotIn {} => Ok(Cmpop::NotIn {}), + } + } + impl Foldable for Comprehension { + type Mapped = Comprehension; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_comprehension(folder, self) + } + } + pub fn fold_comprehension + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Comprehension, + ) -> Result, F::Error> { let Comprehension { target, iter, ifs, is_async, - } = self; + } = node; Ok(Comprehension { - target: target.try_map_ast(f)?, - iter: iter.try_map_ast(f)?, - ifs: ifs.try_map_ast(f)?, - is_async: is_async.try_map_ast(f)?, + target: Foldable::fold(target, folder)?, + iter: Foldable::fold(iter, folder)?, + ifs: Foldable::fold(ifs, folder)?, + is_async: Foldable::fold(is_async, folder)?, }) } -} -impl MapAst for ExcepthandlerKind { - type Mapped = ExcepthandlerKind; - fn try_map_ast Result>(self, f: &mut F) -> Result { - match self { - ExcepthandlerKind::ExceptHandler { type_, name, body } => { - Ok(ExcepthandlerKind::ExceptHandler { - type_: type_.try_map_ast(f)?, - name: name.try_map_ast(f)?, - body: body.try_map_ast(f)?, - }) - } + impl Foldable for Excepthandler { + type Mapped = Excepthandler; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_excepthandler(folder, self) } } -} -impl MapAst for Arguments { - type Mapped = Arguments; - fn try_map_ast Result>(self, f: &mut F) -> Result { + pub fn fold_excepthandler + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Excepthandler, + ) -> Result, F::Error> { + fold_located(folder, node, |folder, node| match node { + ExcepthandlerKind::ExceptHandler { type_, name, body } => { + Ok(ExcepthandlerKind::ExceptHandler { + type_: Foldable::fold(type_, folder)?, + name: Foldable::fold(name, folder)?, + body: Foldable::fold(body, folder)?, + }) + } + }) + } + impl Foldable for Arguments { + type Mapped = Arguments; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_arguments(folder, self) + } + } + pub fn fold_arguments + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Arguments, + ) -> Result, F::Error> { let Arguments { posonlyargs, args, @@ -770,55 +1014,123 @@ impl MapAst for Arguments { kw_defaults, kwarg, defaults, - } = self; + } = node; Ok(Arguments { - posonlyargs: posonlyargs.try_map_ast(f)?, - args: args.try_map_ast(f)?, - vararg: vararg.try_map_ast(f)?, - kwonlyargs: kwonlyargs.try_map_ast(f)?, - kw_defaults: kw_defaults.try_map_ast(f)?, - kwarg: kwarg.try_map_ast(f)?, - defaults: defaults.try_map_ast(f)?, + posonlyargs: Foldable::fold(posonlyargs, folder)?, + args: Foldable::fold(args, folder)?, + vararg: Foldable::fold(vararg, folder)?, + kwonlyargs: Foldable::fold(kwonlyargs, folder)?, + kw_defaults: Foldable::fold(kw_defaults, folder)?, + kwarg: Foldable::fold(kwarg, folder)?, + defaults: Foldable::fold(defaults, folder)?, }) } -} -impl MapAst for ArgData { - type Mapped = ArgData; - fn try_map_ast Result>(self, f: &mut F) -> Result { - let ArgData { - arg, - annotation, - type_comment, - } = self; - Ok(ArgData { - arg: arg.try_map_ast(f)?, - annotation: annotation.try_map_ast(f)?, - type_comment: type_comment.try_map_ast(f)?, + impl Foldable for Arg { + type Mapped = Arg; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_arg(folder, self) + } + } + pub fn fold_arg + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Arg, + ) -> Result, F::Error> { + fold_located(folder, node, |folder, node| { + let ArgData { + arg, + annotation, + type_comment, + } = node; + Ok(ArgData { + arg: Foldable::fold(arg, folder)?, + annotation: Foldable::fold(annotation, folder)?, + type_comment: Foldable::fold(type_comment, folder)?, + }) }) } -} -impl MapAst for KeywordData { - type Mapped = KeywordData; - fn try_map_ast Result>(self, f: &mut F) -> Result { - let KeywordData { arg, value } = self; - Ok(KeywordData { - arg: arg.try_map_ast(f)?, - value: value.try_map_ast(f)?, + impl Foldable for Keyword { + type Mapped = Keyword; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_keyword(folder, self) + } + } + pub fn fold_keyword + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Keyword, + ) -> Result, F::Error> { + fold_located(folder, node, |folder, node| { + let KeywordData { arg, value } = node; + Ok(KeywordData { + arg: Foldable::fold(arg, folder)?, + value: Foldable::fold(value, folder)?, + }) }) } -} -no_user!(Alias); -impl MapAst for Withitem { - type Mapped = Withitem; - fn try_map_ast Result>(self, f: &mut F) -> Result { + impl Foldable for Alias { + type Mapped = Alias; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_alias(folder, self) + } + } + pub fn fold_alias + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Alias, + ) -> Result { + let Alias { name, asname } = node; + Ok(Alias { + name: Foldable::fold(name, folder)?, + asname: Foldable::fold(asname, folder)?, + }) + } + impl Foldable for Withitem { + type Mapped = Withitem; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_withitem(folder, self) + } + } + pub fn fold_withitem + ?Sized>( + #[allow(unused)] folder: &mut F, + node: Withitem, + ) -> Result, F::Error> { let Withitem { context_expr, optional_vars, - } = self; + } = node; Ok(Withitem { - context_expr: context_expr.try_map_ast(f)?, - optional_vars: optional_vars.try_map_ast(f)?, + context_expr: Foldable::fold(context_expr, folder)?, + optional_vars: Foldable::fold(optional_vars, folder)?, }) } + impl Foldable for TypeIgnore { + type Mapped = TypeIgnore; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + fold_type_ignore(folder, self) + } + } + pub fn fold_type_ignore + ?Sized>( + #[allow(unused)] folder: &mut F, + node: TypeIgnore, + ) -> Result { + match node { + TypeIgnore::TypeIgnore { lineno, tag } => Ok(TypeIgnore::TypeIgnore { + lineno: Foldable::fold(lineno, folder)?, + tag: Foldable::fold(tag, folder)?, + }), + } + } } -no_user!(TypeIgnore); diff --git a/ast/src/fold_helpers.rs b/ast/src/fold_helpers.rs new file mode 100644 index 000000000..26a08b2fb --- /dev/null +++ b/ast/src/fold_helpers.rs @@ -0,0 +1,72 @@ +use crate::constant; +use crate::fold::Fold; + +pub(crate) trait Foldable { + type Mapped; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result; +} + +impl Foldable for Vec +where + X: Foldable, +{ + type Mapped = Vec; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + self.into_iter().map(|x| x.fold(folder)).collect() + } +} + +impl Foldable for Option +where + X: Foldable, +{ + type Mapped = Option; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + self.map(|x| x.fold(folder)).transpose() + } +} + +impl Foldable for Box +where + X: Foldable, +{ + type Mapped = Box; + fn fold + ?Sized>( + self, + folder: &mut F, + ) -> Result { + (*self).fold(folder).map(Box::new) + } +} + +macro_rules! simple_fold { + ($($t:ty),+$(,)?) => { + $(impl $crate::fold_helpers::Foldable for $t { + type Mapped = Self; + #[inline] + fn fold + ?Sized>( + self, + _folder: &mut F, + ) -> Result { + Ok(self) + } + })+ + }; +} + +simple_fold!( + usize, + String, + bool, + constant::Constant, + constant::ConversionFlag +); diff --git a/ast/src/lib.rs b/ast/src/lib.rs index 371cb6ac0..83e862315 100644 --- a/ast/src/lib.rs +++ b/ast/src/lib.rs @@ -1,9 +1,7 @@ -// above for the macro -#[macro_use] -pub mod map_ast; - mod ast_gen; mod constant; +#[cfg(feature = "fold")] +mod fold_helpers; mod impls; mod location; diff --git a/ast/src/map_ast.rs b/ast/src/map_ast.rs deleted file mode 100644 index 8e7c885fa..000000000 --- a/ast/src/map_ast.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::convert::Infallible; - -pub trait MapAst: Sized { - type Mapped; - fn try_map_ast Result>(self, f: &mut F) -> Result; - fn map_ast U>(self, mut f: F) -> Self::Mapped { - let result: Result<_, Infallible> = self.try_map_ast(&mut |u| Ok(f(u))); - match result { - Ok(mapped) => mapped, - Err(never) => match never {}, - } - } -} - -macro_rules! no_user { - ($t:ty) => { - impl MapAst for $t { - type Mapped = Self; - #[inline] - fn try_map_ast Result>( - self, - _f: &mut F, - ) -> Result { - Ok(self) - } - } - }; -} - -no_user!(String); -no_user!(crate::Constant); -no_user!(crate::ConversionFlag); -no_user!(bool); -no_user!(usize); - -impl> MapAst for crate::Located { - type Mapped = crate::Located; - fn try_map_ast Result>(self, f: &mut F) -> Result { - Ok(crate::Located { - location: self.location, - custom: f(self.custom)?, - node: self.node.try_map_ast(f)?, - }) - } -} - -impl> MapAst for Vec { - type Mapped = Vec; - fn try_map_ast Result>(self, f: &mut F) -> Result { - self.into_iter().map(|node| node.try_map_ast(f)).collect() - } -} - -impl> MapAst for Option { - type Mapped = Option; - fn try_map_ast Result>(self, f: &mut F) -> Result { - self.map(|node| node.try_map_ast(f)).transpose() - } -} - -impl> MapAst for Box { - type Mapped = Box; - fn try_map_ast Result>(self, f: &mut F) -> Result { - (*self).try_map_ast(f).map(Box::new) - } -} From 0a4141c778c6d628ac83eed528c8597a9e819731 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Sun, 4 Apr 2021 18:03:19 -0500 Subject: [PATCH 2/3] Add an ast optimizer --- ast/src/constant.rs | 128 ++++++++++++++++++++++++++++++++++ compiler/porcelain/src/lib.rs | 10 ++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/ast/src/constant.rs b/ast/src/constant.rs index 579fd47bc..dadff1616 100644 --- a/ast/src/constant.rs +++ b/ast/src/constant.rs @@ -28,6 +28,11 @@ impl From for Constant { Self::Bool(b) } } +impl From for Constant { + fn from(i: BigInt) -> Constant { + Self::Int(i) + } +} /// Transforms a value prior to formatting it. #[derive(Copy, Clone, Debug, PartialEq)] @@ -51,3 +56,126 @@ impl ConversionFlag { } } } + +#[cfg(feature = "constant-optimization")] +#[derive(Default)] +pub struct ConstantOptimizer { + _priv: (), +} + +#[cfg(feature = "constant-optimization")] +impl ConstantOptimizer { + #[inline] + pub fn new() -> Self { + Self { _priv: () } + } +} + +#[cfg(feature = "constant-optimization")] +impl crate::fold::Fold for ConstantOptimizer { + type TargetU = U; + type Error = std::convert::Infallible; + #[inline] + fn map_user(&mut self, user: U) -> Result { + Ok(user) + } + fn fold_expr(&mut self, node: crate::Expr) -> Result, Self::Error> { + match node.node { + crate::ExprKind::Tuple { elts, ctx } => { + let elts = elts + .into_iter() + .map(|x| self.fold_expr(x)) + .collect::, _>>()?; + let expr = if elts + .iter() + .all(|e| matches!(e.node, crate::ExprKind::Constant { .. })) + { + let tuple = elts + .into_iter() + .map(|e| match e.node { + crate::ExprKind::Constant { value, .. } => value, + _ => unreachable!(), + }) + .collect(); + crate::ExprKind::Constant { + value: Constant::Tuple(tuple), + kind: None, + } + } else { + crate::ExprKind::Tuple { elts, ctx } + }; + Ok(crate::Expr { + node: expr, + custom: node.custom, + location: node.location, + }) + } + _ => crate::fold::fold_expr(self, node), + } + } +} + +#[cfg(test)] +mod tests { + #[cfg(feature = "constant-optimization")] + #[test] + fn test_constant_opt() { + use super::*; + use crate::fold::Fold; + use crate::*; + + let location = Location::new(0, 0); + let custom = (); + let ast = Located { + location, + custom, + node: ExprKind::Tuple { + ctx: ExprContext::Load, + elts: vec![ + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(1).into(), + kind: None, + }, + }, + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(2).into(), + kind: None, + }, + }, + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(3).into(), + kind: None, + }, + }, + ], + }, + }; + let new_ast = ConstantOptimizer::new() + .fold_expr(ast) + .unwrap_or_else(|e| match e {}); + assert_eq!( + new_ast, + Located { + location, + custom, + node: ExprKind::Constant { + value: Constant::Tuple(vec![ + BigInt::from(1).into(), + BigInt::from(2).into(), + BigInt::from(3).into() + ]), + kind: None + }, + } + ); + } +} diff --git a/compiler/porcelain/src/lib.rs b/compiler/porcelain/src/lib.rs index c4bd72a68..8308cae0a 100644 --- a/compiler/porcelain/src/lib.rs +++ b/compiler/porcelain/src/lib.rs @@ -1,6 +1,7 @@ use rustpython_bytecode::CodeObject; use rustpython_compiler_core::{compile, symboltable}; -use rustpython_parser::{ast::Location, parser}; +use rustpython_parser::ast::{fold::Fold, ConstantOptimizer, Location}; +use rustpython_parser::parser; use std::fmt; pub use compile::{CompileOpts, Mode}; @@ -80,10 +81,15 @@ pub fn compile( compile::Mode::Eval => parser::Mode::Expression, compile::Mode::Single => parser::Mode::Interactive, }; - let ast = match parser::parse(source, mode) { + let mut ast = match parser::parse(source, mode) { Ok(x) => x, Err(e) => return Err(CompileError::from_parse(e, source, source_path)), }; + if opts.optimize > 0 { + ast = ConstantOptimizer::new() + .fold_mod(ast) + .unwrap_or_else(|e| match e {}); + } compile::compile_top(&ast, source_path, opts).map_err(|e| CompileError::from_compile(e, source)) } From d7d62bfde039ce382e2a5232a403e8f68161020a Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Mon, 5 Apr 2021 10:47:47 -0500 Subject: [PATCH 3/3] Fix nested folding --- ast/asdl_rs.py | 4 ++-- ast/src/ast_gen.rs | 32 ++++++++++++++++---------------- ast/src/constant.rs | 37 +++++++++++++++++++++++++++++++++---- 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/ast/asdl_rs.py b/ast/asdl_rs.py index 95a4ad8b2..3c5d6fd56 100755 --- a/ast/asdl_rs.py +++ b/ast/asdl_rs.py @@ -292,7 +292,7 @@ class FoldImplVisitor(TypeInfoEmitVisitor): self.emit(f"impl Foldable for {enumname}{apply_t} {{", depth) self.emit(f"type Mapped = {enumname}{apply_u};", depth + 1) self.emit("fn fold + ?Sized>(self, folder: &mut F) -> Result {", depth + 1) - self.emit(f"fold_{name}(folder, self)", depth + 2) + self.emit(f"folder.fold_{name}(self)", depth + 2) self.emit("}", depth + 1) self.emit("}", depth) @@ -320,7 +320,7 @@ class FoldImplVisitor(TypeInfoEmitVisitor): self.emit(f"impl Foldable for {structname}{apply_t} {{", depth) self.emit(f"type Mapped = {structname}{apply_u};", depth + 1) self.emit("fn fold + ?Sized>(self, folder: &mut F) -> Result {", depth + 1) - self.emit(f"fold_{name}(folder, self)", depth + 2) + self.emit(f"folder.fold_{name}(self)", depth + 2) self.emit("}", depth + 1) self.emit("}", depth) diff --git a/ast/src/ast_gen.rs b/ast/src/ast_gen.rs index 5d2aa6704..31380e087 100644 --- a/ast/src/ast_gen.rs +++ b/ast/src/ast_gen.rs @@ -477,7 +477,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_mod(folder, self) + folder.fold_mod(self) } } pub fn fold_mod + ?Sized>( @@ -507,7 +507,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_stmt(folder, self) + folder.fold_stmt(self) } } pub fn fold_stmt + ?Sized>( @@ -694,7 +694,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_expr(folder, self) + folder.fold_expr(self) } } pub fn fold_expr + ?Sized>( @@ -838,7 +838,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_expr_context(folder, self) + folder.fold_expr_context(self) } } pub fn fold_expr_context + ?Sized>( @@ -857,7 +857,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_boolop(folder, self) + folder.fold_boolop(self) } } pub fn fold_boolop + ?Sized>( @@ -875,7 +875,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_operator(folder, self) + folder.fold_operator(self) } } pub fn fold_operator + ?Sized>( @@ -904,7 +904,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_unaryop(folder, self) + folder.fold_unaryop(self) } } pub fn fold_unaryop + ?Sized>( @@ -924,7 +924,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_cmpop(folder, self) + folder.fold_cmpop(self) } } pub fn fold_cmpop + ?Sized>( @@ -950,7 +950,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_comprehension(folder, self) + folder.fold_comprehension(self) } } pub fn fold_comprehension + ?Sized>( @@ -976,7 +976,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_excepthandler(folder, self) + folder.fold_excepthandler(self) } } pub fn fold_excepthandler + ?Sized>( @@ -999,7 +999,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_arguments(folder, self) + folder.fold_arguments(self) } } pub fn fold_arguments + ?Sized>( @@ -1031,7 +1031,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_arg(folder, self) + folder.fold_arg(self) } } pub fn fold_arg + ?Sized>( @@ -1057,7 +1057,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_keyword(folder, self) + folder.fold_keyword(self) } } pub fn fold_keyword + ?Sized>( @@ -1078,7 +1078,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_alias(folder, self) + folder.fold_alias(self) } } pub fn fold_alias + ?Sized>( @@ -1097,7 +1097,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_withitem(folder, self) + folder.fold_withitem(self) } } pub fn fold_withitem + ?Sized>( @@ -1119,7 +1119,7 @@ pub mod fold { self, folder: &mut F, ) -> Result { - fold_type_ignore(folder, self) + folder.fold_type_ignore(self) } } pub fn fold_type_ignore + ?Sized>( diff --git a/ast/src/constant.rs b/ast/src/constant.rs index dadff1616..0a627783f 100644 --- a/ast/src/constant.rs +++ b/ast/src/constant.rs @@ -151,9 +151,34 @@ mod tests { Located { location, custom, - node: ExprKind::Constant { - value: BigInt::from(3).into(), - kind: None, + node: ExprKind::Tuple { + ctx: ExprContext::Load, + elts: vec![ + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(3).into(), + kind: None, + }, + }, + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(4).into(), + kind: None, + }, + }, + Located { + location, + custom, + node: ExprKind::Constant { + value: BigInt::from(5).into(), + kind: None, + }, + }, + ], }, }, ], @@ -171,7 +196,11 @@ mod tests { value: Constant::Tuple(vec![ BigInt::from(1).into(), BigInt::from(2).into(), - BigInt::from(3).into() + Constant::Tuple(vec![ + BigInt::from(3).into(), + BigInt::from(4).into(), + BigInt::from(5).into(), + ]) ]), kind: None },