Redesign For-Loops in FuzzIL

Similar to while- and do-while loops, for-loops now consists of multiple blocks:

    BeginForLoopInitializer
        // ...
        // v0 = initial value of the (single) loop variable
    BeginForLoopCondition v0 -> v1
        // v1 = current value of the (single) loop variable
        // ...
    BeginForLoopAfterthought -> v2
        // v2 = current value of the (single) loop variable
        // ...
    BeginForLoopBody -> v3
        // v3 = current value of the (single) loop variable
        // ...
    EndForLoop

A simple for-loop will be lifted to:

    for (let i = init; cond; afterthought) {
        body(i);
    }

However, more (and less) complex for loops are also possible. Some
examples include:

    for (;;) {
        body();
    }

    for (let i4 = f1(), i5 = f2(); f3(i4), f4(i5); f5(i5), f6(i4)) {
        body(i4, i5);
    }

    for (let [i6, i7] = (() => {
            const v1 = f();
            const v3 = g();
            h();
            return [v1, v3];
        })();
        i6 < i7;
        i6 += 2) {
        body(i6, i7);
    }

See the added tests for yet more examples.

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -2190,9 +2190,42 @@ public class ProgramBuilder {
         emit(EndDoWhileLoop(), withInputs: [cond])
     }
 
-    public func buildForLoop(_ start: Variable, _ comparator: Comparator, _ end: Variable, _ op: BinaryOperator, _ rhs: Variable, _ body: (Variable) -> ()) {
-        let i = emit(BeginForLoop(comparator: comparator, op: op), withInputs: [start, end, rhs]).innerOutput
-        body(i)
+    // Build a simple for loop that declares one loop variable.
+    public func buildForLoop(i initializer: () -> Variable, _ cond: (Variable) -> Variable, _ afterthought: (Variable) -> (), _ body: (Variable) -> ()) {
+        emit(BeginForLoopInitializer())
+        let initialValue = initializer()
+        var loopVar = emit(BeginForLoopCondition(numLoopVariables: 1), withInputs: [initialValue]).innerOutput
+        let cond = cond(loopVar)
+        loopVar = emit(BeginForLoopAfterthought(numLoopVariables: 1), withInputs: [cond]).innerOutput
+        afterthought(loopVar)
+        loopVar = emit(BeginForLoopBody(numLoopVariables: 1)).innerOutput
+        body(loopVar)
+        emit(EndForLoop())
+    }
+
+    // Build arbitrarily complex for loops without any loop variables.
+    public func buildForLoop(_ initializer: (() -> ())? = nil, _ cond: (() -> Variable)? = nil, _ afterthought: (() -> ())? = nil, _ body: () -> ()) {
+        emit(BeginForLoopInitializer())
+        initializer?()
+        emit(BeginForLoopCondition(numLoopVariables: 0))
+        let cond = cond?() ?? loadBool(true)
+        emit(BeginForLoopAfterthought(numLoopVariables: 0), withInputs: [cond])
+        afterthought?()
+        emit(BeginForLoopBody(numLoopVariables: 0))
+        body()
+        emit(EndForLoop())
+    }
+
+    // Build arbitrarily complex for loops with one or more loop variables.
+    public func buildForLoop(_ initializer: () -> [Variable], _ cond: (([Variable]) -> Variable)? = nil, _ afterthought: (([Variable]) -> ())? = nil, _ body: ([Variable]) -> ()) {
+        emit(BeginForLoopInitializer())
+        let initialValues = initializer()
+        var loopVars = emit(BeginForLoopCondition(numLoopVariables: initialValues.count), withInputs: initialValues).innerOutputs
+        let cond = cond?(Array(loopVars)) ?? loadBool(true)
+        loopVars = emit(BeginForLoopAfterthought(numLoopVariables: initialValues.count), withInputs: [cond]).innerOutputs
+        afterthought?(Array(loopVars))
+        loopVars = emit(BeginForLoopBody(numLoopVariables: initialValues.count)).innerOutputs
+        body(Array(loopVars))
         emit(EndForLoop())
     }
 
@@ -2209,12 +2242,12 @@ public class ProgramBuilder {
     }
 
     public func buildForOfLoop(_ obj: Variable, selecting indices: [Int64], hasRestElement: Bool = false, _ body: ([Variable]) -> ()) {
-        let instr = emit(BeginForOfWithDestructLoop(indices: indices, hasRestElement: hasRestElement), withInputs: [obj])
+        let instr = emit(BeginForOfLoopWithDestruct(indices: indices, hasRestElement: hasRestElement), withInputs: [obj])
         body(Array(instr.innerOutputs))
         emit(EndForOfLoop())
     }
 
-    public func buildRepeat(n numIterations: Int, _ body: (Variable) -> ()) {
+    public func buildRepeatLoop(n numIterations: Int, _ body: (Variable) -> ()) {
         let i = emit(BeginRepeatLoop(iterations: numIterations)).innerOutput
         body(i)
         emit(EndRepeatLoop())

Sources/Fuzzilli/CodeGen/CodeGenerators.swift

@@ -1176,15 +1176,28 @@ public let CodeGenerators: [CodeGenerator] = [
         }, while: { b.compare(loopVar, with: b.loadInt(Int64.random(in: 0...10)), using: .lessThan) })
     },
 
-    RecursiveCodeGenerator("ForLoopGenerator") { b in
-        let start = b.reuseOrLoadInt(0)
-        let end = b.reuseOrLoadInt(Int64.random(in: 0...10))
-        let step = b.reuseOrLoadInt(1)
-        b.buildForLoop(start, .lessThan, end, .Add, step) { _ in
+    RecursiveCodeGenerator("SimpleForLoopGenerator") { b in
+        b.buildForLoop(i: { b.loadInt(0) }, { i in b.compare(i, with: b.loadInt(Int64.random(in: 0...10)), using: .lessThan) }, { i in b.unary(.PostInc, i) }) { _ in
             b.buildRecursive()
         }
     },
 
+    RecursiveCodeGenerator("ComplexForLoopGenerator") { b in
+        if probability(0.5) {
+            // Generate a for-loop without any loop variables.
+            let counter = b.loadInt(10)
+            b.buildForLoop({}, { b.unary(.PostDec, counter) }) {
+                b.buildRecursive()
+            }
+        } else {
+            // Generate a for-loop with two loop variables.
+            // TODO could also generate loops with even more loop variables?
+            b.buildForLoop({ return [b.loadInt(0), b.loadInt(10)] }, { vs in b.compare(vs[0], with: vs[1], using: .lessThan) }, { vs in b.unary(.PostInc, vs[0]); b.unary(.PostDec, vs[0]) }) { _ in
+                b.buildRecursive()
+            }
+        }
+    },
+
     RecursiveCodeGenerator("ForInLoopGenerator", input: .object()) { b, obj in
         b.buildForInLoop(obj) { _ in
             b.buildRecursive()
@@ -1218,7 +1231,7 @@ public let CodeGenerators: [CodeGenerator] = [
 
     RecursiveCodeGenerator("RepeatLoopGenerator") { b in
         let numIterations = Int.random(in: 2...100)
-        b.buildRepeat(n: numIterations) { _ in
+        b.buildRepeatLoop(n: numIterations) { _ in
             b.buildRecursive()
         }
     },
@@ -1463,7 +1476,7 @@ public let CodeGenerators: [CodeGenerator] = [
             b.buildRecursive(block: 2, of: 3)
             b.doReturn(b.randomVariable())
         }
-        b.buildRepeat(n: numIterations) { i in
+        b.buildRepeatLoop(n: numIterations) { i in
             b.buildIf(b.compare(i, with: lastIteration, using: .equal)) {
                 b.buildRecursive(block: 3, of: 3, n: 3)
             }

Sources/Fuzzilli/CodeGen/ProgramTemplates.swift

@@ -44,7 +44,7 @@ public let ProgramTemplates = [
         b.build(n: genSize)
 
         // trigger JIT
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f, withArgs: b.generateCallArguments(for: signature))
         }
 
@@ -53,7 +53,7 @@ public let ProgramTemplates = [
         b.callFunction(f, withArgs: b.generateCallArguments(for: signature))
 
         // maybe trigger recompilation
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f, withArgs: b.generateCallArguments(for: signature))
         }
 
@@ -99,12 +99,12 @@ public let ProgramTemplates = [
         b.build(n: genSize)
 
         // trigger JIT for first function
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f1, withArgs: b.generateCallArguments(for: signature1))
         }
 
         // trigger JIT for second function
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f2, withArgs: b.generateCallArguments(for: signature2))
         }
 
@@ -115,12 +115,12 @@ public let ProgramTemplates = [
         b.callFunction(f1, withArgs: b.generateCallArguments(for: signature1))
 
         // maybe trigger recompilation
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f1, withArgs: b.generateCallArguments(for: signature1))
         }
 
         // maybe trigger recompilation
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { args in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f2, withArgs: b.generateCallArguments(for: signature2))
         }
 
@@ -156,7 +156,7 @@ public let ProgramTemplates = [
 
         b.callFunction(f, withArgs: initialArgs)
 
-        b.buildForLoop(b.loadInt(0), .lessThan, b.loadInt(100), .Add, b.loadInt(1)) { _ in
+        b.buildRepeatLoop(n: 100) { _ in
             b.callFunction(f, withArgs: optimizationArgs)
         }
 

Sources/Fuzzilli/Compiler/Compiler.swift

@@ -200,55 +200,48 @@ public class JavaScriptCompiler {
             emit(EndDoWhileLoop(), withInputs: [cond])
 
         case .forLoop(let forLoop):
-            // TODO change the IL to avoid this special handling.
-
-            // Process initializer.
-            let initializer = forLoop.init_p;
-            guard initializer.hasValue else {
-                throw CompilerError.invalidNodeError("Expected an initial value in the for loop initializer")
-            }
-            let start = try compileExpression(initializer.value)
+            try enterNewScope {
+                var loopVariables = [String]()
+
+                // Process initializer.
+                var initialLoopVariableValues = [Variable]()
+                emit(BeginForLoopInitializer())
+                if let initializer = forLoop.initializer {
+                    switch initializer {
+                    case .declaration(let declaration):
+                        for declarator in declaration.declarations {
+                            loopVariables.append(declarator.name)
+                            initialLoopVariableValues.append(try compileExpression(declarator.value))
+                        }
+                    case .expression(let expression):
+                        try compileExpression(expression)
+                    }
+                }
 
-            // Process test expression.
-            guard case .binaryExpression(let test) = forLoop.test.expression, let comparator = Comparator(rawValue: test.operator) else {
-                throw CompilerError.invalidNodeError("Expected a comparison as part of the test of a for loop")
-            }
-            guard case .identifier(let identifier) = test.lhs.expression else {
-                throw CompilerError.invalidNodeError("Expected an identifier as lhs of the test expression in a for loop")
-            }
-            guard identifier.name == initializer.name else {
-                throw CompilerError.invalidNodeError("Expected the lhs of the test expression in a for loop to be the loop variable")
-            }
-            let end = try compileExpression(test.rhs)
+                // Process condition.
+                var outputs = emit(BeginForLoopCondition(numLoopVariables: loopVariables.count), withInputs: initialLoopVariableValues).innerOutputs
+                zip(loopVariables, outputs).forEach({ map($0, to: $1 )})
+                let cond: Variable
+                if forLoop.hasCondition {
+                    cond = try compileExpression(forLoop.condition)
+                } else {
+                    cond = emit(LoadBoolean(value: true)).output
+                }
 
-            // Process update expression.
-            guard case .updateExpression(let update) = forLoop.update.expression else {
-                throw CompilerError.invalidNodeError("Expected an update expression as final part of a for loop")
-            }
-            guard case .identifier(let identifier) = update.argument.expression else {
-                throw CompilerError.invalidNodeError("Expected an identifier as argument to the update expression in a for loop")
-            }
-            guard identifier.name == initializer.name else {
-                throw CompilerError.invalidNodeError("Expected the update expression in a for loop to update the loop variable")
-            }
-            let one = emit(LoadInteger(value: 1)).output
-            let op: BinaryOperator
-            switch update.operator {
-            case "++":
-                op = .Add
-            case "--":
-                op = .Sub
-            default:
-                throw CompilerError.invalidNodeError("Unexpected operator in for loop update: \(update.operator)")
-            }
+                // Process afterthought.
+                outputs = emit(BeginForLoopAfterthought(numLoopVariables: loopVariables.count), withInputs: [cond]).innerOutputs
+                zip(loopVariables, outputs).forEach({ remap($0, to: $1 )})
+                if forLoop.hasAfterthought {
+                    try compileExpression(forLoop.afterthought)
+                }
 
-            let loopVar = emit(BeginForLoop(comparator: comparator, op: op), withInputs: [start, end, one]).innerOutput
-            try enterNewScope {
-                map(initializer.name, to: loopVar)
+                // Process body
+                outputs = emit(BeginForLoopBody(numLoopVariables: loopVariables.count)).innerOutputs
+                zip(loopVariables, outputs).forEach({ remap($0, to: $1 )})
                 try compileBody(forLoop.body)
-            }
 
-            emit(EndForLoop())
+                emit(EndForLoop())
+            }
 
         case .forInLoop(let forInLoop):
             let initializer = forInLoop.left;
@@ -282,6 +275,13 @@ public class JavaScriptCompiler {
 
             emit(EndForOfLoop())
 
+        case .breakStatement:
+            // TODO currently we assume this is a LoopBreak, but once we support switch-statements, it could also be a SwitchBreak
+            emit(LoopBreak())
+
+        case .continueStatement:
+            emit(LoopContinue())
+
         case .tryStatement(let tryStatement):
             emit(BeginTry())
             try enterNewScope {
@@ -791,6 +791,11 @@ public class JavaScriptCompiler {
         scopes.top[identifier] = v
     }
 
+    private func remap(_ identifier: String, to v: Variable) {
+        assert(scopes.top[identifier] != nil)
+        scopes.top[identifier] = v
+    }
+
     private func mapParameters(_ parameters: [Compiler_Protobuf_Parameter], to variables: ArraySlice<Variable>) {
         assert(parameters.count == variables.count)
         for (param, v) in zip(parameters, variables) {