[dart2wasm] Generate br_table for some switch statements.

Often switches over enums are exhaustive or nearly exhaustive. The current generator uses identity for the enum values which requires iterating over each value to compare it to the test expression. So each time the switch body is entered is an O(n) operation (where is n is the # of case expressions).

This now generates a br_table using the index of the enum and jumps directly to the correct clause, effectively an O(1) operation.

A similar approach is taken for switches over an integer range. If the range of case expresison values is close in size to the # of values, it is advantageous to normalize the range around 0 and treat the values themselves as table indices.

This approach will also save code size when the index range is dense as the br_table is more compact than the identity/br_if checks. For sparse ranges this may produce a bit more code though only on the order of a few bytes per value in the range.

I've added a denseness heuristic to decide when to revert to the current strategy to avoid the code size penalty.

Golem benchmark: https://golem.corp.goog/Comparison?repository=dart#targetA%3Ddart2wasm-O2-d8%3BmachineTypeA%3Dlinux-x64%3BrevisionA%3D117402%3BpatchA%3Dnatebiggs-dart2wasm-switch-tables3%3BtargetB%3Ddart2wasm-O2-d8%3BmachineTypeB%3Dlinux-x64%3BrevisionB%3D117401%3BpatchB%3DNone

See 100% improvement in SwitchFSM.int and 90% improvement in SwitchFSM.enum.

Change-Id: Ie29e8fd59ef6235044ba5b4a4af04023d702ce57
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/441760
Reviewed-by: Ömer Ağacan <[email protected]>
Commit-Queue: Nate Biggs <[email protected]>

Author: natebiggs

pkg/dart2wasm/lib/code_generator.dart

@@ -1424,29 +1424,44 @@ abstract class AstCodeGenerator
       b.end();
     }
 
