Skip to content

Commit

Permalink
Tree printer for Python AST and its ASR
Browse files Browse the repository at this point in the history
  • Loading branch information
@ansharlubis
ansharlubis committed Mar 18, 2022
1 parent f55ab69 commit 11f7cfa
Show file tree
Hide file tree
Showing 7 changed files with 317 additions and 5 deletions.
261 changes: 259 additions & 2 deletions grammar/asdl_cpp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,264 @@ def visitField(self, field):
self.emit( "this->visit_symbol(*a.second);", 3)
self.emit("}", 2)

# This class generates a visitor that prints the tree structure of AST/ASR
class TreeVisitorVisitor(ASDLVisitor):

def visitModule(self, mod):
self.emit("/" + "*"*78 + "/")
self.emit("// Tree Visitor base class")
self.emit("")
self.emit("template <class Derived>")
self.emit("class TreeBaseVisitor : public BaseVisitor<Derived>")
self.emit("{")
self.emit("private:")
self.emit( "Derived& self() { return static_cast<Derived&>(*this); }", 1)
self.emit("public:")
self.emit( "std::string s, indtd;", 1)
self.emit( "bool use_colors;", 1)
self.emit( "bool start_line = true;", 1)
self.emit( 'bool last, attached;', 1)
self.emit( "int indent_level = 0, indent_spaces = 2, lvl = 0;", 1)
self.emit("public:")
self.emit( "TreeBaseVisitor() : use_colors(false), last(true), attached(false) { s.reserve(100000); }", 1)
self.emit( "void inc_indent() {", 1)
self.emit( "indent_level++;", 2)
self.emit( 'indtd += " ";', 2)
self.emit( "}", 1)
self.emit( "void inc_lindent() {", 1)
self.emit( "indent_level++;", 2)
self.emit( 'indtd += "| ";', 2)
self.emit( "}", 1)
self.emit( "void dec_indent() {", 1)
self.emit( "indent_level--;", 2)
self.emit( "LFORTRAN_ASSERT(indent_level >= 0);", 2)
self.emit( "indtd = indtd.substr(0, indent_level*indent_spaces);",2)
self.emit( "}", 1)
self.mod = mod
super(TreeVisitorVisitor, self).visitModule(mod)
self.emit("};")

def visitType(self, tp):
super(TreeVisitorVisitor, self).visitType(tp, tp.name)

def visitSum(self, sum, *args):
assert isinstance(sum, asdl.Sum)
if is_simple_sum(sum):
name = args[0] + "Type"
self.make_simple_sum_visitor(name, sum.types)
else:
for tp in sum.types:
self.visit(tp, *args)

def visitProduct(self, prod, name):
self.make_visitor(name, prod.fields, False)

def visitConstructor(self, cons, _):
self.make_visitor(cons.name, cons.fields, True)

def make_visitor(self, name, fields, cons):
self.emit("void visit_%s(const %s_t &x) {" % (name, name), 1)
self.emit( 'if(!attached) {', 2)
self.emit( 'if(start_line) {', 3)
self.emit( 'start_line = false;', 4)
self.emit( 's.append(indtd);', 4)
self.emit( '} else {', 3)
self.emit( 's.append("\\n"+indtd);', 4)
self.emit( '}', 3)
self.emit( 'last ? s.append("└-") : s.append("|-");', 3)
self.emit( '}', 2)
self.emit( 'last ? inc_indent() : inc_lindent();', 2)
self.emit( 'attached, last = false;', 2)
if cons:
self.emit( 'if (use_colors) {', 2)
self.emit( 's.append(color(style::bold));', 3)
self.emit( 's.append(color(fg::magenta));', 3)
self.emit( '}', 2)
self.emit( 's.append("%s");' % name, 2)
self.emit( 'if (use_colors) {', 2)
self.emit( 's.append(color(fg::reset));', 3)
self.emit( 's.append(color(style::reset));', 3)
self.emit( '}', 2)
self.used = False
for n, field in enumerate(fields):
self.visitField(field, cons, n == len(fields)-1)
self.emit( 'dec_indent();', 2)
if not self.used:
# Note: a better solution would be to change `&x` to `& /* x */`
# above, but we would need to change emit to return a string.
self.emit("if ((bool&)x) { } // Suppress unused warning", 2)
self.emit("}", 1)

