Feature/switch statement js compiler (#440)

* Implements JS->FuzzIL translatability for switch statements
* Implements compiler tests for the switch statement
* Improves Compiler and JS Parser

Sources/Fuzzilli/Compiler/Compiler.swift

@@ -57,6 +57,9 @@ public class JavaScriptCompiler {
     /// The next free FuzzIL variable.
     private var nextVariable = 0
 
+    /// Context analyzer to track the context of the code being compiled. Used for example to distinguish switch and loop breaks.
+    private var contextAnalyzer = ContextAnalyzer()
+
     public func compile(_ ast: AST) throws -> Program {
         reset()
 
@@ -441,8 +444,16 @@ 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())
+            // If we're in both .loop and .switch context, then the loop must be the most recent context 
+            // (switch blocks don't propagate an outer .loop context) so we just need to check for .loop here
+            // TODO remove this comment once the Analyzer bug fixs has been merged. Until then the code in this switch case is buggy.
+            if contextAnalyzer.context.contains(.loop){
+                emit(LoopBreak())
+            } else if contextAnalyzer.context.contains(.switchBlock){
+                emit(SwitchBreak())
+            } else {
+                throw CompilerError.invalidNodeError("break statement outside of loop or switch")
+            }
 
         case .continueStatement:
             emit(LoopContinue())
@@ -486,6 +497,35 @@ public class JavaScriptCompiler {
                 try compileBody(withStatement.body)
             }
             emit(EndWith())
+        case .switchStatement(let switchStatement):
+            // TODO Replace the precomputation of tests with compilation of the test expressions in the cases.
+            // To do this, we would need to redesign Switch statements in FuzzIL to (for example) have a BeginSwitchCaseHead, BeginSwitchCaseBody, and EndSwitchCase. 
+            // Then the expression would go inside the header.
+            var precomputedTests = [Variable]()
+            for caseStatement in switchStatement.cases {
+                if caseStatement.hasTest {
+                    let test = try compileExpression(caseStatement.test)
+                    precomputedTests.append(test)
+                } 
+            }
+            let discriminant = try compileExpression(switchStatement.discriminant)
+            emit(BeginSwitch(), withInputs: [discriminant])
+            for caseStatement in switchStatement.cases {
+                if caseStatement.hasTest {
+                    emit(BeginSwitchCase(), withInputs: [precomputedTests.removeFirst()])
+                } else {
+                    emit(BeginSwitchDefaultCase())
+                }
+                try enterNewScope {
+                    for statement in caseStatement.consequent {
+                        try compileStatement(statement)
+                    }
+                }
+                // We could also do an optimization here where we check if the last statement in the case is a break, and if so, we drop the last instruction
+                // and set the fallsThrough flag to false.
+                emit(EndSwitchCase(fallsThrough: true)) 
+            }
+            emit(EndSwitch())
         }
     }
 
@@ -999,6 +1039,7 @@ public class JavaScriptCompiler {
         let innerOutputs = (0..<op.numInnerOutputs).map { _ in nextFreeVariable() }
         let inouts = inputs + outputs + innerOutputs
         let instr = Instruction(op, inouts: inouts)
+        contextAnalyzer.analyze(instr)
         return code.append(instr)
     }
 

Sources/Fuzzilli/Compiler/Parser/parser.js

@@ -326,6 +326,18 @@ function parse(script, proto) {
                 withStatement.body = visitStatement(node.body);
                 return makeStatement('WithStatement', withStatement);
             }
+            case 'SwitchStatement': {
+                let switchStatement = {};
+                switchStatement.discriminant = visitExpression(node.discriminant);
+                switchStatement.cases = node.cases.map(visitStatement);
+                return makeStatement('SwitchStatement', switchStatement);
+            }
+            case 'SwitchCase': {
+                let switchCase = {};
+                if (node.test) {switchCase.test = visitExpression(node.test)}
+                switchCase.consequent = node.consequent.map(visitStatement);
+                return switchCase;
+            }
             default: {
                 throw "Unhandled node type " + node.type;
             }

