Format null-aware elements

Author: chloestefantsova

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## 3.0.2-wip
 
+* Format null-aware elements.
+
 * Don't indent conditional branches redundantly after `=`, `:`, and `=>`.
 
   ```dart

lib/src/front_end/ast_node_visitor.dart

@@ -1251,7 +1251,8 @@ final class AstNodeVisitor extends ThrowingAstVisitor<void> with PieceFactory {
 
   @override
   void visitMapLiteralEntry(MapLiteralEntry node) {
-    writeAssignment(node.key, node.separator, node.value);
+    writeAssignment(node.key, node.separator, node.value,
+        includeLeftHandSideQuestion: true, includeRightHandSideQuestion: true);
   }
 
   @override
@@ -1353,6 +1354,11 @@ final class AstNodeVisitor extends ThrowingAstVisitor<void> with PieceFactory {
     writePostfix(node.pattern, node.operator);
   }
 
+  @override
+  void visitNullAwareElement(NullAwareElement node) {
+    writePrefix(node.question, node.value);
+  }
+
   @override
   void visitNullCheckPattern(NullCheckPattern node) {
     writePostfix(node.pattern, node.operator);

lib/src/front_end/piece_factory.dart

@@ -1338,7 +1338,10 @@ mixin PieceFactory {
       AstNode leftHandSide, Token operator, AstNode rightHandSide,
       {bool includeComma = false,
       bool canBlockSplitLeft = false,
-      NodeContext leftHandSideContext = NodeContext.none}) {
+      NodeContext leftHandSideContext = NodeContext.none,
+      bool includeLeftHandSideQuestion = false,
+      bool includeRightHandSideQuestion = false,
+      String? rightHandSidePrefix}) {
     // If an operand can have block formatting, then a newline in it doesn't
     // force the operator to split, as in:
     //
@@ -1383,15 +1386,19 @@ mixin PieceFactory {
       _ => false
     };
 
-    var leftPiece = nodePiece(leftHandSide, context: leftHandSideContext);
+    var leftPiece = nodePiece(leftHandSide,
+        questionBefore: includeLeftHandSideQuestion,
+        context: leftHandSideContext);
 
     var operatorPiece = pieces.build(() {
       if (operator.type != TokenType.COLON) pieces.space();
       pieces.token(operator);
     });
 
     var rightPiece = nodePiece(rightHandSide,
-        commaAfter: includeComma, context: NodeContext.assignment);
+        commaAfter: includeComma,
+        questionBefore: includeRightHandSideQuestion,
+        context: NodeContext.assignment);
 
     pieces.add(AssignPiece(
         left: leftPiece,
@@ -1520,7 +1527,9 @@ mixin PieceFactory {
   /// If [commaAfter] is `true`, looks for a comma token after [node] and
   /// writes it to the piece as well.
   Piece nodePiece(AstNode node,
-      {bool commaAfter = false, NodeContext context = NodeContext.none}) {
+      {bool commaAfter = false,
+      bool questionBefore = false,
+      NodeContext context = NodeContext.none}) {
     var result = pieces.build(() {
       visitNode(node, context);
     });
@@ -1533,6 +1542,14 @@ mixin PieceFactory {
       }
     }
 
+    if (questionBefore) {
+      var previousToken = node.beginToken.previous!;
+      if (previousToken.lexeme == '?') {
+        var question = tokenPiece(previousToken);
+        result = AdjacentPiece([question, result]);
+      }
+    }
+
     return result;
   }
 

test/short/whitespace/collections.stmt

@@ -92,4 +92,34 @@ f() async {
 <<<
 f() async {
   var l = [await for (x in y) x];
-}
+}
+>>> null-aware elements
+foo1(int? x) => [
+  ?x];
+<<<
+foo1(int? x) => [?x];
+>>>
+foo2(int? x) => {
+  ?x};
+<<<
+foo2(int? x) => {?x};
+>>>
+bar0() => {
+  "key": #value};
+<<<
+bar0() => {"key": #value};
+>>>
+bar1(int? x) => {
+  ?x: "value"};
+<<<
+bar1(int? x) => {?x: "value"};
+>>>
+bar2(String? y) => {
+  "key": ?y};
+<<<
+bar2(String? y) => {"key": ?y};
+>>>
+bar3(int? x, String? y) => {
+  ?x: ?y};
+<<<
+bar3(int? x, String? y) => {?x: ?y};