Reland "[pkg/vm] Handle switch statements in unreachable code eliminator."

This is a reland of commit 92bf76d

In the original CL, the changes to the UCE assumed all
SwitchStatements were exhaustive. Now if a SwitchStatement isn't
explicitly or implicitly (via a default case) exhaustive and no case
matches the tested constant, the SwitchStatement is properly removed.

In addition, if a guaranteed to match case is found, more than one
cases remain (thus the SwitchStatement is not removed or replaced with
the single case's body), and the default case is removed, then the
resulting SwitchStatement is marked as explicitly exhaustive, as this
serves as a signal to backends that they do not need to handle the
possibility of no case matching in the absence of a default case.

Original change's description:
> [pkg/vm] Handle switch statements in unreachable code eliminator.
>
> Namely, if the tested expression for a switch statement is constant,
> then we can remove any constant cases where the constants differ,
> and if all but a single case is removed, we can replace the switch
> with the case body.
>
> If constant functions are not enabled, then getters annotated with
> @pragma("vm:platform-const") are still evaluated with the constant
> function evaluation machinery, but only those and no others (including
> any functions called within an annotated getter). This way, functions
> can be annotated with @pragma("vm:platform-const") without having to
> rewrite them to be a single returned expression.
>
> TEST=pkg/vm/test/transformations/unreachable_code_elimination
>      pkg/vm/test/transformations/vm_constant_evaluator
>
> Issue: #50473
> Issue: #31969
> Change-Id: Ie290d2f1f469326238d66c3d9631f8e696685ff0
> Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/332760
> Commit-Queue: Tess Strickland <[email protected]>
> Reviewed-by: Chloe Stefantsova <[email protected]>
> Reviewed-by: Alexander Markov <[email protected]>

TEST=pkg/vm/test/transformations/unreachable_code_elimination
     pkg/vm/test/transformations/vm_constant_evaluator

Issue: #50473
Issue: #31969
Cq-Include-Trybots: luci.dart.try:vm-aot-linux-release-x64-try,vm-aot-mac-release-arm64-try
Change-Id: I557ca933808012e670e306f2d880221a0d7dd670
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/334224
Reviewed-by: Chloe Stefantsova <[email protected]>
Reviewed-by: Alexander Markov <[email protected]>
Commit-Queue: Tess Strickland <[email protected]>

Author: sstrickl

pkg/front_end/lib/src/fasta/kernel/constant_evaluator.dart

@@ -2631,6 +2631,11 @@ class ConstantEvaluator implements ExpressionVisitor<Constant> {
 
   /// Execute a function body using the [StatementConstantEvaluator].
   Constant executeBody(Statement statement) {
+    if (!enableConstFunctions && !inExtensionTypeConstConstructor) {
+      throw new UnsupportedError("Statement evaluation is only supported when "
+          "in extension type const constructors or when the const functions "
+          "feature is enabled.");
+    }
     StatementConstantEvaluator statementEvaluator =
         new StatementConstantEvaluator(this);
     ExecutionStatus status = statement.accept(statementEvaluator);
@@ -2657,6 +2662,11 @@ class ConstantEvaluator implements ExpressionVisitor<Constant> {
   /// Returns [null] on success and an error-"constant" on failure, as such the
   /// return value should be checked.
   AbortConstant? executeConstructorBody(Constructor constructor) {
+    if (!enableConstFunctions && !inExtensionTypeConstConstructor) {
+      throw new UnsupportedError("Statement evaluation is only supported when "
+          "in extension type const constructors or when the const functions "
+          "feature is enabled.");
+    }
     final Statement body = constructor.function.body!;
     StatementConstantEvaluator statementEvaluator =
         new StatementConstantEvaluator(this);
@@ -5477,14 +5487,7 @@ class ConstantEvaluator implements ExpressionVisitor<Constant> {
 class StatementConstantEvaluator implements StatementVisitor<ExecutionStatus> {
   ConstantEvaluator exprEvaluator;
 
-  StatementConstantEvaluator(this.exprEvaluator) {
-    if (!exprEvaluator.enableConstFunctions &&
-        !exprEvaluator.inExtensionTypeConstConstructor) {
-      throw new UnsupportedError("Statement evaluation is only supported when "
-          "in inline class const constructors or when the const functions "
-          "feature is enabled.");
-    }
-  }
+  StatementConstantEvaluator(this.exprEvaluator);
 
   /// Evaluate the expression using the [ConstantEvaluator].
   Constant evaluate(Expression expr) => expr.accept(exprEvaluator);

pkg/vm/lib/transformations/unreachable_code_elimination.dart

@@ -36,25 +36,32 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
     return result;
   }
 
-  bool _isBoolConstant(Expression node) =>
-      node is BoolLiteral ||
-      (node is ConstantExpression && node.constant is BoolConstant);
-
-  bool _getBoolConstantValue(Expression node) {
-    if (node is BoolLiteral) {
-      return node.value;
-    }
-    if (node is ConstantExpression) {
-      final constant = node.constant;
-      if (constant is BoolConstant) {
-        return constant.value;
-      }
+  bool? _getBoolConstantValue(Expression node) {
+    if (node is BoolLiteral) return node.value;
+    if (node is! ConstantExpression) return null;
+    final constant = node.constant;
+    return constant is BoolConstant ? constant.value : null;
+  }
+
+  Expression _makeConstantExpression(Constant constant, Expression node) {
+    if (constant is UnevaluatedConstant &&
+        constant.expression is InvalidExpression) {
+      return constant.expression;
+    }
+    ConstantExpression constantExpression = new ConstantExpression(
+        constant, node.getStaticType(_staticTypeContext!))
+      ..fileOffset = node.fileOffset;
+    if (node is FileUriExpression) {
+      return new FileUriConstantExpression(constantExpression.constant,
+          type: constantExpression.type, fileUri: node.fileUri)
+        ..fileOffset = node.fileOffset;
     }
-    throw 'Expected bool constant: $node';
+    return constantExpression;
   }
 
-  Expression _createBoolLiteral(bool value, int fileOffset) =>
-      new BoolLiteral(value)..fileOffset = fileOffset;
+  Expression _createBoolConstantExpression(bool value, Expression node) =>
+      _makeConstantExpression(
+          constantEvaluator.canonicalize(BoolConstant(value)), node);
 
   Statement _makeEmptyBlockIfEmptyStatement(Statement node, TreeNode parent) =>
       node is EmptyStatement ? (Block(<Statement>[])..parent = parent) : node;
@@ -63,8 +70,8 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
   TreeNode visitIfStatement(IfStatement node, TreeNode? removalSentinel) {
     node.transformOrRemoveChildren(this);
     final condition = node.condition;
-    if (_isBoolConstant(condition)) {
-      final value = _getBoolConstantValue(condition);
+    final value = _getBoolConstantValue(condition);
+    if (value != null) {
       return value
           ? node.then
           : (node.otherwise ?? removalSentinel ?? new EmptyStatement());
@@ -78,8 +85,8 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
       ConditionalExpression node, TreeNode? removalSentinel) {
     node.transformOrRemoveChildren(this);
     final condition = node.condition;
-    if (_isBoolConstant(condition)) {
-      final value = _getBoolConstantValue(condition);
+    final value = _getBoolConstantValue(condition);
+    if (value != null) {
       return value ? node.then : node.otherwise;
     }
     return node;
@@ -89,9 +96,9 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
   TreeNode visitNot(Not node, TreeNode? removalSentinel) {
     node.transformOrRemoveChildren(this);
     final operand = node.operand;
-    if (_isBoolConstant(operand)) {
-      return _createBoolLiteral(
-          !_getBoolConstantValue(operand), node.fileOffset);
+    final value = _getBoolConstantValue(operand);
+    if (value != null) {
+      return _createBoolConstantExpression(!value, node);
     }
     return node;
   }
@@ -100,47 +107,111 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
   TreeNode visitLogicalExpression(
       LogicalExpression node, TreeNode? removalSentinel) {
     node.transformOrRemoveChildren(this);
-    final left = node.left;
-    final right = node.right;
-    final operatorEnum = node.operatorEnum;
-    if (_isBoolConstant(left)) {
-      final leftValue = _getBoolConstantValue(left);
-      if (_isBoolConstant(right)) {
-        final rightValue = _getBoolConstantValue(right);
-        if (operatorEnum == LogicalExpressionOperator.OR) {
-          return _createBoolLiteral(leftValue || rightValue, node.fileOffset);
-        } else if (operatorEnum == LogicalExpressionOperator.AND) {
-          return _createBoolLiteral(leftValue && rightValue, node.fileOffset);
-        } else {
-          throw 'Unexpected LogicalExpression operator ${operatorEnum}: $node';
-        }
-      } else {
-        if (leftValue && operatorEnum == LogicalExpressionOperator.OR) {
-          return _createBoolLiteral(true, node.fileOffset);
-        } else if (!leftValue &&
-            operatorEnum == LogicalExpressionOperator.AND) {
-          return _createBoolLiteral(false, node.fileOffset);
+    bool? value = _getBoolConstantValue(node.left);
+    // Because of short-circuiting, these operators cannot be treated as
+    // symmetric, so a non-constant left and a constant right is left as-is.
+    if (value == null) return node;
+    switch (node.operatorEnum) {
+      case LogicalExpressionOperator.OR:
+        return value ? node.left : node.right;
+      case LogicalExpressionOperator.AND:
+        return value ? node.right : node.left;
+    }
+  }
+
+  @override
+  TreeNode visitSwitchStatement(
+      SwitchStatement node, TreeNode? removalSentinel) {
+    node.transformOrRemoveChildren(this);
+    final tested = node.expression;
+    if (tested is! ConstantExpression) return node;
+
+    // First, keep any reachable case. As a side effect, any expressions that
+    // cannot match in the SwitchCases are removed. An expression cannot match
+    // if it is a non-matching constant expression or it follows a constant
+    // expression that is guaranteed to match.
+    final toKeep = <SwitchCase>{};
+    bool foundMatchingCase = false;
+    for (final c in node.cases) {
+      if (foundMatchingCase) {
+        c.expressions.clear();
+        continue;
+      }
+      c.expressions.retainWhere((e) {
+        if (foundMatchingCase) return false;
+        if (e is! ConstantExpression) return true;
+        foundMatchingCase = e.constant == tested.constant;
+        return foundMatchingCase;
+      });
+      if (c.isDefault || c.expressions.isNotEmpty) {
+        toKeep.add(c);
+      }
+    }
+
+    if (toKeep.isEmpty) {
+      if (node.isExhaustive) {
+        throw 'Expected at least one kept case from exhaustive switch: $node';
+      }
+      return removalSentinel ?? new EmptyStatement();
+    }
+
+    // Now iteratively find additional cases to keep by following targets of
+    // continue statements in kept cases.
+    final worklist = [...toKeep];
+    final collector = ContinueSwitchStatementTargetCollector(node);
+    while (worklist.isNotEmpty) {
+      final next = worklist.removeLast();
+      final targets = collector.collectTargets(next);
+      for (final target in targets) {
+        if (toKeep.add(target)) {
+          worklist.add(target);
         }
       }
     }
+
+    // Finally, remove any cases not marked for keeping. If only one case
+    // is kept, then the switch statement can be replaced with its body.
+    if (toKeep.length == 1) {
+      return toKeep.first.body;
+    }
+    node.cases.retainWhere(toKeep.contains);
+    if (foundMatchingCase && !node.hasDefault) {
+      // While the expression may not be explicitly exhaustive for the type
+      // of the tested expression, it is guaranteed to execute at least one
+      // of the remaining cases, so the backends don't need to handle the case
+      // where no listed case is hit for this switch.
+      //
+      // If the original program has the matching case directly falls through
+      // to the default case for some reason:
+      //
+      // switch (4) {
+      //    ...
+      //    case 4:
+      //    default:
+      //      ...
+      // }
+      //
+      // this means the default case is kept despite finding a guaranteed to
+      // match expression, as it contains that matching expression. If that
+      // happens, then we don't do this, to keep the invariant that
+      // isExplicitlyExhaustive is false if there is a default case.
+      node.isExplicitlyExhaustive = true;
+    }
     return node;
   }
 
   @override
-  visitStaticGet(StaticGet node, TreeNode? removalSentinel) {
+  TreeNode visitStaticGet(StaticGet node, TreeNode? removalSentinel) {
     node.transformOrRemoveChildren(this);
     final target = node.target;
     if (target is Field && target.isConst) {
       throw 'StaticGet from const field $target should be evaluated by front-end: $node';
     }
-    if (constantEvaluator.transformerShouldEvaluateExpression(node)) {
-      final context = _staticTypeContext!;
-      final result = constantEvaluator.evaluate(context, node);
-      assert(result is! UnevaluatedConstant);
-      return new ConstantExpression(result, node.getStaticType(context))
-        ..fileOffset = node.fileOffset;
+    if (!constantEvaluator.transformerShouldEvaluateExpression(node)) {
+      return node;
     }
-    return node;
+    final result = constantEvaluator.evaluate(_staticTypeContext!, node);
+    return _makeConstantExpression(result, node);
   }
 
   @override
@@ -235,3 +306,25 @@ class SimpleUnreachableCodeElimination extends RemovingTransformer {
     return node;
   }
 }
+
+class ContinueSwitchStatementTargetCollector extends RecursiveVisitor {
+  final SwitchStatement parent;
+  late Set<SwitchCase> collected;
+
+  ContinueSwitchStatementTargetCollector(this.parent);
+
+  Set<SwitchCase> collectTargets(SwitchCase node) {
+    collected = {};
+    node.accept(this);
+    return collected;
+  }
+
+  @override
+  void visitContinueSwitchStatement(ContinueSwitchStatement node) {
+    node.visitChildren(this);
+    // Only keep targets that are within the original node being checked.
+    if (node.target.parent == parent) {
+      collected.add(node.target);
+    }
+  }
+}

pkg/vm/lib/transformations/vm_constant_evaluator.dart

@@ -13,10 +13,14 @@ import 'package:front_end/src/base/nnbd_mode.dart';
 import 'package:front_end/src/fasta/kernel/constant_evaluator.dart'
     show
         AbortConstant,
+        AbortStatus,
         ConstantEvaluator,
         ErrorReporter,
         EvaluationMode,
-        SimpleErrorReporter;
+        ProceedStatus,
+        ReturnStatus,
+        SimpleErrorReporter,
+        StatementConstantEvaluator;
 
 import '../target_os.dart';
 
@@ -26,8 +30,10 @@ import '../target_os.dart';
 /// [targetOS] represents the target operating system and is used when
 /// evaluating static fields and getters annotated with "vm:platform-const".
 ///
-/// If [enableConstFunctions] is false, then only getters that return the
-/// result of a single expression can be evaluated.
+/// To avoid restricting getters annotated with "vm:platform-const" to be just
+/// a single return statement whose body is evaluated, we treat annotated
+/// getters as const functions. If [enableConstFunctions] is false, then
+/// only annotated getters are treated this way.
 class VMConstantEvaluator extends ConstantEvaluator {
   final TargetOS? _targetOS;
   final Map<String, Constant> _constantFields = {};
@@ -124,6 +130,19 @@ class VMConstantEvaluator extends ConstantEvaluator {
   bool transformerShouldEvaluateExpression(Expression node) =>
       _targetOS != null && node is StaticGet && _isPlatformConst(node.target);
 
+  Constant _executePlatformConstBody(Statement statement) {
+    final status = statement.accept(StatementConstantEvaluator(this));
+    if (status is AbortStatus) return status.error;
+    // Happens if there is no return statement in a void Function(...) body.
+    if (status is ProceedStatus) return canonicalize(NullConstant());
+    if (status is ReturnStatus) {
+      final value = status.value;
+      return value != null ? value : canonicalize(NullConstant());
+    }
+    // Other statuses denote intermediate states and not final ones.
+    throw 'No valid constant returned after executing $statement';
+  }
+
   @override
   Constant visitStaticGet(StaticGet node) {
     assert(_targetOS != null);
@@ -153,20 +172,7 @@ class VMConstantEvaluator extends ConstantEvaluator {
         result = evaluateExpressionInContext(target, target.initializer!);
       } else if (target is Procedure && target.isGetter) {
         final body = target.function.body!;
-        // If const functions are enabled, execute the getter as if it were
-        // a const function. Otherwise the annotated getter must be a single
-        // return statement whose expression is evaluated.
-        if (enableConstFunctions) {
-          result = executeBody(body);
-        } else if (body is ReturnStatement) {
-          if (body.expression == null) {
-            return canonicalize(NullConstant());
-          }
-          result = evaluateExpressionInContext(target, body.expression!);
-        } else {
-          throw "Cannot evaluate method '$nameText' since it contains more "
-              "than a single return statement.";
-        }
+        result = _executePlatformConstBody(body);
       }
       if (result is AbortConstant) {
         throw "The body or initialization of member '$nameText' does not "