Sources/Fuzzilli/Protobuf/ast.pb.swift

@@ -905,6 +905,55 @@ public struct Compiler_Protobuf_WithStatement: @unchecked Sendable {
   fileprivate var _storage = _StorageClass.defaultInstance
 }
 
+public struct Compiler_Protobuf_SwitchStatement: @unchecked Sendable {
+  // SwiftProtobuf.Message conformance is added in an extension below. See the
+  // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
+  // methods supported on all messages.
+
+  public var discriminant: Compiler_Protobuf_Expression {
+    get {return _storage._discriminant ?? Compiler_Protobuf_Expression()}
+    set {_uniqueStorage()._discriminant = newValue}
+  }
+  /// Returns true if `discriminant` has been explicitly set.
+  public var hasDiscriminant: Bool {return _storage._discriminant != nil}
+  /// Clears the value of `discriminant`. Subsequent reads from it will return its default value.
+  public mutating func clearDiscriminant() {_uniqueStorage()._discriminant = nil}
+
+  public var cases: [Compiler_Protobuf_SwitchCase] {
+    get {return _storage._cases}
+    set {_uniqueStorage()._cases = newValue}
+  }
+
+  public var unknownFields = SwiftProtobuf.UnknownStorage()
+
+  public init() {}
+
+  fileprivate var _storage = _StorageClass.defaultInstance
+}
+
+public struct Compiler_Protobuf_SwitchCase: Sendable {
+  // SwiftProtobuf.Message conformance is added in an extension below. See the
+  // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
+  // methods supported on all messages.
+
+  public var test: Compiler_Protobuf_Expression {
+    get {return _test ?? Compiler_Protobuf_Expression()}
+    set {_test = newValue}
+  }
+  /// Returns true if `test` has been explicitly set.
+  public var hasTest: Bool {return self._test != nil}
+  /// Clears the value of `test`. Subsequent reads from it will return its default value.
+  public mutating func clearTest() {self._test = nil}
+
+  public var consequent: [Compiler_Protobuf_Statement] = []
+
+  public var unknownFields = SwiftProtobuf.UnknownStorage()
+
+  public init() {}
+
+  fileprivate var _test: Compiler_Protobuf_Expression? = nil
+}
+
 public struct Compiler_Protobuf_Statement: @unchecked Sendable {
   // SwiftProtobuf.Message conformance is added in an extension below. See the
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
@@ -1059,6 +1108,14 @@ public struct Compiler_Protobuf_Statement: @unchecked Sendable {
     set {_uniqueStorage()._statement = .withStatement(newValue)}
   }
 
+  public var switchStatement: Compiler_Protobuf_SwitchStatement {
+    get {
+      if case .switchStatement(let v)? = _storage._statement {return v}
+      return Compiler_Protobuf_SwitchStatement()
+    }
+    set {_uniqueStorage()._statement = .switchStatement(newValue)}
+  }
+
   public var unknownFields = SwiftProtobuf.UnknownStorage()
 
   public enum OneOf_Statement: Equatable, Sendable {
@@ -1080,6 +1137,7 @@ public struct Compiler_Protobuf_Statement: @unchecked Sendable {
     case tryStatement(Compiler_Protobuf_TryStatement)
     case throwStatement(Compiler_Protobuf_ThrowStatement)
     case withStatement(Compiler_Protobuf_WithStatement)
+    case switchStatement(Compiler_Protobuf_SwitchStatement)
 
   }
 
@@ -4019,6 +4077,132 @@ extension Compiler_Protobuf_WithStatement: SwiftProtobuf.Message, SwiftProtobuf.
   }
 }
 