-    // Compare against all case values
-    for (SwitchCase c in node.cases) {
-      for (Expression exp in c.expressions) {
-        if (exp is NullLiteral ||
-            exp is ConstantExpression && exp.constant is NullConstant) {
-          // Null already checked, skip
-        } else {
-          switchInfo.compare(
-            switchValueNonNullableLocal,
-            () => translateExpression(exp, switchInfo.nonNullableType),
-          );
-          b.br_if(switchLabels[c]!);
+    final brTable = switchInfo.brTable;
+    if (brTable != null) {
+      // Map each entry in the range to the appropriate jump table entry.
+      final indexBlocks = <w.Label>[];
+      final caseMap = brTable.caseMap;
+      final defaultLabel =
+          defaultCase != null ? switchLabels[defaultCase]! : doneLabel;
+      for (int i = 0; i < brTable.rangeSize; i++) {
+        final c = caseMap[i];
+        indexBlocks.add(c == null ? defaultLabel : switchLabels[c]!);
+      }
+      brTable.brTableExpr(switchValueNonNullableLocal);
+      b.br_table(indexBlocks, defaultLabel);
+    } else {
+      // Compare against all case values
+      for (SwitchCase c in node.cases) {
+        for (Expression exp in c.expressions) {
+          if (exp is NullLiteral ||
+              exp is ConstantExpression && exp.constant is NullConstant) {
+            // Null already checked, skip
+          } else {
+            switchInfo.compare(
+              switchValueNonNullableLocal,
+              () => translateExpression(exp, switchInfo.nonNullableType),
+            );
+            b.br_if(switchLabels[c]!);
+          }
         }
       }
-    }
 
-    // No explicit cases matched
-    if (node.isExplicitlyExhaustive) {
-      b.unreachable();
-    } else {
-      w.Label defaultLabel =
-          defaultCase != null ? switchLabels[defaultCase]! : doneLabel;
-      b.br(defaultLabel);
+      // No explicit cases matched
+      if (node.isExplicitlyExhaustive) {
+        b.unreachable();
+      } else {
+        w.Label defaultLabel =
+            defaultCase != null ? switchLabels[defaultCase]! : doneLabel;
+        b.br(defaultLabel);
+      }
     }
 
     // Emit case bodies
@@ -4297,6 +4312,47 @@ class SwitchBackwardJumpInfo {
       : defaultLoopLabel = null;
 }
 
+/// Info needed to represent a switch statement using a br_table instruction.
+///
+/// This is used for switches on integers and enums (using their indicies). We
+/// map each option to an index in the jump table and then jump directly to the
+/// target case. This is much faster than iteratively comparing each cases's
+/// expression to the switch expression.
+///
+/// Sometimes switches over ranges of ints are sparse enough that a table would
+/// bloat the code compared to the iterative comparison.
+class BrTableInfo {
+  // This is the minimum ratio of br_table indicies to actual mapped cases. If
+  // this gets too large then most of the br_table indicies are unmapped and the
+  // extra code size is not worth using a br_table.
+  static const double _minTableSparseness = 2;
+  // This is the maximum size where it's always worth it to use a br_table.
+  // If the br_table is bigger than this then we start to check sparseness.
+  // Below this point we always accept the potential code size hit.
+  static const int _maxSparseSize = 50;
+
+  final int rangeSize;
+  final void Function(w.Local switchExprLocal) brTableExpr;
+  final Map<int, SwitchCase> caseMap;
+
+  BrTableInfo._(this.rangeSize, this.caseMap, this.brTableExpr)
+      : assert(!isTooSparse(rangeSize, caseMap));
+
+  /// Heuristically validate whether the provided table would be too sparse and
+  /// if so return null. Otherwise return the expected table.
+  static BrTableInfo? build(int rangeSize, Map<int, SwitchCase> caseMap,
+      void Function(w.Local switchExprLocal) brTableExpr) {
+    // Validate the table density and size is worth putting into a br_table.
+    if (isTooSparse(rangeSize, caseMap)) return null;
+
+    return BrTableInfo._(rangeSize, caseMap, brTableExpr);
+  }
+
+  static bool isTooSparse(int rangeSize, Map<int, SwitchCase> caseMap) =>
+      (rangeSize / caseMap.length) > _minTableSparseness &&
+      rangeSize > _maxSparseSize;
+}
+
 class SwitchInfo {
   /// Non-nullable Wasm type of the `switch` expression. Used when the
   /// expression is not nullable, and after the null check.
@@ -4323,6 +4379,11 @@ class SwitchInfo {
   /// The `null: ...` case, if exists.
   late final SwitchCase? nullCase;
 
+  /// Info needed to compile this switch statement into a wasm br_table. If null
+  /// this switch statement should not use a br_table and should use comparison
+  /// based case matching instead.
+  BrTableInfo? brTable;
+
   SwitchInfo(AstCodeGenerator codeGen, SwitchStatement node) {
     final translator = codeGen.translator;
 
@@ -4354,6 +4415,17 @@ class SwitchInfo {
             translator.typeEnvironment.isSubtypeOf(codeGen.dartTypeOf(e),
                 translator.coreTypes.typeNonNullableRawType));
 
+    void useCompareIdentity() {
+      // Object identity switch
+      nonNullableType = translator.topTypeNonNullable;
+      nullableType = translator.topType;
+      compare = (switchExprLocal, pushCaseExpr) {
+        codeGen.b.local_get(switchExprLocal);
+        pushCaseExpr();
+        codeGen.call(translator.coreTypes.identicalProcedure.reference);
+      };
+    }
+
     if (node.cases.every((c) =>
         c.expressions.isEmpty && c.isDefault ||
         c.expressions.every((e) =>
@@ -4453,6 +4525,39 @@ class SwitchInfo {
       nonNullableType = w.NumType.i64;
       nullableType =
           translator.classInfo[translator.boxedIntClass]!.nullableType;
+
+      // Calculate the range covered by the cases and create the jump table.
+      int? minValue;
+      int? maxValue;
+      Map<int, SwitchCase> caseMap = {};
+      for (final c in node.cases) {
+        for (final e in c.expressions) {
+          final value = e is IntLiteral
+              ? e.value
+              : ((e as ConstantExpression).constant as IntConstant).value;
+          caseMap[value] = c;
+          if (minValue == null || value < minValue) minValue = value;
+          if (maxValue == null || value > maxValue) maxValue = value;
+        }
+      }
+      if (maxValue != null) {
+        final range = maxValue - minValue! + 1;
+        if (minValue != 0) {
+          caseMap = caseMap.map((i, c) => MapEntry(i - minValue!, c));
+        }
+        brTable = BrTableInfo.build(range, caseMap, (switchExprLocal) {
+          codeGen.b.local_get(switchExprLocal);
+          // Normalize the value on 0 if necessary.
+          if (minValue != 0) {
+            codeGen.b.i64_const(minValue!);
+            codeGen.b.i64_sub();
+          }
+          // Now that we've normalized to 0 it should be safe to switch to i32.
+          codeGen.b.i32_wrap_i64();
+        });
+      }
+
+      // Provide a compare as a fallback in case the range is too sparse.
       compare = (switchExprLocal, pushCaseExpr) {
         codeGen.b.local_get(switchExprLocal);
         pushCaseExpr();
@@ -4467,15 +4572,61 @@ class SwitchInfo {
         pushCaseExpr();
         codeGen.call(translator.jsStringEquals.reference);
       };
+    } else if (switchExprClass.isEnum) {
+      // If this is an applicable switch over enums, create a jump table.
+      bool isValid = true;
+      final caseMap = <int, SwitchCase>{};
+      int? minIndex;
+      int maxIndex = 0;
+      outer:
+      for (final c in node.cases) {
+        for (final e in c.expressions) {
+          if (e is! ConstantExpression) {
+            isValid = false;
+            break outer;
+          }
+          final constant = e.constant;
+          if (constant is! InstanceConstant) {
+            isValid = false;
+            break outer;
+          }
+          if (constant.classNode != switchExprClass) {
+            isValid = false;
+            break outer;
+          }
+          final enumIndex =
+              (constant.fieldValues[translator.enumIndexField.fieldReference]
+                      as IntConstant)
+                  .value;
+          caseMap[enumIndex] = c;
+          if (enumIndex > maxIndex) maxIndex = enumIndex;
+          if (minIndex == null || enumIndex < minIndex) minIndex = enumIndex;
+        }
+      }
+
+      if (isValid && minIndex != null) {
+        final range = maxIndex - minIndex + 1;
+        brTable = BrTableInfo.build(range, caseMap, (switchExprLocal) {
+          codeGen.b.local_get(switchExprLocal);
+          codeGen.call(translator.enumIndexField.getterReference);
+          if (minIndex != 0) {
+            codeGen.b.i64_const(minIndex!);
+            codeGen.b.i64_sub();
+          }
+          // Now that we've normalized to 0 it should be safe to switch to i32.
+          codeGen.b.i32_wrap_i64();
+        });
+
+        nonNullableType =
+            translator.classInfo[switchExprClass]!.nonNullableType;
+        nullableType = translator.classInfo[switchExprClass]!.nullableType;
+      }
+
+      if (brTable == null) {
+        useCompareIdentity();
+      }
     } else {
-      // Object identity switch
-      nonNullableType = translator.topTypeNonNullable;
-      nullableType = translator.topType;
-      compare = (switchExprLocal, pushCaseExpr) {
-        codeGen.b.local_get(switchExprLocal);
-        pushCaseExpr();
-        codeGen.call(translator.coreTypes.identicalProcedure.reference);
-      };
+      useCompareIdentity();
     }
 
     _initializeSpecialCases(node);

pkg/dart2wasm/lib/kernel_nodes.dart

@@ -57,6 +57,8 @@ mixin KernelNodes {
   late final Class typeErrorClass = index.getClass("dart:core", "_TypeError");
   late final Class javaScriptErrorClass =
       index.getClass("dart:core", "_JavaScriptError");
+  late final Field enumIndexField =
+      index.getField('dart:core', '_Enum', 'index');
 
   // dart:core runtime type classes
   late final Class typeClass = index.getClass("dart:core", "_Type");