def make_simple_sum_visitor(self, name, types):
self.emit("void visit_%s(const %s &x) {" % (name, name), 1)
self.emit( 'if (use_colors) {', 2)
self.emit( 's.append(color(style::bold));', 3)
self.emit( 's.append(color(fg::green));', 3)
self.emit( '}', 2)
self.emit( 'switch (x) {', 2)
for tp in types:
self.emit( 'case (%s::%s) : {' % (name, tp.name), 3)
self.emit( 's.append("%s");' % (tp.name), 4)
self.emit( ' break; }',3)
self.emit( '}', 2)
self.emit( 'if (use_colors) {', 2)
self.emit( 's.append(color(fg::reset));', 3)
self.emit( 's.append(color(style::reset));', 3)
self.emit( '}', 2)
self.emit("}", 1)

def visitField(self, field, cons, last):
arr = '└-' if last else '|-'
if (field.type not in asdl.builtin_types and
field.type not in self.data.simple_types):
self.used = True
level = 2
if field.type in products:
if field.opt:
template = "self().visit_%s(*x.m_%s);" % (field.type, field.name)
else:
template = "self().visit_%s(x.m_%s);" % (field.type, field.name)
else:
template = "self().visit_%s(*x.m_%s);" % (field.type, field.name)
if field.seq:
self.emit('s.append("\\n" + indtd + "%s" + "%s=\u21a7");' % (arr, field.name), level)
self.emit("for (size_t i=0; i<x.n_%s; i++) {" % field.name, level)
self.emit( 'inc_indent();' if last else 'inc_lindent();', level+1)
self.emit( 'last = i == x.n_%s-1;' % field.name, level+1)
self.emit( 'attached = false;', level+1)
if field.type in sums:
self.emit("self().visit_%s(*x.m_%s[i]);" % (field.type, field.name), level+1)
else:
self.emit("self().visit_%s(x.m_%s[i]);" % (field.type, field.name), level+1)
self.emit( 'dec_indent();', level+1)
self.emit("}", level)
elif field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
if last:
self.emit('last, attached = true;', 2)
else:
self.emit('attached = true;', 2)
self.emit("if (x.m_%s) {" % field.name, 2)
self.emit(template, 3)
self.emit("} else {", 2)
self.emit( 's.append("()");', 3)
self.emit( 'last, attached = false;', 3)
self.emit("}", 2)
else:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), level)
if last:
self.emit('last = true;', level)
self.emit('attached = true;', level)
self.emit(template, level)
else:
if field.type == "identifier":
if field.seq:
assert not field.opt
level = 2
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), level)
self.emit("for (size_t i=0; i<x.n_%s; i++) {" % field.name, level)
self.emit( "s.append(x.m_%s[i]);" % (field.name), level+1)
self.emit( 'if (i < x.n_%s-1) s.append(" ");' % (field.name), level+1)
self.emit("}", level)
else:
if field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit("if (x.m_%s) {" % field.name, 2)
self.emit( 's.append(x.m_%s);' % field.name, 3)
self.emit("} else {", 2)
self.emit( 's.append("()");', 3)
self.emit("}", 2)
else:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit('s.append(x.m_%s);' % field.name, 2)
elif field.type == "node":
assert not field.opt
assert field.seq
level = 2
self.emit('s.append("\\n" + indtd + "%s" + "%s=\u21a7");' % (arr, field.name), level)
self.emit('attached = false;', level)
self.emit("for (size_t i=0; i<x.n_%s; i++) {" % field.name, level)
mod_name = self.mod.name.lower()
self.emit( 'inc_indent();' if last else 'inc_lindent();', level+1)
self.emit( 'last = i == x.n_%s-1;' % field.name, level+1)
self.emit( 'attached = false;', level+1)
self.emit( "self().visit_%s(*x.m_%s[i]);" % (mod_name, field.name), level+1)
self.emit( 'dec_indent();', level+1)
self.emit("}", level)
elif field.type == "symbol_table":
assert not field.opt
assert not field.seq
if field.name == "parent_symtab":
level = 2
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), level)
self.emit('s.append(x.m_%s->get_counter());' % field.name, level)
else:
level = 2
self.emit('s.append("\\n" + indtd + "%s");' % arr, level)
self.emit('inc_lindent();', level)
self.emit('if (use_colors) {', level)
self.emit( 's.append(color(fg::yellow));', level+1)
self.emit('}', level)
self.emit('s.append("SymbolTable");', level)
self.emit('if (use_colors) {', level)
self.emit( 's.append(color(fg::reset));', level+1)
self.emit('}', level)
self.emit( 's.append("\\n" + indtd + "|-counter=");', level)
self.emit( 's.append(x.m_%s->get_counter());' % field.name, level)
self.emit('size_t i = 0;', level)
self.emit('s.append("\\n" + indtd + "└-scope=\u21a7");', level)
self.emit('for (auto &a : x.m_%s->scope) {' % field.name, level)
self.emit( 'i++;', level+1)
self.emit( 'inc_indent();', level+1)
self.emit( 'last = i == x.m_%s->scope.size();' % field.name, level+1)
self.emit( 's.append("\\n" + indtd + (last ? "└-" : "|-") + a.first + ": ");', level+1)
self.emit( 'attached = true;', level+1)
self.emit( 'this->visit_symbol(*a.second);', level+1)
self.emit( 'dec_indent();', level+1)
self.emit('}', level)
self.emit('dec_indent();', level)
elif field.type == "string" and not field.seq:
if field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit("if (x.m_%s) {" % field.name, 2)
self.emit( 's.append("\\"" + std::string(x.m_%s) + "\\"");' % field.name, 3)
self.emit("} else {", 2)
self.emit( 's.append("()");', 3)
self.emit("}", 2)
else:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit('s.append("\\"" + std::string(x.m_%s) + "\\"");' % field.name, 2)
elif field.type == "int" and not field.seq:
if field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit("if (x.m_%s) {" % field.name, 2)
self.emit( 's.append(std::to_string(x.m_%s));' % field.name, 3)
self.emit("} else {", 2)
self.emit( 's.append("()");', 3)
self.emit("}", 2)
else:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit('s.append(std::to_string(x.m_%s));' % field.name, 2)
elif field.type == "float" and not field.seq and not field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit('s.append(std::to_string(x.m_%s));' % field.name, 2)
elif field.type == "bool" and not field.seq and not field.opt:
self.emit('s.append("\\n" + indtd + "%s" + "%s=");' % (arr, field.name), 2)
self.emit("if (x.m_%s) {" % field.name, 2)
self.emit( 's.append(".true.");', 3)
self.emit("} else {", 2)
self.emit( 's.append(".false.");', 3)
self.emit("}", 2)
elif field.type in self.data.simple_types:
if field.opt:
self.emit('s.append("Unimplementedopt");', 2)
else:
self.emit('s.append("\\n" + indtd + "%s" + "%sType=");' % (arr, field.type), 2)
self.emit('visit_%sType(x.m_%s);' \
% (field.type, field.name), 2)
else:
self.emit('s.append("Unimplemented' + field.type + '");', 2)