+extension Compiler_Protobuf_SwitchStatement: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
+  public static let protoMessageName: String = _protobuf_package + ".SwitchStatement"
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "discriminant"),
+    2: .same(proto: "cases"),
+  ]
+
+  fileprivate class _StorageClass {
+    var _discriminant: Compiler_Protobuf_Expression? = nil
+    var _cases: [Compiler_Protobuf_SwitchCase] = []
+
+    #if swift(>=5.10)
+      // This property is used as the initial default value for new instances of the type.
+      // The type itself is protecting the reference to its storage via CoW semantics.
+      // This will force a copy to be made of this reference when the first mutation occurs;
+      // hence, it is safe to mark this as `nonisolated(unsafe)`.
+      static nonisolated(unsafe) let defaultInstance = _StorageClass()
+    #else
+      static let defaultInstance = _StorageClass()
+    #endif
+
+    private init() {}
+
+    init(copying source: _StorageClass) {
+      _discriminant = source._discriminant
+      _cases = source._cases
+    }
+  }
+
+  fileprivate mutating func _uniqueStorage() -> _StorageClass {
+    if !isKnownUniquelyReferenced(&_storage) {
+      _storage = _StorageClass(copying: _storage)
+    }
+    return _storage
+  }
+
+  public mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
+    _ = _uniqueStorage()
+    try withExtendedLifetime(_storage) { (_storage: _StorageClass) in
+      while let fieldNumber = try decoder.nextFieldNumber() {
+        // The use of inline closures is to circumvent an issue where the compiler
+        // allocates stack space for every case branch when no optimizations are
+        // enabled. https://github.com/apple/swift-protobuf/issues/1034
+        switch fieldNumber {
+        case 1: try { try decoder.decodeSingularMessageField(value: &_storage._discriminant) }()
+        case 2: try { try decoder.decodeRepeatedMessageField(value: &_storage._cases) }()
+        default: break
+        }
+      }
+    }
+  }
+
+  public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    try withExtendedLifetime(_storage) { (_storage: _StorageClass) in
+      // The use of inline closures is to circumvent an issue where the compiler
+      // allocates stack space for every if/case branch local when no optimizations
+      // are enabled. https://github.com/apple/swift-protobuf/issues/1034 and
+      // https://github.com/apple/swift-protobuf/issues/1182
+      try { if let v = _storage._discriminant {
+        try visitor.visitSingularMessageField(value: v, fieldNumber: 1)
+      } }()
+      if !_storage._cases.isEmpty {
+        try visitor.visitRepeatedMessageField(value: _storage._cases, fieldNumber: 2)
+      }
+    }
+    try unknownFields.traverse(visitor: &visitor)
+  }
+
+  public static func ==(lhs: Compiler_Protobuf_SwitchStatement, rhs: Compiler_Protobuf_SwitchStatement) -> Bool {
+    if lhs._storage !== rhs._storage {
+      let storagesAreEqual: Bool = withExtendedLifetime((lhs._storage, rhs._storage)) { (_args: (_StorageClass, _StorageClass)) in
+        let _storage = _args.0
+        let rhs_storage = _args.1
+        if _storage._discriminant != rhs_storage._discriminant {return false}
+        if _storage._cases != rhs_storage._cases {return false}
+        return true
+      }
+      if !storagesAreEqual {return false}
+    }
+    if lhs.unknownFields != rhs.unknownFields {return false}
+    return true
+  }
+}
+
+extension Compiler_Protobuf_SwitchCase: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
+  public static let protoMessageName: String = _protobuf_package + ".SwitchCase"
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "test"),
+    2: .same(proto: "consequent"),
+  ]
+
+  public mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
+    while let fieldNumber = try decoder.nextFieldNumber() {
+      // The use of inline closures is to circumvent an issue where the compiler
+      // allocates stack space for every case branch when no optimizations are
+      // enabled. https://github.com/apple/swift-protobuf/issues/1034
+      switch fieldNumber {
+      case 1: try { try decoder.decodeSingularMessageField(value: &self._test) }()
+      case 2: try { try decoder.decodeRepeatedMessageField(value: &self.consequent) }()
+      default: break
+      }
+    }
+  }
+
+  public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    // The use of inline closures is to circumvent an issue where the compiler
+    // allocates stack space for every if/case branch local when no optimizations
+    // are enabled. https://github.com/apple/swift-protobuf/issues/1034 and
+    // https://github.com/apple/swift-protobuf/issues/1182
+    try { if let v = self._test {
+      try visitor.visitSingularMessageField(value: v, fieldNumber: 1)
+    } }()
+    if !self.consequent.isEmpty {
+      try visitor.visitRepeatedMessageField(value: self.consequent, fieldNumber: 2)
+    }
+    try unknownFields.traverse(visitor: &visitor)
+  }
+
+  public static func ==(lhs: Compiler_Protobuf_SwitchCase, rhs: Compiler_Protobuf_SwitchCase) -> Bool {
+    if lhs._test != rhs._test {return false}
+    if lhs.consequent != rhs.consequent {return false}
+    if lhs.unknownFields != rhs.unknownFields {return false}
+    return true
+  }
+}
+
 extension Compiler_Protobuf_Statement: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".Statement"
   public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
@@ -4040,6 +4224,7 @@ extension Compiler_Protobuf_Statement: SwiftProtobuf.Message, SwiftProtobuf._Mes
     16: .same(proto: "tryStatement"),
     17: .same(proto: "throwStatement"),
     18: .same(proto: "withStatement"),
+    19: .same(proto: "switchStatement"),
   ]
 
   fileprivate class _StorageClass {
@@ -4311,6 +4496,19 @@ extension Compiler_Protobuf_Statement: SwiftProtobuf.Message, SwiftProtobuf._Mes
             _storage._statement = .withStatement(v)
           }
         }()
+        case 19: try {
+          var v: Compiler_Protobuf_SwitchStatement?
+          var hadOneofValue = false
+          if let current = _storage._statement {
+            hadOneofValue = true
+            if case .switchStatement(let m) = current {v = m}
+          }
+          try decoder.decodeSingularMessageField(value: &v)
+          if let v = v {
+            if hadOneofValue {try decoder.handleConflictingOneOf()}
+            _storage._statement = .switchStatement(v)
+          }
+        }()
         default: break
         }
       }
@@ -4396,6 +4594,10 @@ extension Compiler_Protobuf_Statement: SwiftProtobuf.Message, SwiftProtobuf._Mes
         guard case .withStatement(let v)? = _storage._statement else { preconditionFailure() }
         try visitor.visitSingularMessageField(value: v, fieldNumber: 18)
       }()
+      case .switchStatement?: try {
+        guard case .switchStatement(let v)? = _storage._statement else { preconditionFailure() }
+        try visitor.visitSingularMessageField(value: v, fieldNumber: 19)
+      }()
       case nil: break
       }
     }

Sources/Fuzzilli/Protobuf/ast.proto

@@ -207,6 +207,16 @@ message WithStatement {
     Statement body = 2;
 }
 
+message SwitchStatement {
+    Expression discriminant = 1;
+    repeated SwitchCase cases = 2;
+}
+
+message SwitchCase {
+    Expression test = 1;
+    repeated Statement consequent = 2;
+}
+
 message Statement {
     oneof statement {
         EmptyStatement emptyStatement = 1;
@@ -227,6 +237,7 @@ message Statement {
         TryStatement tryStatement = 16;
         ThrowStatement throwStatement = 17;
         WithStatement withStatement = 18;
+        SwitchStatement switchStatement = 19;
     }
 }
 

Tests/FuzzilliTests/CompilerTests/switch_statements.js

@@ -0,0 +1,26 @@
+let fruit = 'apple';
+for (let i = 0; i < 3; i++) {
+  switch (fruit) {
+    case 'apple': // test if this case falls through
+      console.log('You selected an apple.');
+      for (let j = 0; j < 2; j++) {
+          console.log('Inside apple loop', j);
+          if (j === 1) {
+              break; // test if this break exits the inner loop
+          }
+      }
+    case null:
+        console.log('You selected null.');
+        break; // Babel parses default case as null case. Try to confuse the compiler.
+    default: // test if default case is detected (irrespective of the convention that the last case is the default case)
+      console.log('Unknown fruit selection.'); // test falls through
+      break; // test if this break exits the switch
+    case 'banana':
+      console.log('You selected a banana.');
+      break; // test if this break exits the switch
+
+  }
+  if (i === 2) {
+      break; // test if this break exits the outer loop
+  }
+}