argument and assignment switch assists

pq · Commit Queue · commit b68719c737e4 · 2023-03-15T19:27:39.000Z
Fixes: #49960 Change-Id: Ie9865de244bdb5d4e0e2e045408bd4c18229bd0f Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/287904 Reviewed-by: Brian Wilkerson <brianwilkerson@google.com> Commit-Queue: Phil Quitslund <pquitslund@google.com>
diff --git a/pkg/analysis_server/lib/src/services/correction/dart/convert_to_switch_expression.dart b/pkg/analysis_server/lib/src/services/correction/dart/convert_to_switch_expression.dart
@@ -7,11 +7,19 @@ import 'package:analysis_server/src/services/correction/dart/abstract_producer.d
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/source/source_range.dart';
+import 'package:analyzer/src/dart/element/type_system.dart';
 import 'package:analyzer_plugin/utilities/assist/assist.dart';
 import 'package:analyzer_plugin/utilities/change_builder/change_builder_core.dart';
 import 'package:analyzer_plugin/utilities/range_factory.dart';
 
 class ConvertToSwitchExpression extends CorrectionProducer {
+  /// Local variable reference used in assignment switch expression generation.
+  LocalVariableElement? writeElement;
+
+  /// Function reference used in argument switch expression generation.
+  FunctionElement? functionElement;
+
   @override
   AssistKind get assistKind => DartAssistKind.CONVERT_TO_SWITCH_EXPRESSION;
 
@@ -21,13 +29,120 @@ class ConvertToSwitchExpression extends CorrectionProducer {
     if (node is! SwitchStatement) return;
 
     var expression = node.expression;
-    if (!expression.staticType.isExhaustive) return;
+    if (!isAlwaysExhaustive(expression.staticType)) return;
 
     if (isReturnSwitch(node)) {
       await convertReturnSwitchExpression(builder, node);
+    } else if (isAssignmentSwitch(node)) {
+      await convertAssignmentSwitchExpression(builder, node);
+    } else if (isArgumentSwitch(node)) {
+      await convertArgumentSwitchExpression(builder, node);
     }
   }
 
+  Future<void> convertArgumentSwitchExpression(
+      ChangeBuilder builder, SwitchStatement node) async {
+    await builder.addDartFileEdit(file, (builder) {
+      builder.addSimpleInsertion(node.offset, '${functionElement!.name}(');
+
+      var memberCount = node.members.length;
+      for (var i = 0; i < memberCount; ++i) {
+        // Sure to be a SwitchPatternCase
+        var patternCase = node.members[i] as SwitchPatternCase;
+        builder.addDeletion(
+            range.startStart(patternCase.keyword, patternCase.guardedPattern));
+        var colonRange = range.entity(patternCase.colon);
+        builder.addSimpleReplacement(colonRange, ' => ');
+
+        for (var statement in patternCase.statements) {
+          var hasComment = statement.beginToken.precedingComments != null;
+
+          if (statement is ExpressionStatement) {
+            var invocation = statement.expression;
+            if (invocation is MethodInvocation) {
+              var deletion = !hasComment
+                  ? range.startOffsetEndOffset(
+                      range.offsetBy(colonRange, 1).offset,
+                      invocation.argumentList.leftParenthesis.end)
+                  : range.startOffsetEndOffset(invocation.offset,
+                      invocation.argumentList.leftParenthesis.end);
+              builder.addDeletion(deletion);
+              builder.addDeletion(
+                  range.entity(invocation.argumentList.rightParenthesis));
+            }
+          }
+
+          if (!hasComment && statement.isThrowExpressionStatement) {
+            var deletionRange = range.startOffsetEndOffset(
+                range.offsetBy(colonRange, 1).offset, statement.offset);
+            builder.addDeletion(deletionRange);
+          }
+
+          if (statement is BreakStatement) {
+            var deletion = getBreakRange(statement);
+            builder.addDeletion(deletion);
+          } else {
+            var endToken = i < memberCount - 1 ? ',' : '';
+            builder.addSimpleReplacement(
+                range.entity(statement.endToken), endToken);
+          }
+        }
+      }
+
+      builder.addSimpleInsertion(node.end, ');');
+    });
+  }
+
+  Future<void> convertAssignmentSwitchExpression(
+      ChangeBuilder builder, SwitchStatement node) async {
+    await builder.addDartFileEdit(file, (builder) {
+      builder.addSimpleInsertion(node.offset, '${writeElement!.name} = ');
+
+      var memberCount = node.members.length;
+      for (var i = 0; i < memberCount; ++i) {
+        // todo(pq): extract shared replacement logic
+
+        // Sure to be a SwitchPatternCase
+        var patternCase = node.members[i] as SwitchPatternCase;
+        builder.addDeletion(
+            range.startStart(patternCase.keyword, patternCase.guardedPattern));
+        var colonRange = range.entity(patternCase.colon);
+        builder.addSimpleReplacement(colonRange, ' =>');
+
+        for (var statement in patternCase.statements) {
+          if (statement is ExpressionStatement) {
+            var expression = statement.expression;
+            if (expression is AssignmentExpression) {
+              var hasComment = statement.beginToken.precedingComments != null;
+              var deletion = !hasComment
+                  ? range.startOffsetEndOffset(
+                      range.offsetBy(colonRange, 1).offset,
+                      expression.operator.end)
+                  : range.startOffsetEndOffset(expression.beginToken.offset,
+                      expression.rightHandSide.offset);
+              builder.addDeletion(deletion);
+            } else if (expression is ThrowExpression) {
+              var deletionRange = range.startOffsetEndOffset(
+                  range.offsetBy(colonRange, 1).offset, statement.offset - 1);
+              builder.addDeletion(deletionRange);
+            }
+          }
+
+          if (statement is BreakStatement) {
+            var deletion = getBreakRange(statement);
+            builder.addDeletion(deletion);
+          } else {
+            var endToken = i < memberCount - 1 ? ',' : '';
+            builder.addSimpleReplacement(
+                range.entity(statement.endToken), endToken);
+          }
+        }
+      }
+
+      builder.addSimpleInsertion(node.end, ';');
+    });
+  }
+
   Future<void> convertReturnSwitchExpression(
       ChangeBuilder builder, SwitchStatement node) async {
     await builder.addDartFileEdit(file, (builder) {
@@ -56,13 +171,10 @@ class ConvertToSwitchExpression extends CorrectionProducer {
           builder.addDeletion(deletion);
         }
 
-        if (!hasComment && statement is ExpressionStatement) {
-          var expression = statement.expression;
-          if (expression is ThrowExpression) {
-            var deletionRange = range.startOffsetEndOffset(
-                range.offsetBy(colonRange, 1).offset, statement.offset - 1);
-            builder.addDeletion(deletionRange);
-          }
+        if (!hasComment && statement.isThrowExpressionStatement) {
+          var deletionRange = range.startOffsetEndOffset(
+              range.offsetBy(colonRange, 1).offset, statement.offset - 1);
+          builder.addDeletion(deletionRange);
         }
 
         var endToken = i < memberCount - 1 ? ',' : '';
@@ -72,6 +184,82 @@ class ConvertToSwitchExpression extends CorrectionProducer {
     });
   }
 
+  SourceRange getBreakRange(BreakStatement statement) {
+    var previous = (statement.beginToken.precedingComments ??
+        statement.beginToken.previous)!;
+    var deletion =
+        range.startOffsetEndOffset(previous.end, statement.endToken.end);
+    return deletion;
+  }
+
+  bool isAlwaysExhaustive(DartType? type) {
+    if (type == null) return false;
+    return (typeSystem as TypeSystemImpl).isAlwaysExhaustive(type);
+  }
+
+  // todo(pq): refactor the `is` checks to a single `getSwitchKind`
+  // that only looks at members once
+  // see: https://dart-review.googlesource.com/c/sdk/+/287904
+  bool isArgumentSwitch(SwitchStatement node) {
+    for (var member in node.members) {
+      if (member is! SwitchPatternCase) return false;
+      if (member.labels.isNotEmpty) return false;
+      var statements = member.statements;
+
+      if (statements.length == 1 &&
+          statements.first.isThrowExpressionStatement) {
+        continue;
+      }
+
+      if (statements.length != 2) return false;
+      if (statements[1] is! BreakStatement) return false;
+      var s = statements[0];
+      if (s is! ExpressionStatement) return false;
+      var expression = s.expression;
+      if (expression is! MethodInvocation) return false;
+      var element = expression.methodName.staticElement;
+      if (element is! FunctionElement) return false;
+      if (functionElement == null) {
+        functionElement = element;
+      } else {
+        if (functionElement != element) return false;
+      }
+    }
+
+    return functionElement != null;
+  }
+
+  bool isAssignmentSwitch(SwitchStatement node) {
+    for (var member in node.members) {
+      if (member is! SwitchPatternCase) return false;
+      if (member.labels.isNotEmpty) return false;
+      var statements = member.statements;
+
+      if (statements.length == 1 &&
+          statements.first.isThrowExpressionStatement) {
+        continue;
+      }
+
+      if (statements.length != 2) return false;
+      if (statements[1] is! BreakStatement) return false;
+      var s = statements[0];
+      if (s is! ExpressionStatement) return false;
+      var expression = s.expression;
+      if (expression is! AssignmentExpression) return false;
+      var leftHandSide = expression.leftHandSide;
+      if (leftHandSide is! SimpleIdentifier) return false;
+      if (writeElement == null) {
+        var element = leftHandSide.staticElement;
+        if (element is! LocalVariableElement) return false;
+        writeElement = element;
+      } else {
+        if (writeElement != leftHandSide.staticElement) return false;
+      }
+    }
+
+    return writeElement != null;
+  }
+
   bool isReturnSwitch(SwitchStatement node) {
     for (var member in node.members) {
       if (member is! SwitchPatternCase) return false;
@@ -88,12 +276,10 @@ class ConvertToSwitchExpression extends CorrectionProducer {
   }
 }
 
-extension on DartType? {
-  bool get isExhaustive {
-    var element = this?.element;
-    if (element is EnumElement) return true;
-    if (element is ClassElement) return element.isExhaustive;
-    if (element is MixinElement) return element.isExhaustive;
-    return false;
+extension on Statement {
+  bool get isThrowExpressionStatement {
+    var self = this;
+    if (self is! ExpressionStatement) return false;
+    return self.expression is ThrowExpression;
   }
 }
diff --git a/pkg/analysis_server/test/src/services/correction/assist/convert_to_switch_expression_test.dart b/pkg/analysis_server/test/src/services/correction/assist/convert_to_switch_expression_test.dart
@@ -19,7 +19,6 @@ class ConvertToSwitchExpressionTest extends AssistProcessorTest {
   @override
   AssistKind get kind => DartAssistKind.CONVERT_TO_SWITCH_EXPRESSION;
 
-  @FailingTest(reason: 'Not yet implemented')
   Future<void> test_argument_switchExpression() async {
     await resolveTestCode('''
 enum Color {
@@ -29,14 +28,16 @@ enum Color {
 void f(Color color) {
   switch (color) {
     case Color.red:
-      print('red');
+      print('red'); // Red.
       break;
     case Color.blue:
       print('blue');
       break;
+    // Not green.
     case Color.green:
       throw 'Green is bad';
     case Color.yellow:
+      // Yellow is OK.
       print('yellow');
       break;
   }
@@ -49,16 +50,18 @@ enum Color {
     
 void f(Color color) {
   print(switch (color) {
-    Color.red => 'red',
+    Color.red => 'red', // Red.
     Color.blue => 'blue',
+    // Not green.
     Color.green => throw 'Green is bad',
-    Color.yellow => 'yellow'
+    Color.yellow => 
+      // Yellow is OK.
+      'yellow'
   });
 }
 ''');
   }
 
-  @FailingTest(reason: 'Not yet implemented')
   Future<void> test_assignment_switchExpression() async {
     await resolveTestCode('''
 enum Color {
@@ -72,11 +75,13 @@ String f(Color color) {
       name = 'red';
       break;
     case Color.blue:
-      name = 'blue';
+      name = 'blue'; // Blue!
       break;
+    // Not green.
     case Color.green:
       throw 'Green is bad';
     case Color.yellow:
+      // Yellow is OK.
       name = 'yellow';
       break;
   }
@@ -92,9 +97,12 @@ String f(Color color) {
   var name = '';
   name = switch (color) {
     Color.red => 'red',
-    Color.blue => 'blue',
+    Color.blue => 'blue', // Blue!
+    // Not green.
     Color.green => throw 'Green is bad',
-    Color.yellow => 'yellow'
+    Color.yellow =>
+      // Yellow is OK.
+      'yellow'
   };
   return name;
 }