class PickleVisitorVisitor(ASDLVisitor):

Expand Down Expand Up @@ -1257,7 +1515,6 @@ def add_masks(fields, node):
return down_cast<T>(t);
}
"""

FOOT = r"""} // namespace LFortran::%(MOD)s
Expand All @@ -1267,7 +1524,7 @@ def add_masks(fields, node):

visitors = [ASTNodeVisitor0, ASTNodeVisitor1, ASTNodeVisitor,
ASTVisitorVisitor1, ASTVisitorVisitor1b, ASTVisitorVisitor2,
ASTWalkVisitorVisitor, PickleVisitorVisitor,
ASTWalkVisitorVisitor, TreeVisitorVisitor, PickleVisitorVisitor,
SerializationVisitorVisitor, DeserializationVisitorVisitor]


Expand Down
16 changes: 13 additions & 3 deletions src/bin/lpython.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -527,8 +527,12 @@ int emit_ast(const std::string &infile,
}
LFortran::Python::AST::ast_t* ast = r.result;

std::cout << LFortran::Python::pickle_python(*ast,
compiler_options.use_colors, compiler_options.indent) << std::endl;
if (compiler_options.tree) {
std::cout << LFortran::Python::pickle_tree_python(*ast, compiler_options.use_colors) << std::endl;
} else {
std::cout << LFortran::Python::pickle_python(*ast,
compiler_options.use_colors, compiler_options.indent) << std::endl;
}
return 0;
}

Expand Down Expand Up @@ -559,8 +563,13 @@ int emit_asr(const std::string &infile,
}
LFortran::ASR::TranslationUnit_t* asr = r.result;

