SourceInfo: remove the AST stmt copy

had to change AST_Expression to store a stmt instead of a expr but this has the advantage that memory managment is easier

src/codegen/ast_interpreter.cpp

@@ -2041,7 +2041,7 @@ Box* astInterpretFunction(FunctionMetadata* md, Box* closure, Box* generator, Bo
     // (For instance, throwing the exception will try to fetch the current statement, but we determine
     // that by looking at the cfg.)
     if (!source_info->cfg)
-        source_info->cfg = computeCFG(source_info, source_info->body, md->param_names);
+        source_info->cfg = computeCFG(source_info, md->param_names);
 
     Box** vregs = NULL;
     int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();
@@ -2083,7 +2083,7 @@ Box* astInterpretFunctionEval(FunctionMetadata* md, Box* globals, Box* boxedLoca
     // that by looking at the cfg.)
     SourceInfo* source_info = md->source.get();
     if (!source_info->cfg)
-        source_info->cfg = computeCFG(source_info, source_info->body, md->param_names);
+        source_info->cfg = computeCFG(source_info, md->param_names);
 
     Box** vregs = NULL;
     int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();

src/codegen/codegen.cpp

@@ -93,15 +93,8 @@ void FunctionMetadata::addVersion(CompiledFunction* compiled) {
     }
 }
 
-SourceInfo::SourceInfo(BoxedModule* m, ScopingAnalysis* scoping, FutureFlags future_flags, AST* ast,
-                       std::vector<AST_stmt*> body, BoxedString* fn)
-    : parent_module(m),
-      scoping(scoping),
-      scope_info(NULL),
-      future_flags(future_flags),
-      ast(ast),
-      cfg(NULL),
-      body(std::move(body)) {
+SourceInfo::SourceInfo(BoxedModule* m, ScopingAnalysis* scoping, FutureFlags future_flags, AST* ast, BoxedString* fn)
+    : parent_module(m), scoping(scoping), scope_info(NULL), future_flags(future_flags), ast(ast), cfg(NULL) {
     assert(fn);
 
     // TODO: this is a very bad way of handling this:

src/codegen/cpython_ast.cpp

@@ -740,7 +740,15 @@ public:
             case Expression_kind: {
                 AST_Expression* rtn = new AST_Expression(llvm::make_unique<InternedStringPool>());
                 this->pool = rtn->interned_strings.get();
-                rtn->body = this->convert(mod->v.Expression.body);
+
+                // instead of storing the expression inside the AST node we convert it directly to a return statement
+                AST_expr* expr = this->convert(mod->v.Expression.body);
+                auto rtn_stmt = new AST_Return;
+                rtn_stmt->lineno = expr->lineno;
+                rtn_stmt->col_offset = expr->col_offset;
+                rtn_stmt->value = expr;
+
+                rtn->body = rtn_stmt;
                 return rtn;
             }
             default:

src/codegen/irgen/future.cpp

@@ -67,7 +67,7 @@ inline bool is_stmt_string(AST_stmt* stmt) {
     return stmt->type == AST_TYPE::Expr && static_cast<AST_Expr*>(stmt)->value->type == AST_TYPE::Str;
 }
 
-FutureFlags getFutureFlags(std::vector<AST_stmt*> const& body, const char* file) {
+FutureFlags getFutureFlags(llvm::ArrayRef<AST_stmt*> body, const char* file) {
     FutureFlags ff = 0;
 
     // Set the defaults for the future flags depending on what version we are

src/codegen/irgen/future.h

@@ -15,7 +15,7 @@
 #ifndef PYSTON_CODEGEN_IRGEN_FUTURE_H
 #define PYSTON_CODEGEN_IRGEN_FUTURE_H
 
-#include <vector>
+#include "llvm/ADT/ArrayRef.h"
 
 #include "core/types.h"
 
@@ -25,7 +25,7 @@ namespace pyston {
 // bad __future__ imports. Returns the futures that are turned on. This is used
 // for irgeneration; the parser still has to handle some futures on its own,
 // when they are relevant for the parser.
-FutureFlags getFutureFlags(std::vector<AST_stmt*> const& body, const char* file);
+FutureFlags getFutureFlags(llvm::ArrayRef<AST_stmt*> body, const char* file);
 }
 
 #endif

src/codegen/irgen/hooks.cpp

@@ -103,6 +103,21 @@ InternedString SourceInfo::mangleName(InternedString id) {
     return getScopeInfo()->mangleName(id);
 }
 
+llvm::ArrayRef<AST_stmt*> SourceInfo::getBody() const {
+    switch (ast->type) {
+        case AST_TYPE::ClassDef:
+            return ((AST_ClassDef*)ast)->body;
+        case AST_TYPE::Expression:
+            return ((AST_Expression*)ast)->body;
+        case AST_TYPE::FunctionDef:
+            return ((AST_FunctionDef*)ast)->body;
+        case AST_TYPE::Module:
+            return ((AST_Module*)ast)->body;
+        default:
+            RELEASE_ASSERT(0, "unknown %d", ast->type);
+    };
+}
+
 InternedStringPool& SourceInfo::getInternedStrings() {
     return scoping->getInternedStrings();
 }
@@ -135,8 +150,7 @@ BORROWED(BoxedString*) SourceInfo::getName() noexcept {
 }
 
 Box* SourceInfo::getDocString() {
-    AST_Str* first_str = NULL;
-
+    auto body = getBody();
     if (body.size() > 0 && body[0]->type == AST_TYPE::Expr
         && static_cast<AST_Expr*>(body[0])->value->type == AST_TYPE::Str) {
         return boxString(static_cast<AST_Str*>(static_cast<AST_Expr*>(body[0])->value)->str_data);
@@ -275,7 +289,7 @@ CompiledFunction* compileFunction(FunctionMetadata* f, FunctionSpecialization* s
 
     // Do the analysis now if we had deferred it earlier:
     if (source->cfg == NULL) {
-        source->cfg = computeCFG(source, source->body, f->param_names);
+        source->cfg = computeCFG(source, f->param_names);
     }
 
 
@@ -336,7 +350,7 @@ void compileAndRunModule(AST_Module* m, BoxedModule* bm) {
 
         auto fn_str = boxString(fn);
         AUTO_DECREF(fn_str);
-        std::unique_ptr<SourceInfo> si(new SourceInfo(bm, scoping, future_flags, m, m->body, fn_str));
+        std::unique_ptr<SourceInfo> si(new SourceInfo(bm, scoping, future_flags, m, fn_str));
 
         static BoxedString* doc_str = getStaticString("__doc__");
         bm->setattr(doc_str, autoDecref(si->getDocString()), NULL);
@@ -375,7 +389,7 @@ Box* evalOrExec(FunctionMetadata* md, Box* globals, Box* boxedLocals) {
     return astInterpretFunctionEval(md, globals, boxedLocals);
 }
 
-static FunctionMetadata* compileForEvalOrExec(AST* source, std::vector<AST_stmt*> body, BoxedString* fn,
+static FunctionMetadata* compileForEvalOrExec(AST* source, llvm::ArrayRef<AST_stmt*> body, BoxedString* fn,
                                               PyCompilerFlags* flags) {
     LOCK_REGION(codegen_rwlock.asWrite());
 
@@ -394,8 +408,7 @@ static FunctionMetadata* compileForEvalOrExec(AST* source, std::vector<AST_stmt*
         flags->cf_flags = future_flags;
     }
 
-    std::unique_ptr<SourceInfo> si(
-        new SourceInfo(getCurrentModule(), scoping, future_flags, source, std::move(body), fn));
+    std::unique_ptr<SourceInfo> si(new SourceInfo(getCurrentModule(), scoping, future_flags, source, fn));
 
     FunctionMetadata* md = new FunctionMetadata(0, false, false, std::move(si));
     return md;
@@ -406,14 +419,7 @@ static FunctionMetadata* compileExec(AST_Module* parsedModule, BoxedString* fn,
 }
 
 static FunctionMetadata* compileEval(AST_Expression* parsedExpr, BoxedString* fn, PyCompilerFlags* flags) {
-    // We need body (list of statements) to compile.
-    // Obtain this by simply making a single statement which contains the expression.
-    AST_Return* stmt = new AST_Return();
-    stmt->lineno = parsedExpr->body->lineno;
-    stmt->value = parsedExpr->body;
-    std::vector<AST_stmt*> body = { stmt };
-
-    return compileForEvalOrExec(parsedExpr, std::move(body), fn, flags);
+    return compileForEvalOrExec(parsedExpr, parsedExpr->body, fn, flags);
 }
 
 extern "C" PyCodeObject* PyAST_Compile(struct _mod* _mod, const char* filename, PyCompilerFlags* flags,

src/codegen/irgen/irgenerator.cpp

@@ -1754,7 +1754,7 @@ private:
             // I think it's better to just not generate bad speculations:
             if (rtn->canConvertTo(speculated_type)) {
                 auto source = irstate->getSourceInfo();
-                printf("On %s:%d, function %s:\n", source->getFn()->c_str(), source->body[0]->lineno,
+                printf("On %s:%d, function %s:\n", source->getFn()->c_str(), source->getBody()[0]->lineno,
                        source->getName()->c_str());
                 irstate->getSourceInfo()->cfg->print();
             }
@@ -3268,7 +3268,7 @@ FunctionMetadata* wrapFunction(AST* node, AST_arguments* args, const std::vector
     FunctionMetadata*& md = made[node];
     if (md == NULL) {
         std::unique_ptr<SourceInfo> si(
-            new SourceInfo(source->parent_module, source->scoping, source->future_flags, node, body, source->getFn()));
+            new SourceInfo(source->parent_module, source->scoping, source->future_flags, node, source->getFn()));
         if (args)
             md = new FunctionMetadata(args->args.size(), args->vararg, args->kwarg, std::move(si));
         else

src/core/ast.h

@@ -503,7 +503,8 @@ class AST_Expression : public AST {
 public:
     std::unique_ptr<InternedStringPool> interned_strings;
 
-    AST_expr* body;
+    // this should be an expr but we convert it into a AST_Return(AST_expr) to make the code simpler
+    AST_stmt* body;
 
     virtual void accept(ASTVisitor* v);
 

src/core/cfg.cpp

@@ -253,7 +253,7 @@ private:
 
     unsigned int next_var_index = 0;
 
-    friend CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names);
+    friend CFG* computeCFG(SourceInfo* source, const ParamNames& param_names);
 
 public:
     CFGVisitor(SourceInfo* source, AST_TYPE::AST_TYPE root_type, FutureFlags future_flags,
@@ -2862,9 +2862,11 @@ void VRegInfo::assignVRegs(CFG* cfg, const ParamNames& param_names, ScopeInfo* s
 #endif
 }
 
-CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names) {
+CFG* computeCFG(SourceInfo* source, const ParamNames& param_names) {
     STAT_TIMER(t0, "us_timer_computecfg", 0);
 
+    auto body = source->getBody();
+
     CFG* rtn = new CFG();
 
     ScopingAnalysis* scoping_analysis = source->scoping;

src/core/cfg.h

@@ -320,7 +320,7 @@ public:
 };
 
 class SourceInfo;
-CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names);
+CFG* computeCFG(SourceInfo* source, const ParamNames& param_names);
 void printCFG(CFG* cfg);
 }
 

src/core/types.h

@@ -421,20 +421,17 @@ public:
     ScopeInfo* getScopeInfo();
     LivenessAnalysis* getLiveness();
 
-    // TODO we're currently copying the body of the AST into here, since lambdas don't really have a statement-based
-    // body and we have to create one.  Ideally, we'd be able to avoid the space duplication for non-lambdas.
-    const std::vector<AST_stmt*> body;
-
     // does not throw CXX or CAPI exceptions:
     BORROWED(BoxedString*) getName() noexcept;
     BORROWED(BoxedString*) getFn();
 
     InternedString mangleName(InternedString id);
 
+    llvm::ArrayRef<AST_stmt*> getBody() const;
+
     Box* getDocString();
 
-    SourceInfo(BoxedModule* m, ScopingAnalysis* scoping, FutureFlags future_flags, AST* ast,
-               std::vector<AST_stmt*> body, BoxedString* fn);
+    SourceInfo(BoxedModule* m, ScopingAnalysis* scoping, FutureFlags future_flags, AST* ast, BoxedString* fn);
     ~SourceInfo();
 
 private:

src/runtime/util.cpp

@@ -251,7 +251,7 @@ extern "C" void dumpEx(void* p, int levels) {
             FunctionMetadata* md = f->md;
             if (md->source) {
                 printf("User-defined function '%s'\n", md->source->getName()->c_str());
-                printf("Defined at %s:%d\n", md->source->getFn()->c_str(), md->source->body[0]->lineno);
+                printf("Defined at %s:%d\n", md->source->getFn()->c_str(), md->source->getBody()[0]->lineno);
 
                 if (md->source->cfg && levels > 0) {
                     md->source->cfg->print();

test/unittests/analysis.cpp

@@ -40,11 +40,10 @@ TEST_F(AnalysisTest, augassign) {
 
     FutureFlags future_flags = getFutureFlags(module->body, fn.c_str());
 
-    SourceInfo* si = new SourceInfo(createModule(boxString("augassign"), fn.c_str()), scoping, future_flags, func,
-            func->body, boxString(fn));
+    SourceInfo* si = new SourceInfo(createModule(boxString("augassign"), fn.c_str()), scoping, future_flags, func, boxString(fn));
 
     ParamNames param_names(si->ast, si->getInternedStrings());
-    CFG* cfg = computeCFG(si, func->body, param_names);
+    CFG* cfg = computeCFG(si, param_names);
     std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg);
     auto&& vregs = cfg->getVRegInfo();
 
@@ -73,10 +72,10 @@ void doOsrTest(bool is_osr, bool i_maybe_undefined) {
 
     ScopeInfo* scope_info = scoping->getScopeInfoForNode(func);
     std::unique_ptr<SourceInfo> si(new SourceInfo(createModule(boxString("osr" + std::to_string((is_osr << 1) + i_maybe_undefined)),
-                    fn.c_str()), scoping, future_flags, func, func->body, boxString(fn)));
+                    fn.c_str()), scoping, future_flags, func, boxString(fn)));
     FunctionMetadata* clfunc = new FunctionMetadata(0, false, false, std::move(si));
 
-    CFG* cfg = computeCFG(clfunc->source.get(), func->body, clfunc->param_names);
+    CFG* cfg = computeCFG(clfunc->source.get(), clfunc->param_names);
     clfunc->source->cfg = cfg;
     std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg);