Hug lambda arguments in google-java-format.

core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java

@@ -167,6 +167,7 @@
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -522,7 +523,7 @@
         if (afterFirstToken) {
           builder.forcedBreak();
         }
         builder.open(row.iterator().next().getKind() == NEW_ARRAY || cols == 1 ? ZERO : plusFour);
         boolean firstInRow = true;
         for (ExpressionTree item : row) {
           if (!firstInRow) {
@@ -762,7 +763,7 @@
         CollapseEmptyOrNot.YES,
         AllowLeadingBlankLine.YES,
         AllowTrailingBlankLine.YES);
     if (node.getStatement().getKind() == BLOCK) {
       builder.space();
     } else {
       builder.breakOp(" ");
@@ -1078,11 +1079,11 @@
     builder.open(plusFour);
     builder.open(
         node.getInitializer().size() > 1
                 && node.getInitializer().get(0).getKind() == Tree.Kind.EXPRESSION_STATEMENT
             ? plusFour
             : ZERO);
     if (!node.getInitializer().isEmpty()) {
       if (node.getInitializer().get(0).getKind() == VARIABLE) {
         PeekingIterator<StatementTree> it =
             Iterators.peekingIterator(node.getInitializer().iterator());
         visitVariables(
@@ -1146,7 +1147,7 @@
     while (true) {
       expressions.add(node.getCondition());
       statements.add(node.getThenStatement());
       if (node.getElseStatement() != null && node.getElseStatement().getKind() == IF) {
         node = (IfTree) node.getElseStatement();
       } else {
         break;
@@ -1356,6 +1357,7 @@
     builder.close();
     builder.space();
     builder.op("->");
+    markHugPoint(node);
     builder.open(statementBody ? ZERO : plusFour);
     if (statementBody) {
       builder.space();
@@ -2985,7 +2987,7 @@
     builder.close();
   }
 
   /** Collapse chains of {@code .} operators, across multiple {@link ASTNode} types. */
 
   /// Output a "." node.
   ///
@@ -3368,10 +3370,16 @@
   /// @param arguments the arguments
   /// @param plusIndent the extra indent for the arguments
   private void addArguments(List<? extends ExpressionTree> arguments, Indent plusIndent) {
-    builder.open(plusIndent);
+    boolean tabular = arguments.size() % 2 == 0 && argumentsAreTabular(arguments) == 2;
+    boolean hug = canHugOnlyArgument(arguments);
+    if (hug) {
+      builder.openHuggable(plusIndent, hugPoint(getOnlyElement(arguments)));
+    } else {
+      builder.open(plusIndent);
+    }
     token("(");
     if (!arguments.isEmpty()) {
-      if (arguments.size() % 2 == 0 && argumentsAreTabular(arguments) == 2) {
+      if (tabular) {
         builder.forcedBreak();
         builder.open(ZERO);
         boolean afterFirstToken = false;
@@ -3403,15 +3411,65 @@
         builder.close();
       } else {
         builder.breakOp();
-        argList(arguments);
+        argList(arguments, hug);
       }
     }
     token(")");
     builder.close();
   }
 
+  /**
+   * The {@link LayoutBuilder.HugPoint}s of constructs that may be hugged, keyed by the construct.
+   * They are marked when the construct is visited.
+   */
+  private final IdentityHashMap<Tree, LayoutBuilder.HugPoint> hugPoints = new IdentityHashMap<>();
+
+  /** Returns the hug point for {@code tree}, which is marked when {@code tree} is visited. */
+  private LayoutBuilder.HugPoint hugPoint(Tree tree) {
+    return hugPoints.computeIfAbsent(tree, t -> LayoutBuilder.newHugPoint());
+  }
+
+  /** Marks the hug point of {@code tree}, if it may be hugged. */
+  private void markHugPoint(Tree tree) {
+    LayoutBuilder.HugPoint hugPoint = hugPoints.remove(tree);
+    if (hugPoint != null) {
+      builder.hugPoint(hugPoint);
+    }
+  }
+
+  /**
+   * Returns true for constructs that may be hugged, i.e. laid out starting on the line of the
+   * preceding content: lambdas.
+   */
+  private static boolean isHuggable(ExpressionTree tree) {
+    return tree instanceof LambdaExpressionTree;
+  }
+
+  /**
+   * Returns true if the argument list is a single argument that may be hugged, e.g. {@code f(x -> {
+   * ... })}. Lambdas preceded by other arguments aren't hugged: the preceding arguments would share
+   * a line with the lambda's head, and wouldn't line up with the lambda's body.
+   */
+  private static boolean canHugOnlyArgument(List<? extends ExpressionTree> arguments) {
+    return arguments.size() == 1 && isHuggable(getOnlyElement(arguments));
+  }
+
   private void argList(List<? extends ExpressionTree> arguments) {
-    builder.open(ZERO);
+    argList(arguments, /* hug= */ false);
+  }
+
+  /**
+   * Adds an argument list.
+   *
+   * @param hug whether the only argument may be hugged, i.e. laid out starting on the line of the
+   *     call; see {@link LayoutBuilder#openHuggable}
+   */
+  private void argList(List<? extends ExpressionTree> arguments, boolean hug) {
+    if (hug) {
+      builder.openHuggedArguments();
+    } else {
+      builder.open(ZERO);
+    }
     boolean afterFirstToken = false;
     FillMode fillMode = hasOnlyShortItems(arguments) ? FillMode.INDEPENDENT : FillMode.UNIFIED;
     for (ExpressionTree argument : arguments) {

core/src/main/java/com/google/googlejavaformat/layout/Doc.java

@@ -53,7 +53,10 @@
 /// group. Trailing text that must stay on the same line as a group (e.g. a `;`) is instead placed
 /// inside the group by [LayoutBuilder]. These are the rules of the existing
 /// [com.google.googlejavaformat.Doc] engine, and are kept for output compatibility.
-public sealed interface Doc permits Doc.Text, Doc.Tok, Doc.Space, Doc.Break, Doc.Group {
+///
+/// A [Choice] offers a preferred alternative layout for a group, used when the group doesn't fit
+/// flat but the preferred layout's first line does.
+public sealed interface Doc permits Doc.Text, Doc.Tok, Doc.Space, Doc.Break, Doc.Group, Doc.Choice {
 
   /// Width used for documents that cannot be laid out on a single line.
   int MAX_LINE_WIDTH = 1000;
@@ -260,4 +263,41 @@ public void appendFlat(StringBuilder sb) {
       }
     }
   }
+
+  /// A choice between two layouts of the same content.
+  ///
+  /// If the content fits flat it is laid out flat, as for a [Group]. Otherwise, `fallback` is used
+  /// if taking a single direct break in `fallback` allows the rest of it to fit flat on the next
+  /// line (e.g. `f(\n x -> g(x))` rather than breaking inside the lambda). Otherwise, `preferred`
+  /// is used if, when laid out, its first line fits and contains at least its first `prefixWidth`
+  /// columns; `fallback` is used if not.
+  ///
+  /// Both alternatives must have the same flat text, so the choice never affects layouts in which
+  /// it is flat, or the fit decisions of enclosing groups. They usually share subtrees.
+  ///
+  /// @param preferred the preferred layout
+  /// @param prefixWidth the width of the part of `preferred` that must be on the current line, e.g.
+  ///   `f(x ->` for a hugged lambda argument
+  /// @param fallback the layout to use if `preferred` isn't chosen
+  record Choice(Group preferred, int prefixWidth, Group fallback) implements Doc {
+    public Choice {
+      requireNonNull(preferred, "preferred");
+      requireNonNull(fallback, "fallback");
+    }
+
+    @Override
+    public int width() {
+      return fallback.width();
+    }
+
+    @Override
+    public void appendFlat(StringBuilder sb) {
+      fallback.appendFlat(sb);
+    }
+
+    @Override
+    public Range<Integer> range() {
+      return fallback.range();
+    }
+  }
 }

core/src/main/java/com/google/googlejavaformat/layout/LayoutBuilder.java

@@ -127,8 +127,47 @@ public void drain() {
   }
 
   public void open(Indent plusIndent) {
+    open(plusIndent, HugRole.NONE, null);
+  }
+
+  private void open(Indent plusIndent, HugRole hugRole, HugPoint hugPoint) {
     depth++;
-    events.add(new Open(plusIndent));
+    events.add(new Open(plusIndent, hugRole, hugPoint));
+  }
+
+  /// The start of the body of a construct that can be hugged, i.e. laid out starting on the same
+  /// line as preceding content, e.g. the body of a lambda. See [#openHuggable].
+  public static final class HugPoint {
+    /// A zero-width marker for this point in the [Doc].
+    private final Doc marker = new Doc.Text("");
+
+    private HugPoint() {}
+  }
+
+  /// Returns a new [HugPoint], to be passed to [#hugPoint] and an enclosing level.
+  public static HugPoint newHugPoint() {
+    return new HugPoint();
+  }
+
+  /// Marks the position of `hugPoint`.
+  public void hugPoint(HugPoint hugPoint) {
+    events.add(new HugPointOp(hugPoint));
+  }
+
+  /// Opens a level for an argument list whose argument may be hugged, e.g. `f(x -> { ... })`. The
+  /// level's direct breaks, and the direct breaks of a nested level opened with
+  /// [#openHuggedArguments], are the ones that are not taken when hugging, and everything up to
+  /// `hugPoint` must then fit on the current line.
+  ///
+  /// When hugged, the hugged content is indented relative to the line containing the call, instead
+  /// of by the level's extra indent.
+  public void openHuggable(Indent plusIndent, HugPoint hugPoint) {
+    open(plusIndent, HugRole.HUGGABLE, hugPoint);
+  }
+
+  /// Opens the level holding the arguments of an enclosing [#openHuggable] level.
+  public void openHuggedArguments() {
+    open(ZERO, HugRole.ARGUMENTS, null);
   }
 
   public void close() {
@@ -493,8 +532,9 @@ private List<Event> spliceEvents(ListMultimap<Integer, Event> tokEvents) {
         continue;
       }
       result.add(ev);
-      // An `Open` doesn't output anything, so it doesn't end the run that follows a forced break.
-      if (!(ev instanceof Open)) {
+      // An `Open` or hug point doesn't output anything, so it doesn't end the run that follows a
+      // forced break.
+      if (!(ev instanceof Open || ev instanceof HugPointOp)) {
         afterForcedBreak = isForcedBreak(ev);
       }
     }
@@ -534,13 +574,14 @@ private static boolean isSpaceLike(Event ev) {
   /// closed group) into the preceding group, so it's accounted for when deciding whether that group
   /// fits on the line.
   private static Doc buildTree(List<Event> events) {
-    LevelBuilder base = new LevelBuilder(ZERO);
+    LevelBuilder base = new LevelBuilder(ZERO, HugRole.NONE, null);
     ArrayDeque<LevelBuilder> stack = new ArrayDeque<>();
     stack.addLast(base);
     LevelBuilder appendLevel = base;
     for (Event ev : events) {
       switch (ev) {
-        case Open open -> stack.addLast(new LevelBuilder(open.plusIndent()));
+        case Open open ->
+            stack.addLast(new LevelBuilder(open.plusIndent(), open.hugRole(), open.hugPoint()));
         case Close close -> {
           LevelBuilder top = stack.removeLast();
           stack.getLast().add(top);
@@ -556,6 +597,7 @@ private static Doc buildTree(List<Event> events) {
                   tok.getOriginalText(), Range.singleton(tok.getIndex()).canonical(INTEGERS)));
         }
         case TokOp tokOp -> appendLevel.add(new Doc.Tok(tokOp.tok()));
+        case HugPointOp hugPointOp -> appendLevel.add(hugPointOp.hugPoint().marker);
         case SpaceOp spaceOp -> appendLevel.add(Doc.Space.INSTANCE);
       }
     }
@@ -577,28 +619,127 @@ private static LineBreak forcedLineBreak(Indent plusIndent) {
   /// A mutable group under construction; see [#buildTree].
   private static final class LevelBuilder {
     final Indent plusIndent;
+    final HugRole hugRole;
+    final HugPoint hugPoint;
     final List<Object> items = new ArrayList<>();
 
-    LevelBuilder(Indent plusIndent) {
+    LevelBuilder(Indent plusIndent, HugRole hugRole, HugPoint hugPoint) {
       this.plusIndent = plusIndent;
+      this.hugRole = hugRole;
+      this.hugPoint = hugPoint;
     }
 
     void add(Object item) {
       items.add(item);
     }
 
-    Doc.Group build() {
+    Doc build() {
+      return hugRole == HugRole.HUGGABLE
+          ? buildHuggable()
+          : new Doc.Group(plusIndent, buildItems());
+    }
+
+    ImmutableList<Doc> buildItems() {
       ImmutableList.Builder<Doc> docs = ImmutableList.builder();
       for (Object item : items) {
         docs.add(item instanceof LevelBuilder lb ? lb.build() : (Doc) item);
       }
-      return new Doc.Group(plusIndent, docs.build());
+      return docs.build();
+    }
+
+    /// Builds a choice between hugging and the regular layout.
+    ///
+    /// The hugged layout has the same content as the regular one, but the breaks of this level and
+    /// of its [HugRole#ARGUMENTS] level (if any) are replaced by their flat text, and this level's
+    /// extra indent is dropped. Everything up to the hug point must fit on the current line.
+    private Doc buildHuggable() {
+      ImmutableList.Builder<Doc> fallback = ImmutableList.builder();
+      ImmutableList.Builder<Doc> preferred = ImmutableList.builder();
+      boolean canHug = true;
+      for (Object item : items) {
+        ImmutableList<Doc> flattened;
+        Doc fallbackDoc;
+        switch (item) {
+          case LevelBuilder lb when lb.hugRole == HugRole.ARGUMENTS -> {
+            flattened = lb.buildItems();
+            fallbackDoc = new Doc.Group(lb.plusIndent, flattened);
+          }
+          case LevelBuilder lb -> {
+            fallbackDoc = lb.build();
+            flattened = ImmutableList.of(fallbackDoc);
+          }
+          case Doc doc -> {
+            fallbackDoc = doc;
+            flattened = ImmutableList.of(doc);
+          }
+          default -> throw new AssertionError(item);
+        }
+        fallback.add(fallbackDoc);
+        for (Doc doc : flattened) {
+          if (doc instanceof Doc.Break b) {
+            canHug &= canFlatten(b);
+            preferred.add(new Doc.Text(b.flatText()));
+          } else {
+            preferred.add(doc);
+          }
+        }
+      }
+      Doc.Group fallbackGroup = new Doc.Group(plusIndent, fallback.build());
+      Doc.Group preferredGroup = new Doc.Group(ZERO, preferred.build());
+      int prefixWidth = canHug ? widthBefore(preferredGroup, hugPoint.marker) : -1;
+      if (prefixWidth < 0) {
+        return fallbackGroup;
+      }
+      return new Doc.Choice(preferredGroup, prefixWidth, fallbackGroup);
+    }
+
+    /// Forced breaks (e.g. after line comments) and tagged breaks must not be flattened.
+    private static boolean canFlatten(Doc.Break b) {
+      return b.mode() != Doc.BreakMode.FORCED && b.tag().isEmpty();
+    }
+
+    /// Returns the flat width of the content of `doc` before `marker`, or -1 if `doc` doesn't
+    /// contain `marker` within fewer than [Doc#MAX_LINE_WIDTH] columns.
+    // Hug point markers are equal to any other empty Text, so they must be found by identity.
+    @SuppressWarnings("ReferenceEquality")
+    private static int widthBefore(Doc doc, Doc marker) {
+      if (doc == marker) {
+        return 0;
+      }
+      return switch (doc) {
+        case Doc.Group g -> {
+          int width = 0;
+          for (Doc child : g.docs()) {
+            int before = widthBefore(child, marker);
+            if (before >= 0) {
+              yield min(width + before, Doc.MAX_LINE_WIDTH);
+            }
+            width += child.width();
+            if (width >= Doc.MAX_LINE_WIDTH) {
+              yield -1;
+            }
+          }
+          yield -1;
+        }
+        case Doc.Choice c -> widthBefore(c.fallback(), marker);
+        default -> -1;
+      };
     }
   }
 
-  private sealed interface Event permits Open, Close, LineBreak, TokenOp, TokOp, SpaceOp {}
+  private sealed interface Event
+      permits Open, Close, LineBreak, TokenOp, TokOp, SpaceOp, HugPointOp {}
+
+  private record Open(Indent plusIndent, HugRole hugRole, HugPoint hugPoint) implements Event {}
 
-  private record Open(Indent plusIndent) implements Event {}
+  private record HugPointOp(HugPoint hugPoint) implements Event {}
+
+  /// The role of a level in a huggable argument list; see [#openHuggable].
+  private enum HugRole {
+    NONE,
+    HUGGABLE,
+    ARGUMENTS
+  }
 
   private record Close() implements Event {}