std::cout << LFortran::pickle(*asr, compiler_options.use_colors, compiler_options.indent,
if (compiler_options.tree) {
std::cout << LFortran::pickle_tree(*asr, compiler_options.use_colors,
with_intrinsic_modules) << std::endl;
} else {
std::cout << LFortran::pickle(*asr, compiler_options.use_colors, compiler_options.indent,
with_intrinsic_modules) << std::endl;
}
return 0;
}

Expand Down Expand Up @@ -1285,6 +1294,7 @@ int main(int argc, char *argv[])
app.add_flag("--with-intrinsic-mods", with_intrinsic_modules, "Show intrinsic modules in ASR");
app.add_flag("--no-color", arg_no_color, "Turn off colored AST/ASR");
app.add_flag("--indent", compiler_options.indent, "Indented print ASR/AST");
app.add_flag("--tree", compiler_options.tree, "Tree structure print ASR/AST");
app.add_option("--pass", arg_pass, "Apply the ASR pass and show ASR (implies --show-asr)");
app.add_flag("--symtab-only", compiler_options.symtab_only, "Only create symbol tables in ASR (skip executable stmt)");
app.add_flag("--time-report", time_report, "Show compilation time report");
Expand Down
1 change: 1 addition & 0 deletions src/libasr/utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ struct CompilerOptions {
bool show_stacktrace = false;
bool use_colors = true;
bool indent = false;
bool tree = false;
bool fast = false;
bool openmp = false;
bool no_warnings = false;
Expand Down
24 changes: 24 additions & 0 deletions src/lpython/pickle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -270,4 +270,28 @@ std::string pickle(LFortran::ASR::TranslationUnit_t &asr, bool colors, bool inde
return pickle((ASR::asr_t &)asr, colors, indent, show_intrinsic_modules);
}

class ASRTreeVisitor :
public LFortran::ASR::TreeBaseVisitor<ASRTreeVisitor>
{
public:
bool show_intrinsic_modules;

std::string get_str() {
return s;
}

};

std::string pickle_tree(LFortran::ASR::asr_t &asr, bool colors, bool show_intrinsic_modules) {
ASRTreeVisitor v;
v.use_colors = colors;
v.show_intrinsic_modules = show_intrinsic_modules;
v.visit_asr(asr);
return v.get_str();
}

std::string pickle_tree(LFortran::ASR::TranslationUnit_t &asr, bool colors, bool show_intrinsic_modules) {
return pickle_tree((ASR::asr_t &)asr, colors, show_intrinsic_modules);
}

}
4 changes: 4 additions & 0 deletions src/lpython/pickle.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@ namespace LFortran {
std::string pickle(ASR::TranslationUnit_t &asr, bool colors=false,
bool indent=false, bool show_intrinsic_modules=false);

// Print the tree structure
std::string pickle_tree(LFortran::ASR::asr_t &asr, bool colors, bool show_intrinsic_modules);
std::string pickle_tree(LFortran::ASR::TranslationUnit_t &asr, bool colors, bool show_intrinsic_modules);

}

#endif // LFORTRAN_PICKLE_H
15 changes: 15 additions & 0 deletions src/lpython/semantics/python_ast_to_asr.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2433,6 +2433,21 @@ std::string pickle_python(AST::ast_t &ast, bool colors, bool indent) {
return v.get_str();
}

class TreeVisitor : public AST::TreeBaseVisitor<TreeVisitor>
{
public:
std::string get_str() {
return s;
}
};

std::string pickle_tree_python(AST::ast_t &ast, bool colors) {
TreeVisitor v;
v.use_colors = colors;
v.visit_ast(ast);
return v.get_str();
}

Result<ASR::TranslationUnit_t*> python_ast_to_asr(Allocator &al,
AST::ast_t &ast, diag::Diagnostics &diagnostics, bool main_module,
bool symtab_only)
Expand Down
1 change: 1 addition & 0 deletions src/lpython/semantics/python_ast_to_asr.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
namespace LFortran::Python {

std::string pickle_python(AST::ast_t &ast, bool colors=false, bool indent=false);
std::string pickle_tree_python(AST::ast_t &ast, bool colors=true);
Result<ASR::TranslationUnit_t*> python_ast_to_asr(Allocator &al,
Python::AST::ast_t &ast, diag::Diagnostics &diagnostics, bool main_module,
bool symtab_only);
Expand Down

0 comments on commit 11f7cfa

Please sign in to comment.