Create classdef scopes when evaluating classes

Not completely correct/working yet, but finally start
doing classes the right way.

src/analysis/scoping_analysis.cpp

@@ -41,6 +41,8 @@ public:
     }
     virtual bool refersToClosure(const std::string name) { return false; }
     virtual bool saveInClosure(const std::string name) { return false; }
+
+    virtual const std::unordered_set<std::string>& getClassDefLocalNames() { RELEASE_ASSERT(0, ""); }
 };
 
 struct ScopingAnalysis::ScopeNameUsage {
@@ -58,7 +60,12 @@ struct ScopingAnalysis::ScopeNameUsage {
     StrSet referenced_from_nested;
     StrSet got_from_closure;
 
-    ScopeNameUsage(AST* node, ScopeNameUsage* parent) : node(node), parent(parent) {}
+    ScopeNameUsage(AST* node, ScopeNameUsage* parent) : node(node), parent(parent) {
+        if (node->type == AST_TYPE::ClassDef) {
+            // classes have an implicit write to "__module__"
+            written.insert("__module__");
+        }
+    }
 };
 
 class ScopeInfoBase : public ScopeInfo {
@@ -101,6 +108,11 @@ public:
             return false;
         return usage->referenced_from_nested.count(name) != 0;
     }
+
+    virtual const std::unordered_set<std::string>& getClassDefLocalNames() {
+        RELEASE_ASSERT(usage->node->type == AST_TYPE::ClassDef, "");
+        return usage->written;
+    }
 };
 
 class NameCollectorVisitor : public ASTVisitor {

src/analysis/scoping_analysis.h

@@ -33,6 +33,11 @@ public:
     virtual bool refersToGlobal(const std::string& name) = 0;
     virtual bool refersToClosure(const std::string name) = 0;
     virtual bool saveInClosure(const std::string name) = 0;
+
+    // Get the names set within a classdef that should be forwarded on to
+    // the metaclass constructor.
+    // An error to call this on a non-classdef node.
+    virtual const std::unordered_set<std::string>& getClassDefLocalNames() = 0;
 };
 
 class ScopingAnalysis {

src/analysis/type_analysis.cpp

@@ -459,7 +459,9 @@ private:
     }
 
     virtual void visit_classdef(AST_ClassDef* node) {
-        CompilerType* t = typeFromClass(type_cls);
+        // TODO should we speculate that classdefs will generally return a class?
+        // CompilerType* t = typeFromClass(type_cls);
+        CompilerType* t = UNKNOWN;
         _doSet(node->name, t);
     }
 

src/codegen/irgen/hooks.cpp

