[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]>

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,89 @@ 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;
+
+    // Set of SwitchCases that should be retained, either because they are
+    // guaranteed to match, may match and no case before them is guaranteed
+    // to match, or because they are the target of a case that may match or
+    // is guaranteed to match.
+    final toKeep = <SwitchCase>{};
+    // Whether an expression from a previous case is guaranteed to match.
+    bool foundMatchingCase = false;
+    for (final c in node.cases) {
+      // Trim any constant expressions that don't match, or any expression
+      // that follows a guaranteed-to-match case. Perform this trimming even
+      // on cases that follow a guaranteed to match case and thus are not
+      // in the initial set to retain, since they may be the target of continue
+      // statements and thus added to the set to retain later.
+      bool containsMatchingCase = false;
+      c.expressions.retainWhere((e) {
+        if (foundMatchingCase) return false;
+        if (e is! ConstantExpression) return true;
+        containsMatchingCase = e.constant == tested.constant;
+        return containsMatchingCase;
+      });
+      if (foundMatchingCase) continue;
+      foundMatchingCase = containsMatchingCase;
+      if (c.isDefault || c.expressions.isNotEmpty) {
+        toKeep.add(c);
+      }
+    }
+
+    if (toKeep.isEmpty) {
+      throw 'Expected at least one possibly matching case: $node';
+    }
+
+    // Now iteratively keep cases that are targets of continue statements
+    // within the cases we've already marked to keep.
+    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);
         }
       }
     }
+
+    node.cases.retainWhere((c) => toKeep.contains(c));
+    if (node.cases.length == 1) {
+      return node.cases.first.body;
+    }
     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 +284,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 "