Don't reimplement comment threads for attention set

As we extracted the logic for comment threads with Ie2713ebdb, we can
use it for attention set. There was also a corner case which wasn't
covered by the logic which was previously used. We add a test case for
it to avoid future regressions.

In AttentionSetIT, we also switched to a different approach for looking
up the ChangeData instance. That was necessary due to the different
types used in the new test but also had a different reason: There's
no guarantee that the ChangeData object returned by PushOneCommit.Result
is updated after the push happened. We skipped additional improvements
like looking up attention set updates via a cleaner/better way than a
ChangeData instance as that's much beyond the scope of this change.

Change-Id: I9d8091c3ad6575304de9a5698cb1c94e32421e7b

java/com/google/gerrit/server/CommentsUtil.java

@@ -21,7 +21,6 @@ import static java.util.stream.Collectors.toCollection;
 import static java.util.stream.Collectors.toList;
 
 import com.google.common.collect.ComparisonChain;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
 import com.google.gerrit.common.Nullable;
@@ -51,12 +50,8 @@ import java.io.IOException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
-import java.util.stream.Collectors;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.Repository;
@@ -354,43 +349,6 @@ public class CommentsUtil {
     update.deleteCommentByRewritingHistory(commentKey.uuid, newMessage);
   }
 
-  /**
-   * Gets all of the {@link HumanComment} in the comment threads that received a reply.
-   *
-   * @param changeNotes notes of this change.
-   * @param newComments set of all the new comments added on the change by the current user.
-   * @return set of all comments in the comments thread that received a reply.
-   */
-  public Set<HumanComment> getAllHumanCommentsInCommentThreads(
-      ChangeNotes changeNotes, ImmutableSet<HumanComment> newComments) {
-    Map<String, HumanComment> uuidToComment =
-        publishedHumanCommentsByChange(changeNotes).stream()
-            .collect(Collectors.toMap(c -> c.key.uuid, c -> c));
-
-    // Copy the set so that it won't be mutated.
-    List<HumanComment> toTraverse = new ArrayList<>(newComments);
-    Set<String> seen = new HashSet<>();
-    Set<HumanComment> allCommentsInCommentThreads = new HashSet<>();
-    while (!toTraverse.isEmpty()) {
-      HumanComment current = toTraverse.remove(0);
-      allCommentsInCommentThreads.add(current);
-
-      if (current.parentUuid != null) {
-        HumanComment parent = uuidToComment.get(current.parentUuid);
-        if (parent == null) {
-          // If we can't find the parent within the human comments, the parent must be a robot
-          // comment and can be ignored.
-          continue;
-        }
-        if (!seen.contains(current.parentUuid)) {
-          toTraverse.add(parent);
-          seen.add(current.parentUuid);
-        }
-      }
-    }
-    return allCommentsInCommentThreads;
-  }
-
   private static List<HumanComment> commentsOnFile(
       Collection<HumanComment> allComments, String file) {
     List<HumanComment> result = new ArrayList<>(allComments.size());

java/com/google/gerrit/server/change/CommentThreads.java

@@ -17,6 +17,7 @@ package com.google.gerrit.server.change;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.util.stream.Collectors.groupingBy;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.gerrit.entities.Comment;
@@ -24,6 +25,7 @@ import java.util.Comparator;
 import java.util.Map;
 import java.util.PriorityQueue;
 import java.util.Queue;
+import java.util.function.Function;
 
 /**
  * Identifier of comment threads.
@@ -40,31 +42,74 @@ import java.util.Queue;
  */
 public class CommentThreads<T extends Comment> {
 
-  private final Iterable<T> comments;
+  private final ImmutableMap<String, T> commentPerUuid;
+  private final Map<String, ImmutableSet<T>> childrenPerParent;
 
-  private CommentThreads(Iterable<T> comments) {
-    this.comments = comments;
+  public CommentThreads(
+      ImmutableMap<String, T> commentPerUuid, Map<String, ImmutableSet<T>> childrenPerParent) {
+    this.commentPerUuid = commentPerUuid;
+    this.childrenPerParent = childrenPerParent;
   }
 
   public static <T extends Comment> CommentThreads<T> forComments(Iterable<T> comments) {
-    return new CommentThreads<>(comments);
-  }
-
-  public ImmutableSet<CommentThread<T>> getThreads() {
-    ImmutableSet<String> commentUuids =
-        Streams.stream(comments).map(comment -> comment.key.uuid).collect(toImmutableSet());
-
-    ImmutableSet<T> roots =
+    ImmutableMap<String, T> commentPerUuid =
         Streams.stream(comments)
-            .filter(
-                comment -> comment.parentUuid == null || !commentUuids.contains(comment.parentUuid))
-            .collect(toImmutableSet());
+            .distinct()
+            .collect(ImmutableMap.toImmutableMap(comment -> comment.key.uuid, Function.identity()));
 
     Map<String, ImmutableSet<T>> childrenPerParent =
-        Streams.stream(comments)
+        commentPerUuid.values().stream()
             .filter(comment -> comment.parentUuid != null)
             .collect(groupingBy(comment -> comment.parentUuid, toImmutableSet()));
+    return new CommentThreads<>(commentPerUuid, childrenPerParent);
+  }
 
+  /**
+   * Returns all comments organized into threads.
+   *
+   * <p>Comments appear only once.
+   */
+  public ImmutableSet<CommentThread<T>> getThreads() {
+    ImmutableSet<T> roots =
+        commentPerUuid.values().stream().filter(this::isRoot).collect(toImmutableSet());
+
+    return buildThreadsOf(roots);
+  }
+
+  /**
+   * Returns only the comment threads to which the specified comments are a reply.
+   *
+   * <p>If the specified child comments are part of the comments originally provided to {@link
+   * CommentThreads#forComments(Iterable)}, they will also appear in the returned comment threads.
+   * They don't need to be part of the originally provided comments, though, but should refer to one
+   * of these comments via their {@link Comment#parentUuid}. Child comments not referring to any
+   * known comments will be ignored.
+   *
+   * @param childComments comments for which the matching threads should be determined
+   * @return threads to which the provided child comments are a reply
+   */
+  public ImmutableSet<CommentThread<T>> getThreadsForChildren(Iterable<? extends T> childComments) {
+    ImmutableSet<T> relevantRoots =
+        Streams.stream(childComments)
+            .map(this::findRoot)
+            .filter(root -> commentPerUuid.containsKey(root.key.uuid))
+            .collect(toImmutableSet());
+    return buildThreadsOf(relevantRoots);
+  }
+
+  private T findRoot(T comment) {
+    T current = comment;
+    while (!isRoot(current)) {
+      current = commentPerUuid.get(current.parentUuid);
+    }
+    return current;
+  }
+
+  private boolean isRoot(T current) {
+    return current.parentUuid == null || !commentPerUuid.containsKey(current.parentUuid);
+  }
+
+  private ImmutableSet<CommentThread<T>> buildThreadsOf(ImmutableSet<T> roots) {
     return roots.stream()
         .map(root -> buildCommentThread(root, childrenPerParent))
         .collect(toImmutableSet());

java/com/google/gerrit/server/restapi/change/ReplyAttentionSetUpdates.java

@@ -17,6 +17,8 @@ package com.google.gerrit.server.restapi.change;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
+import com.google.common.collect.Sets.SetView;
 import com.google.gerrit.entities.Account;
 import com.google.gerrit.entities.HumanComment;
 import com.google.gerrit.entities.PatchSet;
@@ -32,6 +34,8 @@ import com.google.gerrit.server.account.AccountResolver;
 import com.google.gerrit.server.account.ServiceUserClassifier;
 import com.google.gerrit.server.change.AddToAttentionSetOp;
 import com.google.gerrit.server.change.AttentionSetUnchangedOp;
+import com.google.gerrit.server.change.CommentThread;
+import com.google.gerrit.server.change.CommentThreads;
 import com.google.gerrit.server.change.RemoveFromAttentionSetOp;
 import com.google.gerrit.server.notedb.ChangeNotes;
 import com.google.gerrit.server.notedb.ChangeUpdate;
@@ -44,6 +48,7 @@ import com.google.gerrit.server.util.time.TimeUtil;
 import com.google.inject.Inject;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -208,19 +213,22 @@ public class ReplyAttentionSetUpdates {
   /** Adds all authors of all comment threads that received a reply during this update */
   private void addAllAuthorsOfCommentThreads(
       BatchUpdate bu, ChangeNotes changeNotes, ImmutableSet<HumanComment> allNewComments) {
-    Set<HumanComment> allCommentsInCommentThreads =
-        commentsUtil.getAllHumanCommentsInCommentThreads(changeNotes, allNewComments);
-    // Copy the set to make it mutable, so that we can delete users that were already added.
-    Set<Account.Id> possibleUsersToAdd =
-        new HashSet<>(approvalsUtil.getReviewers(changeNotes).all());
+    List<HumanComment> publishedComments = commentsUtil.publishedHumanCommentsByChange(changeNotes);
+    ImmutableSet<CommentThread<HumanComment>> repliedToCommentThreads =
+        CommentThreads.forComments(publishedComments).getThreadsForChildren(allNewComments);
 
-    for (HumanComment comment : allCommentsInCommentThreads) {
-      Account.Id author = comment.author.getId();
-      if (possibleUsersToAdd.contains(author)) {
-        addToAttentionSet(
-            bu, changeNotes, author, "Someone else replied on a comment you posted", false);
-        possibleUsersToAdd.remove(author);
-      }
+    ImmutableSet<Account.Id> repliedToUsers =
+        repliedToCommentThreads.stream()
+            .map(CommentThread::comments)
+            .flatMap(Collection::stream)
+            .map(comment -> comment.author.getId())
+            .collect(toImmutableSet());
+    ImmutableSet<Account.Id> possibleUsersToAdd = approvalsUtil.getReviewers(changeNotes).all();
+    SetView<Account.Id> usersToAdd = Sets.intersection(possibleUsersToAdd, repliedToUsers);
+
+    for (Account.Id user : usersToAdd) {
+      addToAttentionSet(
+          bu, changeNotes, user, "Someone else replied on a comment you posted", false);
     }
   }