more weird js syntax handling

2 files changed, 23 insertions, 2 deletions

diff --git a/dbtests/jstests.cpp b/dbtests/jstests.cpp
index 5521e63118f..300b69a69fd 100644
--- a/dbtests/jstests.cpp
+++ b/dbtests/jstests.cpp
@@ -139,9 +139,12 @@ namespace JSTests {
 
             s->invoke( "function (){ return this.x == 17; }" , BSONObj() );
             ASSERT_EQUALS( true , s->getBoolean( "return" ) );
-
+            
             s->invoke( "function z(){ return this.x == 18; }" , BSONObj() );
             ASSERT_EQUALS( false , s->getBoolean( "return" ) );
+
+            s->invoke( "x = 5; for( ; x <10; x++){ a = 1; }" , BSONObj() );
+            ASSERT_EQUALS( 10 , s->getNumber( "x" ) );
             
             delete s;
         }
diff --git a/scripting/engine_spidermonkey.cpp b/scripting/engine_spidermonkey.cpp
index 2ebbe88e4cb..06e224a0070 100644
--- a/scripting/engine_spidermonkey.cpp
+++ b/scripting/engine_spidermonkey.cpp
@@ -208,11 +208,29 @@ namespace mongo {
             return code[8] == ' ' || code[8] == '(';
         }
 
+        bool isSimpleStatement( const string& code ){
+            if ( code.find( "return" ) != string::npos )
+                return false;
+
+            if ( code.find( ";" ) != string::npos &&
+                 code.find( ";" ) != code.rfind( ";" ) )
+                return false;
+
+            if ( code.find( "for(" ) != string::npos ||
+                 code.find( "for (" ) != string::npos ||
+                 code.find( "while (" ) != string::npos ||
+                 code.find( "while(" ) != string::npos )
+                return false;
+
+            return true;
+        }
+
         JSFunction * compileFunction( const char * code ){
             if ( ! hasFunctionIdentifier( code ) ){
                 string s = code;
-                if ( strstr( code , "return" ) == 0 )
+                if ( isSimpleStatement( s ) ){
                     s = "return " + s;
+                }
                 return JS_CompileFunction( _context , 0 , "anonymous" , 0 , 0 , s.c_str() , strlen( s.c_str() ) , "nofile_a" , 0 );
             }