2 files changed, 92 insertions, 2 deletions
diff --git a/src/declarative/qml/qmlrewrite.cpp b/src/declarative/qml/qmlrewrite.cpp
index a41b571..19b88cc 100644
--- a/src/declarative/qml/qmlrewrite.cpp
+++ b/src/declarative/qml/qmlrewrite.cpp
@@ -111,12 +111,80 @@ bool RewriteBinding::visit(AST::Block *ast)
 
 bool RewriteBinding::visit(AST::ExpressionStatement *ast)
 {
-    unsigned startOfExpressionStatement = ast->firstSourceLocation().begin() - _position;
-    _writer->replace(startOfExpressionStatement, 0, QLatin1String("return "));
+    if (_loopStack.isEmpty()) {
+        unsigned startOfExpressionStatement = ast->firstSourceLocation().begin() - _position;
+        _writer->replace(startOfExpressionStatement, 0, QLatin1String("return "));
+    }
 
     return false;
 }
 
+bool RewriteBinding::visit(AST::DoWhileStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::DoWhileStatement *)
+{
+    _loopStack.removeLast();
+}
+
+bool RewriteBinding::visit(AST::WhileStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::WhileStatement *)
+{
+    _loopStack.removeLast();
+}
+
+bool RewriteBinding::visit(AST::ForStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::ForStatement *)
+{
+    _loopStack.removeLast();
+}
+
+bool RewriteBinding::visit(AST::LocalForStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::LocalForStatement *)
+{
+    _loopStack.removeLast();
+}
+
+bool RewriteBinding::visit(AST::ForEachStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::ForEachStatement *)
+{
+    _loopStack.removeLast();
+}
+
+bool RewriteBinding::visit(AST::LocalForEachStatement *ast)
+{
+    _loopStack.append(ast);
+    return true;
+}
+
+void RewriteBinding::endVisit(AST::LocalForEachStatement *)
+{
+    _loopStack.removeLast();
+}
+
 } // namespace QmlRewrite
 
 QT_END_NAMESPACE
diff --git a/src/declarative/qml/qmlrewrite_p.h b/src/declarative/qml/qmlrewrite_p.h
index b6fe017..449b219 100644
--- a/src/declarative/qml/qmlrewrite_p.h
+++ b/src/declarative/qml/qmlrewrite_p.h
@@ -76,8 +76,30 @@ protected:
 
     void accept(AST::Node *node);
     QString rewrite(QString code, unsigned position, AST::Statement *node);
+
     virtual bool visit(AST::Block *ast);
     virtual bool visit(AST::ExpressionStatement *ast);
+
+    virtual bool visit(AST::DoWhileStatement *ast);
+    virtual void endVisit(AST::DoWhileStatement *ast);
+
+    virtual bool visit(AST::WhileStatement *ast);
+    virtual void endVisit(AST::WhileStatement *ast);
+
+    virtual bool visit(AST::ForStatement *ast);
+    virtual void endVisit(AST::ForStatement *ast);
+
+    virtual bool visit(AST::LocalForStatement *ast);
+    virtual void endVisit(AST::LocalForStatement *ast);
+
+    virtual bool visit(AST::ForEachStatement *ast);
+    virtual void endVisit(AST::ForEachStatement *ast);
+
+    virtual bool visit(AST::LocalForEachStatement *ast);
+    virtual void endVisit(AST::LocalForEachStatement *ast);
+
+private:
+    QList<AST::Statement *> _loopStack;
 };
 
 } // namespace QmlRewrite