@@ -43,7 +43,7 @@ namespace pyston {
 
 // TODO terrible place for these!
 SourceInfo::ArgNames::ArgNames(AST* ast) {
-    if (ast->type == AST_TYPE::Module) {
+    if (ast->type == AST_TYPE::Module || ast->type == AST_TYPE::ClassDef) {
         args = NULL;
         kwarg = vararg = NULL;
     } else if (ast->type == AST_TYPE::FunctionDef) {
@@ -59,6 +59,8 @@ SourceInfo::ArgNames::ArgNames(AST* ast) {
 const std::string SourceInfo::getName() {
     assert(ast);
     switch (ast->type) {
+        case AST_TYPE::ClassDef:
+            return ast_cast<AST_ClassDef>(ast)->name;
         case AST_TYPE::FunctionDef:
             return ast_cast<AST_FunctionDef>(ast)->name;
         case AST_TYPE::Module:
@@ -71,6 +73,8 @@ const std::string SourceInfo::getName() {
 const std::vector<AST_stmt*>& SourceInfo::getBody() {
     assert(ast);
     switch (ast->type) {
+        case AST_TYPE::ClassDef:
+            return ast_cast<AST_ClassDef>(ast)->body;
         case AST_TYPE::FunctionDef:
             return ast_cast<AST_FunctionDef>(ast)->body;
         case AST_TYPE::Module:
@@ -165,7 +169,7 @@ CompiledFunction* compileFunction(CLFunction* f, FunctionSpecialization* spec, E
     // Do the analysis now if we had deferred it earlier:
     if (source->cfg == NULL) {
         assert(source->ast);
-        source->cfg = computeCFG(source->ast->type, source->getBody());
+        source->cfg = computeCFG(source, source->getBody());
         source->liveness = computeLivenessInfo(source->cfg);
         source->phis = computeRequiredPhis(source->arg_names, source->cfg, source->liveness,
                                            source->scoping->getScopeInfoForNode(source->ast));
@@ -228,7 +232,7 @@ void compileAndRunModule(AST_Module* m, BoxedModule* bm) {
         ScopingAnalysis* scoping = runScopingAnalysis(m);
 
         SourceInfo* si = new SourceInfo(bm, scoping, m);
-        si->cfg = computeCFG(AST_TYPE::Module, m->body);
+        si->cfg = computeCFG(si, m->body);
         si->liveness = computeLivenessInfo(si->cfg);
         si->phis = computeRequiredPhis(si->arg_names, si->cfg, si->liveness, si->scoping->getScopeInfoForNode(si->ast));
 

src/core/ast.h

@@ -688,6 +688,8 @@ public:
     virtual void* accept_expr(ExprVisitor* v);
 
     AST_Str() : AST_expr(AST_TYPE::Str) {}
+    AST_Str(const std::string& s) : AST_expr(AST_TYPE::Str), s(s) {}
+    AST_Str(const std::string&& s) : AST_expr(AST_TYPE::Str), s(std::move(s)) {}
 
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Str;
 };

src/core/cfg.cpp

@@ -19,8 +19,11 @@
 #include <cstdio>
 #include <cstdlib>
 
+#include "analysis/scoping_analysis.h"
 #include "core/ast.h"
 #include "core/options.h"
+#include "core/types.h"
+#include "runtime/types.h"
 
 //#undef VERBOSITY
 //#define VERBOSITY(x) 2
@@ -49,6 +52,15 @@ void CFGBlock::unconnectFrom(CFGBlock* successor) {
                                   successor->predecessors.end());
 }
 
+static AST_Name* makeName(const std::string& id, AST_TYPE::AST_TYPE ctx_type, int lineno = 0, int col_offset = 0) {
+    AST_Name* name = new AST_Name();
+    name->id = id;
+    name->col_offset = col_offset;
+    name->lineno = lineno;
+    name->ctx_type = ctx_type;
+    return name;
+}
+
 class CFGVisitor : public ASTVisitor {
 private:
     AST_TYPE::AST_TYPE root_type;
@@ -327,15 +339,6 @@ private:
         return call;
     }
 
-    AST_Name* makeName(const std::string& id, AST_TYPE::AST_TYPE ctx_type, int lineno = 0, int col_offset = 0) {
-        AST_Name* name = new AST_Name();
-        name->id = id;
-        name->col_offset = col_offset;
-        name->lineno = lineno;
-        name->ctx_type = ctx_type;
-        return name;
-    }
-
     AST_stmt* makeAssign(AST_expr* target, AST_expr* val) {
         AST_Assign* assign = new AST_Assign();
         assign->targets.push_back(target);
@@ -1636,20 +1639,71 @@ void CFG::print() {
     delete pv;
 }
 
-CFG* computeCFG(AST_TYPE::AST_TYPE root_type, std::vector<AST_stmt*> body) {
+CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
     CFG* rtn = new CFG();
-    CFGVisitor visitor(root_type, rtn);
+    CFGVisitor visitor(source->ast->type, rtn);
+
+    // In a classdef, the "__module__" attribute is immediately available:
+    if (source->ast->type == AST_TYPE::ClassDef) {
+        Box* module_name = source->parent_module->getattr("__name__", NULL, NULL);
+        assert(module_name->cls == str_cls);
+        AST_Assign* module_assign = new AST_Assign();
+        module_assign->targets.push_back(makeName("__module__", AST_TYPE::Store));
+        module_assign->value = new AST_Str(static_cast<BoxedString*>(module_name)->s);
+
+        visitor.push_back(module_assign);
+    }
+
     for (int i = 0; i < body.size(); i++) {
         body[i]->accept(&visitor);
     }
 
-    // Put a fake "return" statement at the end of every function just to make sure they all have one;
-    // we already have to support multiple return statements in a function, but this way we can avoid
-    // having to support not having a return statement:
-    AST_Return* return_stmt = new AST_Return();
-    return_stmt->lineno = return_stmt->col_offset = 0;
-    return_stmt->value = NULL;
-    visitor.push_back(return_stmt);
+    // The functions we create for classdefs are supposed to return a dictionary of their locals.
+    // This is the place that we add all of that:
+    if (source->ast->type == AST_TYPE::ClassDef) {
+        ScopeInfo* scope_info = source->scoping->getScopeInfoForNode(source->ast);
+
+        auto written_names = scope_info->getClassDefLocalNames();
+        AST_Dict* rtn_dict = new AST_Dict();
+
+        // It'd be ok to add __doc__ to the dict multiple times, since the last one would win
+        if (written_names.count("__doc__") == 0) {
+            bool found_docstr = false;
+            if (body.size() && body[0]->type == AST_TYPE::Expr) {
+                AST_Expr* first_expr = ast_cast<AST_Expr>(body[0]);
+                if (first_expr->value->type == AST_TYPE::Str) {
+                    found_docstr = true;
+                    rtn_dict->keys.push_back(new AST_Str("__doc__"));
+                    rtn_dict->values.push_back(first_expr->value);
+                }
+            }
+            if (!found_docstr) {
+                rtn_dict->keys.push_back(new AST_Str("__doc__"));
+                rtn_dict->values.push_back(makeName("None", AST_TYPE::Load));
+            }
+        }
+
+        // Even if the user never explicitly wrote to __module__, there was an
+        // implicit write:
+        assert(written_names.count("__module__"));
+
+        for (auto s : written_names) {
+            rtn_dict->keys.push_back(new AST_Str(s));
+            rtn_dict->values.push_back(makeName(s, AST_TYPE::Load));
+        }
+
+        AST_Return* rtn = new AST_Return();
+        rtn->value = rtn_dict;
+        visitor.push_back(rtn);
+    } else {
+        // Put a fake "return" statement at the end of every function just to make sure they all have one;
+        // we already have to support multiple return statements in a function, but this way we can avoid
+        // having to support not having a return statement:
+        AST_Return* return_stmt = new AST_Return();
+        return_stmt->lineno = return_stmt->col_offset = 0;
+        return_stmt->value = NULL;
+        visitor.push_back(return_stmt);
+    }
 
     if (VERBOSITY("cfg") >= 2) {
         printf("Before cfg checking and transformations:\n");

src/core/cfg.h

@@ -93,7 +93,8 @@ public:
     void print();
 };
 
-CFG* computeCFG(AST_TYPE::AST_TYPE root_type, std::vector<AST_stmt*> body);
+class SourceInfo;
+CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body);
 }
 
 #endif

src/runtime/types.cpp

@@ -291,11 +291,14 @@ const AllocationKind conservative_kind(&conservativeGCHandler, NULL);
 BoxedTuple* EmptyTuple;
 }
 
-extern "C" Box* createUserClass(std::string* name, Box* _base, BoxedModule* parent_module) {
+extern "C" Box* createUserClass(std::string* name, Box* _base, Box* _attr_dict) {
     assert(_base);
     assert(isSubclass(_base->cls, type_cls));
     BoxedClass* base = static_cast<BoxedClass*>(_base);
 
+    ASSERT(_attr_dict->cls == dict_cls, "%s", getTypeName(_attr_dict)->c_str());
+    BoxedDict* attr_dict = static_cast<BoxedDict*>(_attr_dict);
+
     BoxedClass* made;
 
     if (base->instancesHaveAttrs()) {
@@ -305,10 +308,14 @@ extern "C" Box* createUserClass(std::string* name, Box* _base, BoxedModule* pare
         made = new BoxedClass(base, base->instance_size, base->instance_size + sizeof(HCAttrs), true);
     }
 
-    made->giveAttr("__name__", boxString(*name));
+    for (const auto& p : attr_dict->d) {
+        assert(p.first->cls == str_cls);
+        made->giveAttr(static_cast<BoxedString*>(p.first)->s, p.second);
+    }
+
+    // Note: make sure to do this after assigning the attrs, since it will overwrite any defined __name__
+    made->setattr("__name__", boxString(*name), NULL);
 
-    Box* modname = parent_module->getattr("__name__", NULL, NULL);
-    made->giveAttr("__module__", modname);
 
     return made;
 }
@@ -604,6 +611,8 @@ BoxedModule* createModule(const std::string& name, const std::string& fn) {
     Box* b_name = boxStringPtr(&name);
     assert(d->d.count(b_name) == 0);
     d->d[b_name] = module;
+
+    module->giveAttr("__doc__", None);
     return module;
 }
 

src/runtime/types.h

@@ -90,7 +90,7 @@ extern "C" void listAppendInternal(Box* self, Box* v);
 extern "C" void listAppendArrayInternal(Box* self, Box** v, int nelts);
 extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, std::initializer_list<Box*> defaults);
 extern "C" CLFunction* unboxCLFunction(Box* b);
-extern "C" Box* createUserClass(std::string* name, Box* base, BoxedModule* parent_module);
+extern "C" Box* createUserClass(std::string* name, Box* base, Box* attr_dict);
 extern "C" double unboxFloat(Box* b);
 extern "C" Box* createDict();
 extern "C" Box* createList();

test/tests/classdef_arbitrary.py

-# expected: fail
-# - arbitrary stuff in classes
-
 # You can put arbitrary stuff in class definitions, which end up being added as class attributes
 
 class C(object):
@@ -27,6 +24,7 @@ class C(object):
         pass
 
     [123]
+print C.__module__
 
 class D(object):
     x = 1

test/tests/docstrings.py

+print __doc__
+__doc__ = "module_doc"
+print __doc__
+
 class C1(object):
+    print 1, __doc__
     "hello world"
+    print 2, __doc__
 print C1.__doc__
 
 class C2(object):
@@ -8,8 +14,21 @@ class C2(object):
 print C2.__doc__
 
 class C3(object):
+    print __doc__
     pass
-print C3.__doc__
+print "C3", C3.__doc__
+
+class C4(object):
+    print 1, __doc__
+    "doc1"
+    print 2, __doc__
+    __doc__ = "doc2"
+    print 3, __doc__
+print C4.__doc__
+
+class C5(object):
+    __doc__ = "doc2"
+print C5.__doc__
 
 """
 # Not supported yet:

test/tests/return_in_class.py

-# expected: fail
-# - arbitrary stuff in classes
-
 class C(object):
     return

test/tests/scoping_classes.py

 # expected: fail
-# - arbitrary stuff in classes
+# - WIP
 
 X = 0
 Y = 0
+Z = 0
+W = 0
+
+def f(val, desc):
+    print desc
+    return val
+
 def wrapper():
     X = 1
     Y = 1
-    class C(object):
+    Z = 1
+    W = 1
+
+    print "starting classdef"
+    class C(f(object, "evaluating bases")):
+        print __doc__, __module__, __name__ # "None __main__ __main__"
+
+        print "inside classdef"
         global Y
         X = 2
         Y = 2
-        def f(self):
+
+        if 0:
+            Z = 2
+
+        # class scopes have different scoping rules than function scopes (!!):
+        # In a function scope, if the name has a store but isn't set on this path,
+        # referencing it raises a NameError.
+        # In a class scope, the lookup resolves to the global scope.
+        print Z, W # "0 1".
+
+        # The defaults for a and b should resolve to the classdef definitions, and the default
+        # for c should resolve to the wrapper() definition
+        def f(self, a=X, b=Y, c=Z, d=W):
+            print a, b, c, W # "2 2 0 1"
             # These references should skip all of the classdef directives,
             # and hit the definitions in the wrapper() function
-            print X
-            print Y
+            print X, Y # "1 1"
+    print "done with classdef"
+
     return C
 
 wrapper()().f()
 
-print X
-print Y # got changed in classdef for C
+print X # "0"
+print Y # "2" -- got changed in classdef for C
+
+print
+class C2(object):
+    print __name__
+    __name__ = 1
+    print __name__
+print C2.__name__
+
+print
+class C3(object):
+    print __module__
+    __module__ = 1
+    print __module__
+print C3.__module__
+
+"""
+# not supported (del)
+print
+class C4(object):
+    print __module__
+    del __module__
+    try:
+        print __module__ # this should throw a NameError
+        assert 0
+    except NameError:
+        pass
+print C4.__module__
+"""
+
+class C5(object):
+    try:
+        print not_defined
+    except NameError:
+        print "threw NameError as expected"
+
+def make_class(add_z):
+    class C(object):
+        if add_z:
+            Z = 1
+    return C
+print hasattr(make_class(False), "Z")
+print hasattr(make_class(True), "Z")
+

test/unittests/analysis.cpp

@@ -14,8 +14,16 @@
 
 using namespace pyston;
 
-TEST(func_analysis, augassign) {
-    AST_Module* module = caching_parse("../test/unittests/analysis_listcomp.py");
+class AnalysisTest : public ::testing::Test {
+protected:
+    virtual void SetUp() {
+        initCodegen();
+    }
+};
+
+TEST_F(AnalysisTest, augassign) {
+    const std::string fn("../test/unittests/analysis_listcomp.py");
+    AST_Module* module = caching_parse(fn.c_str());
     assert(module);
 
     ScopingAnalysis *scoping = runScopingAnalysis(module);
@@ -27,7 +35,9 @@ TEST(func_analysis, augassign) {
     ASSERT_FALSE(scope_info->refersToGlobal("a"));
     ASSERT_FALSE(scope_info->refersToGlobal("b"));
 
-    CFG* cfg = computeCFG(func->type, func->body);
+    SourceInfo* si = new SourceInfo(createModule("__main__", fn), scoping, func);
+
+    CFG* cfg = computeCFG(si, func->body);
     LivenessAnalysis* liveness = computeLivenessInfo(cfg);
 
     //cfg->print();

test/unittests/unittests.h

@@ -18,6 +18,7 @@
 #include "gtest/gtest.h"
 
 #include "core/threading.h"
+#include "codegen/entry.h"
 
 namespace pyston {