Implement AugAssign (ie '+=' and friends)

src/Makefile

@@ -208,9 +208,9 @@ OPTIONAL_SRCS := codegen/profiling/oprofile.cpp codegen/profiling/pprof.cpp
 TOOL_SRCS := $(wildcard $(TOOLS_DIR)/*.cpp)
 NONSTDLIB_SRCS := $(MAIN_SRCS) $(OPTIONAL_SRCS) $(TOOL_SRCS)
 
-# .DEFAULT_GOAL := pyston_dbg
-_ :
-	$(MAKE) pyston_dbg || (clear; $(MAKE) pyston_dbg -j1 ERROR_LIMIT=1)
+.DEFAULT_GOAL := pyston_dbg
+# _ :
+	# $(MAKE) pyston_dbg || (clear; $(MAKE) pyston_dbg -j1 ERROR_LIMIT=1)
 
 .PHONY: all _all
 all: llvm

src/analysis/function_analysis.cpp

@@ -204,6 +204,11 @@ class DefinednessVisitor : public ASTVisitor {
             return true;
         }
 
+        virtual bool visit_augassign(AST_AugAssign *node) {
+            _doSet(node->target);
+            return true;
+        }
+
         virtual bool visit_arguments(AST_arguments *node) {
             if (node->kwarg) _doSet(node->kwarg);
             if (node->vararg.size()) _doSet(node->vararg);

src/analysis/scoping_analysis.cpp

@@ -163,6 +163,7 @@ class NameCollectorVisitor : public ASTVisitor {
 
         virtual bool visit_arguments(AST_arguments *node) { return false; }
         virtual bool visit_assign(AST_Assign *node) { return false; }
+        virtual bool visit_augassign(AST_AugAssign *node) { return false; }
         virtual bool visit_attribute(AST_Attribute *node) { return false; }
         virtual bool visit_binop(AST_BinOp *node) { return false; }
         virtual bool visit_boolop(AST_BoolOp *node) { return false; }

src/analysis/type_analysis.cpp

@@ -346,6 +346,22 @@ class BasicBlockTypePropagator : public ExprVisitor, public StmtVisitor {
             }
         }
 
+        virtual void visit_augassign(AST_AugAssign* node) {
+            CompilerType *t = getType(node->target);
+            CompilerType *v = getType(node->value);
+
+            // TODO this isn't the right behavior
+            std::string name = getOpName(node->op_type);
+            name = "__i" + name.substr(2);
+            CompilerType *attr_type = t->getattrType(name);
+
+            std::vector<CompilerType*> arg_types;
+            arg_types.push_back(v);
+            CompilerType *rtn = attr_type->callType(arg_types);
+
+            _doSet(node->target, rtn);
+        }
+
         virtual void visit_branch(AST_Branch* node) {
             if (EXPAND_UNNEEDED) {
                 getType(node->test);

src/codegen/irgen.cpp

@@ -519,6 +519,7 @@ class IRGenerator {
         }
 
         enum BinExpType {
+            AugAssign,
             BinOp,
             Compare,
         };
@@ -536,7 +537,7 @@ class IRGenerator {
                     v = emitter.getBuilder()->CreateCall2(g.funcs.div_i64_i64, converted_left->getValue(), converted_right->getValue());
                 } else if (type == AST_TYPE::Pow) {
                     v = emitter.getBuilder()->CreateCall2(g.funcs.pow_i64_i64, converted_left->getValue(), converted_right->getValue());
-                } else if (exp_type == BinOp) {
+                } else if (exp_type == BinOp || exp_type == AugAssign) {
                     llvm::Instruction::BinaryOps binopcode;
                     switch (type) {
                         case AST_TYPE::Add:
@@ -626,7 +627,7 @@ class IRGenerator {
                     v = emitter.getBuilder()->CreateCall2(g.funcs.div_float_float, converted_left->getValue(), converted_right->getValue());
                 } else if (type == AST_TYPE::Pow) {
                     v = emitter.getBuilder()->CreateCall2(g.funcs.pow_float_float, converted_left->getValue(), converted_right->getValue());
-                } else if (exp_type == BinOp) {
+                } else if (exp_type == BinOp || exp_type == AugAssign) {
                     llvm::Instruction::BinaryOps binopcode;
                     switch (type) {
                         case AST_TYPE::Add:
@@ -707,6 +708,9 @@ class IRGenerator {
             if (exp_type == BinOp) {
                 rt_func = g.funcs.binop;
                 rt_func_addr = (void*)binop;
+            } else if (exp_type == AugAssign) {
+                rt_func = g.funcs.augassign;
+                rt_func_addr = (void*)augassign;
             } else {
                 rt_func = g.funcs.compare;
                 rt_func_addr = (void*)compare;
@@ -1485,6 +1489,29 @@ class IRGenerator {
             val->decvref(emitter);
         }
 
+        void doAugAssign(AST_AugAssign *node) {
+            CompilerVariable *target = evalExpr(node->target);
+            _setFake(_nodeFakeName(0, node), target); // 'fakes' are for handling deopt entries
+
+            CompilerVariable *val = evalExpr(node->value);
+            _setFake(_nodeFakeName(1, node), val);
+
+            if (state == PARTIAL) {
+                _clearFake(_nodeFakeName(0, node));
+                _clearFake(_nodeFakeName(1, node));
+                return;
+            }
+            target = _getFake(_nodeFakeName(0, node));
+            val = _getFake(_nodeFakeName(1, node));
+
+            CompilerVariable *rtn = this->_evalBinExp(target, val, node->op_type, AugAssign);
+            target->decvref(emitter);
+            val->decvref(emitter);
+
+            _doSet(node->target, rtn);
+            rtn->decvref(emitter);
+        }
+
         void doClassDef(AST_ClassDef *node) {
             if (state == PARTIAL)
                 return;
@@ -1621,7 +1648,7 @@ class IRGenerator {
 
             endBlock(DEAD);
 
-            assert(rtn->getVrefs() == 1);
+            ASSERT(rtn->getVrefs() == 1, "%d", rtn->getVrefs());
             emitter.getBuilder()->CreateRet(rtn->getValue());
         }
 
@@ -1820,6 +1847,9 @@ class IRGenerator {
                 case AST_TYPE::Assign:
                     doAssign(static_cast<AST_Assign*>(node));
                     break;
+                case AST_TYPE::AugAssign:
+                    doAugAssign(static_cast<AST_AugAssign*>(node));
+                    break;
                 case AST_TYPE::ClassDef:
                     doClassDef(static_cast<AST_ClassDef*>(node));
                     break;

src/codegen/parser.cpp

@@ -179,6 +179,17 @@ AST_Assign* read_assign(BufferedReader *reader) {
     return rtn;
 }
 
+AST_AugAssign* read_augassign(BufferedReader *reader) {
+    AST_AugAssign *rtn = new AST_AugAssign();
+
+    rtn->col_offset = readColOffset(reader);
+    rtn->lineno = reader->readULL();
+    rtn->op_type = (AST_TYPE::AST_TYPE)reader->readByte();
+    rtn->target = readASTExpr(reader);
+    rtn->value = readASTExpr(reader);
+    return rtn;
+}
+
 AST_Attribute* read_attribute(BufferedReader *reader) {
     AST_Attribute *rtn = new AST_Attribute();
 
@@ -622,6 +633,8 @@ AST_stmt* readASTStmt(BufferedReader *reader) {
     switch (type) {
         case AST_TYPE::Assign:
             return read_assign(reader);
+        case AST_TYPE::AugAssign:
+            return read_augassign(reader);
         case AST_TYPE::Break:
             return read_break(reader);
         case AST_TYPE::ClassDef:

src/codegen/runtime_hooks.cpp

@@ -151,6 +151,7 @@ void initGlobalFuncs(GlobalState &g) {
     GET(getGlobal);
     GET(binop);
     GET(compare);
+    GET(augassign);
     GET(nonzero);
     GET(print);
     GET(unboxedLen);

src/codegen/runtime_hooks.h

@@ -21,7 +21,7 @@ struct GlobalFuncs {
     llvm::Value *printf, *my_assert, *malloc, *free;
 
     llvm::Value *boxInt, *unboxInt, *boxFloat, *unboxFloat, *boxStringPtr, *boxCLFunction, *unboxCLFunction, *boxInstanceMethod, *boxBool, *unboxBool, *createTuple, *createDict, *createList, *createSlice, *createClass;
-    llvm::Value *getattr, *setattr, *print, *nonzero, *binop, *compare, *unboxedLen, *getitem, *getclsattr, *getGlobal, *setitem, *unaryop, *import;
+    llvm::Value *getattr, *setattr, *print, *nonzero, *binop, *compare, *augassign, *unboxedLen, *getitem, *getclsattr, *getGlobal, *setitem, *unaryop, *import;
     llvm::Value *checkUnpackingLength, *raiseAttributeError, *raiseAttributeErrorStr, *raiseNotIterableError, *assertNameDefined;
     llvm::Value *printFloat, *listAppendInternal;
     llvm::Value *dump;

src/core/ast.cpp

@@ -81,6 +81,10 @@ std::string getOpSymbol(int op_type) {
     }
 }
 
+std::string getInplaceOpSymbol(int op_type) {
+    return getOpSymbol(op_type) + '=';
+}
+
 std::string getOpName(int op_type) {
     assert(op_type != AST_TYPE::Is);
     assert(op_type != AST_TYPE::IsNot);
@@ -139,6 +143,11 @@ std::string getOpName(int op_type) {
     }
 }
 
+std::string getInplaceOpName(int op_type) {
+    std::string normal_name = getOpName(op_type);
+    return "__i" + normal_name.substr(2);
+}
+
 std::string getReverseOpName(int op_type) {
     if (op_type == AST_TYPE::Lt)
         return getOpName(AST_TYPE::GtE);
@@ -196,6 +205,18 @@ void AST_Assign::accept_stmt(StmtVisitor *v) {
     v->visit_assign(this);
 }
 
+void AST_AugAssign::accept(ASTVisitor *v) {
+    bool skip = v->visit_augassign(this);
+    if (skip) return;
+
+    value->accept(v);
+    target->accept(v);
+}
+
+void AST_AugAssign::accept_stmt(StmtVisitor *v) {
+    v->visit_augassign(this);
+}
+
 void AST_Attribute::accept(ASTVisitor *v) {
     bool skip = v->visit_attribute(this);
     if (skip) return;
@@ -616,16 +637,8 @@ bool PrintVisitor::visit_assign(AST_Assign *node) {
     return true;
 }
 
-bool PrintVisitor::visit_attribute(AST_Attribute *node) {
-    node->value->accept(this);
-    putchar('.');
-    printf("%s", node->attr.c_str());
-    return true;
-}
-
-bool PrintVisitor::visit_binop(AST_BinOp *node) {
-    node->left->accept(this);
-    switch (node->op_type) {
+static void printOp(AST_TYPE::AST_TYPE op_type) {
+    switch (op_type) {
         case AST_TYPE::Add:
             putchar('+');
             break;
@@ -660,9 +673,29 @@ bool PrintVisitor::visit_binop(AST_BinOp *node) {
             putchar('-');
             break;
         default:
-            printf("<%d>", node->op_type);
+            printf("<%d>", op_type);
             break;
     }
+}
+
+bool PrintVisitor::visit_augassign(AST_AugAssign *node) {
+    node->target->accept(this);
+    printOp(node->op_type);
+    putchar('=');
+    node->value->accept(this);
+    return true;
+}
+
+bool PrintVisitor::visit_attribute(AST_Attribute *node) {
+    node->value->accept(this);
+    putchar('.');
+    printf("%s", node->attr.c_str());
+    return true;
+}
+
+bool PrintVisitor::visit_binop(AST_BinOp *node) {
+    node->left->accept(this);
+    printOp(node->op_type);
     node->right->accept(this);
     return true;
 }
@@ -1071,6 +1104,7 @@ class FlattenVisitor : public ASTVisitor {
         virtual bool visit_alias(AST_alias *node) { output->push_back(node); return false; }
         virtual bool visit_arguments(AST_arguments *node) { output->push_back(node); return false; }
         virtual bool visit_assign(AST_Assign *node) { output->push_back(node); return false; }
+        virtual bool visit_augassign(AST_AugAssign *node) { output->push_back(node); return false; }
         virtual bool visit_attribute(AST_Attribute *node) { output->push_back(node); return false; }
         virtual bool visit_binop(AST_BinOp *node) { output->push_back(node); return false; }
         virtual bool visit_boolop(AST_BoolOp *node) { output->push_back(node); return false; }

src/core/ast.h

@@ -185,6 +185,18 @@ class AST_Assign : public AST_stmt {
         AST_Assign() : AST_stmt(AST_TYPE::Assign) {}
 };
 
+class AST_AugAssign : public AST_stmt {
+    public:
+        AST_expr* value;
+        AST_expr* target;
+        AST_TYPE::AST_TYPE op_type;
+
+        virtual void accept(ASTVisitor *v);
+        virtual void accept_stmt(StmtVisitor *v);
+
+        AST_AugAssign() : AST_stmt(AST_TYPE::AugAssign) {}
+};
+
 class AST_Attribute : public AST_expr {
     public:
         AST_expr* value;
@@ -577,6 +589,7 @@ class ASTVisitor {
         virtual bool visit_alias(AST_alias *node) { assert(0); abort(); }
         virtual bool visit_arguments(AST_arguments *node) { assert(0); abort(); }
         virtual bool visit_assign(AST_Assign *node) { assert(0); abort(); }
+        virtual bool visit_augassign(AST_AugAssign *node) { assert(0); abort(); }
         virtual bool visit_attribute(AST_Attribute *node) { assert(0); abort(); }
         virtual bool visit_binop(AST_BinOp *node) { assert(0); abort(); }
         virtual bool visit_boolop(AST_BoolOp *node) { assert(0); abort(); }
@@ -622,6 +635,7 @@ class NoopASTVisitor : public ASTVisitor {
         virtual bool visit_alias(AST_alias *node) { return false; }
         virtual bool visit_arguments(AST_arguments *node) { return false; }
         virtual bool visit_assign(AST_Assign *node) { return false; }
+        virtual bool visit_augassign(AST_AugAssign *node) { return false; }
         virtual bool visit_attribute(AST_Attribute *node) { return false; }
         virtual bool visit_binop(AST_BinOp *node) { return false; }
         virtual bool visit_boolop(AST_BoolOp *node) { return false; }
@@ -688,6 +702,7 @@ class StmtVisitor {
         virtual ~StmtVisitor() {}
 
         virtual void visit_assign(AST_Assign *node) { assert(0); abort(); }
+        virtual void visit_augassign(AST_AugAssign *node) { assert(0); abort(); }
         virtual void visit_break(AST_Break *node) { assert(0); abort(); }
         virtual void visit_classdef(AST_ClassDef *node) { assert(0); abort(); }
         virtual void visit_continue(AST_Continue *node) { assert(0); abort(); }
@@ -719,6 +734,7 @@ class PrintVisitor : public ASTVisitor {
         virtual bool visit_alias(AST_alias *node);
         virtual bool visit_arguments(AST_arguments *node);
         virtual bool visit_assign(AST_Assign *node);
+        virtual bool visit_augassign(AST_AugAssign *node);
         virtual bool visit_attribute(AST_Attribute *node);
         virtual bool visit_binop(AST_BinOp *node);
         virtual bool visit_boolop(AST_BoolOp *node);

src/core/cfg.cpp

@@ -202,6 +202,7 @@ class CFGVisitor : public ASTVisitor {
         }
 
         virtual bool visit_assign(AST_Assign* node) { push_back(node); return true; }
+        virtual bool visit_augassign(AST_AugAssign* node) { push_back(node); return true; }
         virtual bool visit_classdef(AST_ClassDef* node) { push_back(node); return true; }
         virtual bool visit_expr(AST_Expr* node) { push_back(node); return true; }
         virtual bool visit_functiondef(AST_FunctionDef* node) { push_back(node); return true; }

src/core/types.h

@@ -224,8 +224,10 @@ extern "C" CompiledFunction* resolveCLFunc(CLFunction *f, int64_t nargs, Box* ar
 extern "C" Box* callCompiledFunc(CompiledFunction *cf, int64_t nargs, Box* arg1, Box* arg2, Box* arg3, Box** args);
 
 std::string getOpName(int op_type);
-std::string getOpSymbol(int op_type);
 std::string getReverseOpName(int op_type);
+std::string getInplaceOpName(int op_type);
+std::string getOpSymbol(int op_type);
+std::string getInplaceOpSymbol(int op_type);
 
 typedef bool i1;
 typedef int64_t i64;

src/runtime/inline/link_forcer.cpp

@@ -62,6 +62,7 @@ void force() {
     FORCE(nonzero);
     FORCE(binop);
     FORCE(compare);
+    FORCE(augassign);
     FORCE(unboxedLen);
     FORCE(getitem);
     FORCE(getclsattr);

src/runtime/list.cpp

@@ -230,6 +230,43 @@ Box* listMul(BoxedList* self, Box* rhs) {
     return rtn;
 }
 
+Box* listIAdd(BoxedList* self, Box* _rhs) {
+    if (_rhs->cls != list_cls) {
+        fprintf(stderr, "TypeError: can only concatenate list (not \"%s\") to list\n", getTypeName(_rhs)->c_str());
+        raiseExc();
+    }
+
+    BoxedList* rhs = static_cast<BoxedList*>(_rhs);
+
+    int s1 = self->size;
+    int s2 = rhs->size;
+    self->ensure(s1 + s2);
+
+    memcpy(self->elts->elts + s1, rhs->elts->elts, sizeof(rhs->elts->elts[0]) * s2);
+    self->size = s1 + s2;
+    return self;
+}
+
+Box* listAdd(BoxedList* self, Box* _rhs) {
+    if (_rhs->cls != list_cls) {
+        fprintf(stderr, "TypeError: can only concatenate list (not \"%s\") to list\n", getTypeName(_rhs)->c_str());
+        raiseExc();
+    }
+
+    BoxedList* rhs = static_cast<BoxedList*>(_rhs);
+
+    BoxedList* rtn = new BoxedList();
+
+    int s1 = self->size;
+    int s2 = rhs->size;
+    rtn->ensure(s1 + s2);
+
+    memcpy(rtn->elts->elts, self->elts->elts, sizeof(self->elts->elts[0]) * s1);
+    memcpy(rtn->elts->elts + s1, rhs->elts->elts, sizeof(rhs->elts->elts[0]) * s2);
+    rtn->size = s1 + s2;
+    return rtn;
+}
+
 BoxedClass *list_iterator_cls = NULL;
 extern "C" void listIteratorGCHandler(GCVisitor *v, void* p) {
     boxGCHandler(v, p);
@@ -289,6 +326,9 @@ void setupList() {
     list_cls->giveAttr("insert", new BoxedFunction(boxRTFunction((void*)listInsert, NULL, 3, false)));
     list_cls->giveAttr("__mul__", new BoxedFunction(boxRTFunction((void*)listMul, NULL, 2, false)));
 
+    list_cls->giveAttr("__iadd__", new BoxedFunction(boxRTFunction((void*)listIAdd, NULL, 2, false)));
+    list_cls->giveAttr("__add__", new BoxedFunction(boxRTFunction((void*)listAdd, NULL, 2, false)));
+
     CLFunction *new_ = boxRTFunction((void*)listNew1, NULL, 1, false);
     addRTFunction(new_, (void*)listNew2, NULL, 2, false);
     list_cls->giveAttr("__new__", new BoxedFunction(new_));

src/runtime/objmodel.h

@@ -54,6 +54,7 @@ extern "C" void dump(Box* obj);
 //extern "C" Box* trap();
 extern "C" i64 unboxedLen(Box* obj);
 extern "C" Box* binop(Box* lhs, Box* rhs, int op_type);
+extern "C" Box* augassign(Box* lhs, Box* rhs, int op_type);
 extern "C" Box* getGlobal(BoxedModule* m, std::string *name, bool from_global);
 extern "C" Box* getitem(Box* value, Box* slice);
 extern "C" void setitem(Box* target, Box* slice, Box* value);

src/runtime/str.cpp

@@ -30,9 +30,15 @@
 
 namespace pyston {
 
-extern "C" BoxedString* strAdd(BoxedString* lhs, BoxedString* rhs) {
+extern "C" BoxedString* strAdd(BoxedString* lhs, Box* _rhs) {
     assert(lhs->cls == str_cls);
-    assert(rhs->cls == str_cls);
+
+    if (_rhs->cls != str_cls) {
+        fprintf(stderr, "TypeError: cannot concatenate 'str' and '%s' objects", getTypeName(_rhs)->c_str());
+        raiseExc();
+    }
+
+    BoxedString* rhs = static_cast<BoxedString*>(_rhs);
     return new BoxedString(lhs->s + rhs->s);
 }
 

test/tests/.gitignore

+t.py
+t2.py

test/tests/augassign.py

+# Augassign
+
+def f(a, b):
+    a += b
+    a -= b
+    a *= b
+    a /= b
+    a //= b
+    a **= b
+    a %= b
+
+    a |= b
+    a &= b
+    a ^= b
+    a <<= b
+    a >>= b
+    return a
+
+l = range(10)
+l2 = l
+l += l
+print l
+print l2
+
+l = range(10)
+l2 = l
+l = l + l
+print l
+print l2
+
+print f(4, 2)
+print f(4.1, 2.3)

test/tests/augassign_reverse.py

+class IntLike(object):
+    def __radd__(self, rhs):
+        return "hello world"
+
+def f():
+    i = 1
+    print i
+
+    # Augassigns can change the type of the variable:
+    i += IntLike()
+    print i
+    i + 1
+f()

test/tests/t.py

-def f():
-    return range(10)
-
-for i in xrange(1000000):
-    f()
-
-print "This will break"

tools/tester.py

@@ -246,6 +246,8 @@ if __name__ == "__main__":
             raise Exception((t, v))
 
     TEST_DIR = patterns[0]
+    assert os.path.isdir(TEST_DIR), "%s doesn't look like a directory with tests in it" % TEST_DIR
+
     patterns = patterns[1:]
 
     TOSKIP = ["%s/%s.py" % (TEST_DIR, i) for i in (