notedb: Rewrite comments to be keyed by RevId

Rather than maintain two separate maps, one for each side, use a
single multimap of SHA-1 (RevId) to PatchLineComment. Callers are
responsible for looking up the RevId they need (e.g. with the
PatchListCache), but in practice they rarely do.

Change-Id: I5597cc6c12a881e7cf7f3c786b9697b4f7b9749f

gerrit-server/src/main/java/com/google/gerrit/server/PatchLineCommentsUtil.java

@@ -14,6 +14,8 @@
 
 package com.google.gerrit.server;
 
+import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
@@ -105,8 +107,7 @@ public class PatchLineCommentsUtil {
 
     notes.load();
     List<PatchLineComment> comments = Lists.newArrayList();
-    comments.addAll(notes.getBaseComments().values());
+    comments.addAll(notes.getComments().values());
-    comments.addAll(notes.getPatchSetComments().values());
     return sort(comments);
   }
 
@@ -165,13 +166,7 @@ public class PatchLineCommentsUtil {
       return sort(
           db.patchComments().publishedByChangeFile(changeId, file).toList());
     }
-    notes.load();
+    return commentsOnFile(notes.load().getComments().values(), file);
-    List<PatchLineComment> comments = Lists.newArrayList();
-
-    addCommentsOnFile(comments, notes.getBaseComments().values(), file);
-    addCommentsOnFile(comments, notes.getPatchSetComments().values(),
-        file);
-    return sort(comments);
   }
 
   public List<PatchLineComment> publishedByPatchSet(ReviewDb db,
@@ -180,11 +175,7 @@ public class PatchLineCommentsUtil {
       return sort(
           db.patchComments().publishedByPatchSet(psId).toList());
     }
-    notes.load();
+    return commentsOnPatchSet(notes.load().getComments().values(), psId);
-    List<PatchLineComment> comments = new ArrayList<>();
-    comments.addAll(notes.getPatchSetComments().get(psId));
-    comments.addAll(notes.getBaseComments().get(psId));
-    return sort(comments);
   }
 
   public List<PatchLineComment> draftByPatchSetAuthor(ReviewDb db,
@@ -194,11 +185,8 @@ public class PatchLineCommentsUtil {
       return sort(
           db.patchComments().draftByPatchSetAuthor(psId, author).toList());
     }
-
+    return commentsOnPatchSet(
-    List<PatchLineComment> comments = Lists.newArrayList();
+        notes.load().getDraftComments(author).values(), psId);
-    comments.addAll(notes.getDraftBaseComments(author).row(psId).values());
-    comments.addAll(notes.getDraftPsComments(author).row(psId).values());
-    return sort(comments);
   }
 
   public List<PatchLineComment> draftByChangeFileAuthor(ReviewDb db,
@@ -210,12 +198,8 @@ public class PatchLineCommentsUtil {
             .draftByChangeFileAuthor(notes.getChangeId(), file, author)
             .toList());
     }
-    List<PatchLineComment> comments = Lists.newArrayList();
+    return commentsOnFile(
-    addCommentsOnFile(comments, notes.getDraftBaseComments(author).values(),
+        notes.load().getDraftComments(author).values(), file);
-        file);
-    addCommentsOnFile(comments, notes.getDraftPsComments(author).values(),
-        file);
-    return sort(comments);
   }
 
   public List<PatchLineComment> draftByChangeAuthor(ReviewDb db,
@@ -225,8 +209,7 @@ public class PatchLineCommentsUtil {
       return sort(db.patchComments().byChange(notes.getChangeId()).toList());
     }
     List<PatchLineComment> comments = Lists.newArrayList();
-    comments.addAll(notes.getDraftBaseComments(author).values());
+    comments.addAll(notes.getDraftComments(author).values());
-    comments.addAll(notes.getDraftPsComments(author).values());
     return sort(comments);
   }
 
@@ -236,9 +219,8 @@ public class PatchLineCommentsUtil {
       return sort(db.patchComments().draftByAuthor(author).toList());
     }
 
-    Set<String> refNames =
+    // TODO(dborowitz): Just scan author space.
-        getRefNamesAllUsers(RefNames.REFS_DRAFT_COMMENTS);
+    Set<String> refNames = getRefNamesAllUsers(RefNames.REFS_DRAFT_COMMENTS);
-
     List<PatchLineComment> comments = Lists.newArrayList();
     for (String refName : refNames) {
       Account.Id id = Account.Id.fromRefPart(refName);
@@ -246,10 +228,8 @@ public class PatchLineCommentsUtil {
         continue;
       }
       Change.Id changeId = Change.Id.parse(refName);
-      DraftCommentNotes draftNotes =
+      comments.addAll(
-          draftFactory.create(changeId, author).load();
+          draftFactory.create(changeId, author).load().getComments().values());
-      comments.addAll(draftNotes.getDraftBaseComments().values());
-      comments.addAll(draftNotes.getDraftPsComments().values());
     }
     return sort(comments);
   }
@@ -286,33 +266,45 @@ public class PatchLineCommentsUtil {
     db.patchComments().delete(comments);
   }
 
-  private static Collection<PatchLineComment> addCommentsOnFile(
+  private static List<PatchLineComment> commentsOnFile(
-      Collection<PatchLineComment> commentsOnFile,
       Collection<PatchLineComment> allComments,
       String file) {
+    List<PatchLineComment> result = new ArrayList<>(allComments.size());
     for (PatchLineComment c : allComments) {
       String currentFilename = c.getKey().getParentKey().getFileName();
       if (currentFilename.equals(file)) {
-        commentsOnFile.add(c);
+        result.add(c);
       }
     }
-    return commentsOnFile;
+    return sort(result);
   }
 
-  public static void setCommentRevId(PatchLineComment c,
+  private static List<PatchLineComment> commentsOnPatchSet(
+      Collection<PatchLineComment> allComments,
+      PatchSet.Id psId) {
+    List<PatchLineComment> result = new ArrayList<>(allComments.size());
+    for (PatchLineComment c : allComments) {
+      if (getCommentPsId(c).equals(psId)) {
+        result.add(c);
+      }
+    }
+    return sort(result);
+  }
+
+  public static RevId setCommentRevId(PatchLineComment c,
       PatchListCache cache, Change change, PatchSet ps) throws OrmException {
-    if (c.getRevId() != null) {
+    if (c.getRevId() == null) {
-      return;
+      try {
+        // TODO(dborowitz): Bypass cache if side is REVISION.
+        PatchList patchList = cache.get(change, ps);
+        c.setRevId((c.getSide() == (short) 0)
+          ? new RevId(ObjectId.toString(patchList.getOldId()))
+          : new RevId(ObjectId.toString(patchList.getNewId())));
+      } catch (PatchListNotAvailableException e) {
+        throw new OrmException(e);
+      }
     }
-    PatchList patchList;
+    return c.getRevId();
-    try {
-      patchList = cache.get(change, ps);
-    } catch (PatchListNotAvailableException e) {
-      throw new OrmException(e);
-    }
-    c.setRevId((c.getSide() == (short) 0)
-      ? new RevId(ObjectId.toString(patchList.getOldId()))
-      : new RevId(ObjectId.toString(patchList.getNewId())));
   }
 
   private Set<String> getRefNamesAllUsers(String prefix) throws OrmException {

gerrit-server/src/main/java/com/google/gerrit/server/change/PostReview.java

@@ -370,8 +370,7 @@ public class PostReview implements RestModifyView<RevisionResource, ReviewInput>
         e.setStatus(PatchLineComment.Status.PUBLISHED);
         e.setWrittenOn(timestamp);
         e.setSide(c.side == Side.PARENT ? (short) 0 : (short) 1);
-        setCommentRevId(e, patchListCache, rsrc.getChange(),
+        setCommentRevId(e, patchListCache, rsrc.getChange(), rsrc.getPatchSet());
-            rsrc.getPatchSet());
         e.setMessage(c.message);
         if (c.range != null) {
           e.setRange(new CommentRange(
@@ -396,8 +395,7 @@ public class PostReview implements RestModifyView<RevisionResource, ReviewInput>
         for (PatchLineComment e : drafts.values()) {
           e.setStatus(PatchLineComment.Status.PUBLISHED);
           e.setWrittenOn(timestamp);
-          setCommentRevId(e, patchListCache, rsrc.getChange(),
+          setCommentRevId(e, patchListCache, rsrc.getChange(), rsrc.getPatchSet());
-              rsrc.getPatchSet());
           ups.add(e);
         }
         break;

gerrit-server/src/main/java/com/google/gerrit/server/change/PutDraftComment.java

@@ -96,8 +96,7 @@ public class PutDraftComment implements RestModifyView<DraftCommentResource, Dra
           Collections.singleton(update(c, in)));
     } else {
       if (c.getRevId() == null) {
-        setCommentRevId(c, patchListCache, rsrc.getChange(),
+        setCommentRevId(c, patchListCache, rsrc.getChange(), rsrc.getPatchSet());
-            rsrc.getPatchSet());
       }
       plcUtil.updateComments(db.get(), update,
           Collections.singleton(update(c, in)));

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeDraftUpdate.java

@@ -16,14 +16,15 @@ package com.google.gerrit.server.notedb;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.gerrit.server.notedb.CommentsInNotesUtil.addCommentToMap;
 import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
 
+import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Table;
+import com.google.common.collect.Sets;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchLineComment.Status;
-import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.Project;
 import com.google.gerrit.reviewdb.client.RefNames;
 import com.google.gerrit.reviewdb.client.RevId;
@@ -46,8 +47,12 @@ import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.RevCommit;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -190,76 +195,55 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
       noteMap = NoteMap.newEmptyMap();
     }
 
-    Table<PatchSet.Id, String, PatchLineComment> baseDrafts =
+    Map<RevId, List<PatchLineComment>> allComments = new HashMap<>();
-        draftNotes.getDraftBaseComments();
-    Table<PatchSet.Id, String, PatchLineComment> psDrafts =
-        draftNotes.getDraftPsComments();
-
-    boolean draftsEmpty = baseDrafts.isEmpty() && psDrafts.isEmpty();
-
-    // There is no need to rewrite the note for one of the sides of the patch
-    // set if all of the modifications were made to the comments of one side,
-    // so we set these flags to potentially save that extra work.
-    boolean baseSideChanged = false;
-    boolean revisionSideChanged = false;
-
-    // We must define these RevIds so that if this update deletes all
-    // remaining comments on a given side, then we can remove that note.
-    // However, if this update doesn't delete any comments, it is okay for these
-    // to be null because they won't be used.
-    RevId baseRevId = null;
-    RevId psRevId = null;
 
+    boolean hasComments = false;
+    int n = deleteComments.size() + upsertComments.size();
+    Set<RevId> updatedRevs = Sets.newHashSetWithExpectedSize(n);
+    Set<PatchLineComment.Key> updatedKeys = Sets.newHashSetWithExpectedSize(n);
     for (PatchLineComment c : deleteComments) {
-      if (c.getSide() == (short) 0) {
+      allComments.put(c.getRevId(), new ArrayList<PatchLineComment>());
-        baseSideChanged = true;
+      updatedRevs.add(c.getRevId());
-        baseRevId = c.getRevId();
+      updatedKeys.add(c.getKey());
-        baseDrafts.remove(psId, c.getKey().get());
-      } else {
-        revisionSideChanged = true;
-        psRevId = c.getRevId();
-        psDrafts.remove(psId, c.getKey().get());
-      }
     }
 
     for (PatchLineComment c : upsertComments) {
-      if (c.getSide() == (short) 0) {
+      hasComments = true;
-        baseSideChanged = true;
+      addCommentToMap(allComments, c);
-        baseDrafts.put(psId, c.getKey().get(), c);
+      updatedRevs.add(c.getRevId());
-      } else {
+      updatedKeys.add(c.getKey());
-        revisionSideChanged = true;
+    }
-        psDrafts.put(psId, c.getKey().get(), c);
+
+    // Re-add old comments for updated revisions so the new note contents
+    // includes both old and new comments merged in the right order.
+    //
+    // writeCommentsToNoteMap doesn't touch notes for SHA-1s that are not
+    // mentioned in the input map, so by omitting comments for those revisions,
+    // we avoid the work of having to re-serialize identical comment data for
+    // those revisions.
+    ListMultimap<RevId, PatchLineComment> existing =
+        draftNotes.getComments();
+    for (Map.Entry<RevId, PatchLineComment> e : existing.entries()) {
+      PatchLineComment c = e.getValue();
+      if (updatedRevs.contains(c.getRevId())
+          && !updatedKeys.contains(c.getKey())) {
+        hasComments = true;
+        addCommentToMap(allComments, e.getValue());
       }
     }
 
-    List<PatchLineComment> newBaseDrafts =
+    // If we touched every revision and there are no comments left, set the flag
-        Lists.newArrayList(baseDrafts.row(psId).values());
+    // for the caller to delete the entire ref.
-    List<PatchLineComment> newPsDrafts =
+    boolean touchedAllRevs = updatedRevs.equals(existing.keySet());
-        Lists.newArrayList(psDrafts.row(psId).values());
+    if (touchedAllRevs && !hasComments) {
-
+      removedAllComments.set(touchedAllRevs && !hasComments);
-    updateNoteMap(baseSideChanged, noteMap, newBaseDrafts,
+      return null;
-        baseRevId);
+    }
-    updateNoteMap(revisionSideChanged, noteMap, newPsDrafts,
-        psRevId);
-
-    removedAllComments.set(
-        baseDrafts.isEmpty() && psDrafts.isEmpty() && !draftsEmpty);
 
+    commentsUtil.writeCommentsToNoteMap(noteMap, allComments, inserter);
     return noteMap.writeTree(inserter);
   }
 
-  private void updateNoteMap(boolean changed, NoteMap noteMap,
-      List<PatchLineComment> comments, RevId commitId)
-      throws IOException {
-    if (changed) {
-      if (comments.isEmpty()) {
-        commentsUtil.removeNote(noteMap, commitId);
-      } else {
-        commentsUtil.writeCommentsToNoteMap(noteMap, comments, inserter);
-      }
-    }
-  }
-
   public RevCommit commit() throws IOException {
     BatchMetaDataUpdate batch = openUpdate();
     try {
@@ -279,13 +263,11 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
     if (migration.writeChanges()) {
       AtomicBoolean removedAllComments = new AtomicBoolean();
       ObjectId treeId = storeCommentsInNotes(removedAllComments);
-      if (treeId != null) {
+      if (removedAllComments.get()) {
-        if (removedAllComments.get()) {
+        batch.removeRef(getRefName());
-          batch.removeRef(getRefName());
+      } else if (treeId != null) {
-        } else {
+        builder.setTreeId(treeId);
-          builder.setTreeId(treeId);
+        batch.write(builder);
-          batch.write(builder);
-        }
       }
     }
   }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeNotes.java

@@ -15,7 +15,6 @@
 package com.google.gerrit.server.notedb;
 
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.GERRIT_PLACEHOLDER_HOST;
-import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
@@ -26,7 +25,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Ordering;
-import com.google.common.collect.Table;
 import com.google.common.primitives.Ints;
 import com.google.gerrit.common.data.SubmitRecord;
 import com.google.gerrit.reviewdb.client.Account;
@@ -36,6 +34,7 @@ import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.PatchSetApproval;
 import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.config.AllUsersName;
 import com.google.gerrit.server.config.AllUsersNameProvider;
 import com.google.gerrit.server.git.GitRepositoryManager;
@@ -53,7 +52,6 @@ import org.eclipse.jgit.revwalk.RevWalk;
 import java.io.IOException;
 import java.sql.Timestamp;
 import java.util.Comparator;
-import java.util.List;
 import java.util.Map;
 
 /** View of a single {@link Change} based on the log of its notes branch. */
@@ -138,8 +136,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   private ImmutableList<Account.Id> allPastReviewers;
   private ImmutableList<SubmitRecord> submitRecords;
   private ImmutableListMultimap<PatchSet.Id, ChangeMessage> changeMessages;
-  private ImmutableListMultimap<PatchSet.Id, PatchLineComment> commentsForBase;
+  private ImmutableListMultimap<RevId, PatchLineComment> comments;
-  private ImmutableListMultimap<PatchSet.Id, PatchLineComment> commentsForPS;
   private ImmutableSet<String> hashtags;
   NoteMap noteMap;
 
@@ -194,28 +191,15 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
     return changeMessages;
   }
 
-  /** @return inline comments on each patchset's base (side == 0). */
+  /** @return inline comments on each revision. */
-  public ImmutableListMultimap<PatchSet.Id, PatchLineComment>
+  public ImmutableListMultimap<RevId, PatchLineComment> getComments() {
-      getBaseComments() {
+    return comments;
-    return commentsForBase;
   }
 
-  /** @return inline comments on each patchset (side == 1). */
+  public ImmutableListMultimap<RevId, PatchLineComment> getDraftComments(
-  public ImmutableListMultimap<PatchSet.Id, PatchLineComment>
-      getPatchSetComments() {
-    return commentsForPS;
-  }
-
-  public Table<PatchSet.Id, String, PatchLineComment> getDraftBaseComments(
       Account.Id author) throws OrmException {
     loadDraftComments(author);
-    return draftCommentNotes.getDraftBaseComments();
+    return draftCommentNotes.getComments();
-  }
-
-  public Table<PatchSet.Id, String, PatchLineComment> getDraftPsComments(
-      Account.Id author) throws OrmException {
-    loadDraftComments(author);
-    return draftCommentNotes.getDraftPsComments();
   }
 
   /**
@@ -234,6 +218,11 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
     }
   }
 
+  @VisibleForTesting
+  DraftCommentNotes getDraftCommentNotes() {
+    return draftCommentNotes;
+  }
+
   public boolean containsComment(PatchLineComment c) throws OrmException {
     if (containsCommentPublished(c)) {
       return true;
@@ -243,11 +232,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   }
 
   public boolean containsCommentPublished(PatchLineComment c) {
-    PatchSet.Id psId = getCommentPsId(c);
+    for (PatchLineComment l : getComments().values()) {
-    List<PatchLineComment> list = (c.getSide() == (short) 0)
-        ? getBaseComments().get(psId)
-        : getPatchSetComments().get(psId);
-    for (PatchLineComment l : list) {
       if (c.getKey().equals(l.getKey())) {
         return true;
       }
@@ -282,8 +267,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
       }
       approvals = parser.buildApprovals();
       changeMessages = parser.buildMessages();
-      commentsForBase = ImmutableListMultimap.copyOf(parser.commentsForBase);
+      comments = ImmutableListMultimap.copyOf(parser.comments);
-      commentsForPS = ImmutableListMultimap.copyOf(parser.commentsForPs);
       noteMap = parser.commentNoteMap;
 
       if (parser.hashtags != null) {
@@ -310,8 +294,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
     reviewers = ImmutableSetMultimap.of();
     submitRecords = ImmutableList.of();
     changeMessages = ImmutableListMultimap.of();
-    commentsForBase = ImmutableListMultimap.of();
+    comments = ImmutableListMultimap.of();
-    commentsForPS = ImmutableListMultimap.of();
     hashtags = ImmutableSet.of();
   }
 

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeNotesParser.java

@@ -44,6 +44,7 @@ import com.google.gerrit.reviewdb.client.LabelId;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.PatchSetApproval;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.util.LabelVote;
 
@@ -72,8 +73,7 @@ class ChangeNotesParser implements AutoCloseable {
   final Map<Account.Id, ReviewerState> reviewers;
   final List<Account.Id> allPastReviewers;
   final List<SubmitRecord> submitRecords;
-  final Multimap<PatchSet.Id, PatchLineComment> commentsForPs;
+  final Multimap<RevId, PatchLineComment> comments;
-  final Multimap<PatchSet.Id, PatchLineComment> commentsForBase;
   NoteMap commentNoteMap;
   Change.Status status;
   Set<String> hashtags;
@@ -99,8 +99,7 @@ class ChangeNotesParser implements AutoCloseable {
     allPastReviewers = Lists.newArrayList();
     submitRecords = Lists.newArrayListWithExpectedSize(1);
     changeMessages = LinkedListMultimap.create();
-    commentsForPs = ArrayListMultimap.create();
+    comments = ArrayListMultimap.create();
-    commentsForBase = ArrayListMultimap.create();
   }
 
   @Override
@@ -275,7 +274,7 @@ class ChangeNotesParser implements AutoCloseable {
       throws IOException, ConfigInvalidException {
     commentNoteMap = CommentsInNotesUtil.parseCommentsFromNotes(repo,
         ChangeNoteUtil.changeRefName(changeId), walk, changeId,
-        commentsForBase, commentsForPs, PatchLineComment.Status.PUBLISHED);
+        comments, PatchLineComment.Status.PUBLISHED);
   }
 
   private void parseApproval(PatchSet.Id psId, Account.Id accountId,

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeUpdate.java

@@ -20,6 +20,7 @@ import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_LABEL;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_PATCH_SET;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_STATUS;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_SUBMITTED_WITH;
+import static com.google.gerrit.server.notedb.CommentsInNotesUtil.addCommentToMap;
 import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -28,14 +29,13 @@ import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
-import com.google.common.collect.Multimap;
 import com.google.gerrit.common.data.SubmitRecord;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchLineComment.Status;
-import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.GerritPersonIdent;
 import com.google.gerrit.server.account.AccountCache;
 import com.google.gerrit.server.config.AnonymousCowardName;
@@ -57,6 +57,7 @@ import org.eclipse.jgit.revwalk.RevCommit;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.Date;
 import java.util.List;
@@ -91,8 +92,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   private String subject;
   private List<SubmitRecord> submitRecords;
   private final CommentsInNotesUtil commentsUtil;
-  private List<PatchLineComment> commentsForBase;
+  private List<PatchLineComment> comments;
-  private List<PatchLineComment> commentsForPs;
   private Set<String> hashtags;
   private String changeMessage;
   private ChangeNotes notes;
@@ -161,8 +161,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     this.commentsUtil = commentsUtil;
     this.approvals = Maps.newTreeMap(labelNameComparator);
     this.reviewers = Maps.newLinkedHashMap();
-    this.commentsForPs = Lists.newArrayList();
+    this.comments = Lists.newArrayList();
-    this.commentsForBase = Lists.newArrayList();
   }
 
   public void setStatus(Change.Status status) {
@@ -238,11 +237,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
           "A comment already exists with the same key as the following comment,"
           + " so we cannot insert this comment: %s", c);
     }
-    if (c.getSide() == 0) {
+    comments.add(c);
-      commentsForBase.add(c);
-    } else {
-      commentsForPs.add(c);
-    }
   }
 
   private void insertDraftComment(PatchLineComment c) throws OrmException {
@@ -263,11 +258,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
       checkArgument(!notes.containsCommentPublished(c),
           "Cannot update a comment that has already been published and saved");
     }
-    if (c.getSide() == 0) {
+    comments.add(c);
-      commentsForBase.add(c);
-    } else {
-      commentsForPs.add(c);
-    }
   }
 
   private void upsertDraftComment(PatchLineComment c) {
@@ -286,11 +277,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
       checkArgument(!notes.containsCommentPublished(c),
           "Cannot update a comment that has already been published and saved");
     }
-    if (c.getSide() == 0) {
+    comments.add(c);
-      commentsForBase.add(c);
-    } else {
-      commentsForPs.add(c);
-    }
   }
 
   private void updateDraftComment(PatchLineComment c) throws OrmException {
@@ -356,31 +343,23 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     if (noteMap == null) {
       noteMap = NoteMap.newEmptyMap();
     }
-    if (commentsForPs.isEmpty() && commentsForBase.isEmpty()) {
+    if (comments.isEmpty()) {
       return null;
     }
 
-    Multimap<PatchSet.Id, PatchLineComment> allCommentsOnBases =
+    Map<RevId, List<PatchLineComment>> allComments = Maps.newHashMap();
-        notes.getBaseComments();
+    for (Map.Entry<RevId, Collection<PatchLineComment>> e
-    Multimap<PatchSet.Id, PatchLineComment> allCommentsOnPs =
+        : notes.getComments().asMap().entrySet()) {
-        notes.getPatchSetComments();
+      List<PatchLineComment> comments = new ArrayList<>();
-
+      for (PatchLineComment c : e.getValue()) {
-    // This writes all comments for the base of this PS to the note map.
+        comments.add(c);
-    if (!commentsForBase.isEmpty()) {
+      }
-      List<PatchLineComment> baseCommentsForThisPs =
+      allComments.put(e.getKey(), comments);
-          new ArrayList<>(allCommentsOnBases.get(psId));
-      baseCommentsForThisPs.addAll(commentsForBase);
-      commentsUtil.writeCommentsToNoteMap(noteMap, baseCommentsForThisPs,
-          inserter);
     }
-
+    for (PatchLineComment c : comments) {
-    // This write all comments for this PS to the note map.
+      addCommentToMap(allComments, c);
-    if (!commentsForPs.isEmpty()) {
-      List<PatchLineComment> commentsForThisPs =
-          new ArrayList<>(allCommentsOnPs.get(psId));
-      commentsForThisPs.addAll(commentsForPs);
-      commentsUtil.writeCommentsToNoteMap(noteMap, commentsForThisPs, inserter);
     }
+    commentsUtil.writeCommentsToNoteMap(noteMap, allComments, inserter);
     return noteMap.writeTree(inserter);
   }
 
@@ -504,8 +483,7 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   private boolean isEmpty() {
     return approvals.isEmpty()
         && changeMessage == null
-        && commentsForBase.isEmpty()
+        && comments.isEmpty()
-        && commentsForPs.isEmpty()
         && reviewers.isEmpty()
         && status == null
         && subject == null

gerrit-server/src/main/java/com/google/gerrit/server/notedb/CommentsInNotesUtil.java

@@ -63,9 +63,11 @@ import java.io.PrintWriter;
 import java.nio.charset.Charset;
 import java.sql.Timestamp;
 import java.text.ParseException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Utility functions to parse PatchLineComments out of a note byte array and
@@ -86,8 +88,7 @@ public class CommentsInNotesUtil {
 
   public static NoteMap parseCommentsFromNotes(Repository repo, String refName,
       RevWalk walk, Change.Id changeId,
-      Multimap<PatchSet.Id, PatchLineComment> commentsForBase,
+      Multimap<RevId, PatchLineComment> comments,
-      Multimap<PatchSet.Id, PatchLineComment> commentsForPs,
       Status status)
       throws IOException, ConfigInvalidException {
     Ref ref = repo.getRef(refName);
@@ -99,20 +100,14 @@ public class CommentsInNotesUtil {
     RevCommit commit = walk.parseCommit(ref.getObjectId());
     NoteMap noteMap = NoteMap.read(reader, commit);
 
-    for (Note note: noteMap) {
+    for (Note note : noteMap) {
       byte[] bytes =
           reader.open(note.getData(), OBJ_BLOB).getCachedBytes(MAX_NOTE_SZ);
       List<PatchLineComment> result = parseNote(bytes, changeId, status);
       if (result == null || result.isEmpty()) {
         continue;
       }
-      PatchSet.Id psId = result.get(0).getKey().getParentKey().getParentKey();
+      comments.putAll(new RevId(note.name()), result);
-      short side = result.get(0).getSide();
-      if (side == 0) {
-        commentsForBase.putAll(psId, result);
-      } else {
-        commentsForPs.putAll(psId, result);
-      }
     }
     return noteMap;
   }
@@ -524,19 +519,47 @@ public class CommentsInNotesUtil {
     return buf.toByteArray();
   }
 
+  /**
+   * Write comments for multiple revisions to a note map.
+   * <p>
+   * Mutates the map in-place. only notes for SHA-1s found as keys in the map
+   * are modified; all other notes are left untouched.
+   *
+   * @param noteMap note map to modify.
+   * @param allComments map of revision to all comments for that revision;
+   *     callers are responsible for reading the original comments and applying
+   *     any changes. Differs from a multimap in that present-but-empty values
+   *     are significant, and indicate the note for that SHA-1 should be
+   *     deleted.
+   * @param inserter object inserter for writing notes.
+   * @throws IOException if an error occurred.
+   */
   public void writeCommentsToNoteMap(NoteMap noteMap,
-      List<PatchLineComment> allComments, ObjectInserter inserter)
+      Map<RevId, List<PatchLineComment>> allComments, ObjectInserter inserter)
-        throws IOException {
+      throws IOException {
-    checkArgument(!allComments.isEmpty(),
+    for (Map.Entry<RevId, List<PatchLineComment>> e : allComments.entrySet()) {
-        "No comments to write; to delete, use removeNoteFromNoteMap().");
+      List<PatchLineComment> comments = e.getValue();
-    ObjectId commit =
+      ObjectId commit = ObjectId.fromString(e.getKey().get());
-        ObjectId.fromString(allComments.get(0).getRevId().get());
+      if (comments.isEmpty()) {
-    Collections.sort(allComments, ChangeNotes.PLC_ORDER);
+        noteMap.remove(commit);
-    noteMap.set(commit, inserter.insert(OBJ_BLOB, buildNote(allComments)));
+        continue;
+      }
+      Collections.sort(comments, ChangeNotes.PLC_ORDER);
+      // We allow comments for multiple commits to be written in the same
+      // update, even though the rest of the metadata update is associated with
+      // a single patch set.
+      noteMap.set(commit, inserter.insert(OBJ_BLOB, buildNote(comments)));
+    }
   }
 
-  public void removeNote(NoteMap noteMap, RevId commitId)
+  static void addCommentToMap(Map<RevId, List<PatchLineComment>> map,
-      throws IOException {
+      PatchLineComment c) {
-    noteMap.remove(ObjectId.fromString(commitId.get()));
+    List<PatchLineComment> list = map.get(c.getRevId());
+    if (list == null) {
+      list = new ArrayList<>();
+      map.put(c.getRevId(), list);
+    }
+    list.add(c);
   }
+
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/DraftCommentNotes.java

@@ -14,18 +14,14 @@
 
 package com.google.gerrit.server.notedb;
 
-import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
-
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.ImmutableListMultimap;
-import com.google.common.collect.Multimap;
-import com.google.common.collect.Table;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
-import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.Project;
 import com.google.gerrit.reviewdb.client.RefNames;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.config.AllUsersName;
 import com.google.gerrit.server.config.AllUsersNameProvider;
 import com.google.gerrit.server.git.GitRepositoryManager;
@@ -70,8 +66,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   private final AllUsersName draftsProject;
   private final Account.Id author;
 
-  private final Table<PatchSet.Id, String, PatchLineComment> draftBaseComments;
+  private ImmutableListMultimap<RevId, PatchLineComment> comments;
-  private final Table<PatchSet.Id, String, PatchLineComment> draftPsComments;
   private NoteMap noteMap;
 
   DraftCommentNotes(GitRepositoryManager repoManager, NotesMigration migration,
@@ -79,9 +74,6 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     super(repoManager, migration, changeId);
     this.draftsProject = draftsProject;
     this.author = author;
-
-    this.draftBaseComments = HashBasedTable.create();
-    this.draftPsComments = HashBasedTable.create();
   }
 
   public NoteMap getNoteMap() {
@@ -92,32 +84,18 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     return author;
   }
 
-  /**
+  public ImmutableListMultimap<RevId, PatchLineComment> getComments() {
-   * @return a defensive copy of the table containing all draft comments
+    // TODO(dborowitz): Defensive copy?
-   *    on this change with side == 0. The row key is the comment's PatchSet.Id,
+    return comments;
-   *    the column key is the comment's UUID, and the value is the comment.
-   */
-  public Table<PatchSet.Id, String, PatchLineComment>
-      getDraftBaseComments() {
-    return HashBasedTable.create(draftBaseComments);
-  }
-
-  /**
-   * @return a defensive copy of the table containing all draft comments
-   *    on this change with side == 1. The row key is the comment's PatchSet.Id,
-   *    the column key is the comment's UUID, and the value is the comment.
-   */
-  public Table<PatchSet.Id, String, PatchLineComment>
-      getDraftPsComments() {
-    return HashBasedTable.create(draftPsComments);
   }
 
   public boolean containsComment(PatchLineComment c) {
-    Table<PatchSet.Id, String, PatchLineComment> t =
+    for (PatchLineComment existing : comments.values()) {
-        c.getSide() == (short) 0
+      if (c.getKey().equals(existing.getKey())) {
-        ? getDraftBaseComments()
+        return true;
-        : getDraftPsComments();
+      }
-    return t.contains(getCommentPsId(c), c.getKey().get());
+    }
+    return false;
   }
 
   @Override
@@ -129,6 +107,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   protected void onLoad() throws IOException, ConfigInvalidException {
     ObjectId rev = getRevision();
     if (rev == null) {
+      loadDefaults();
       return;
     }
 
@@ -137,8 +116,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
           getChangeId(), walk, rev, repoManager, draftsProject, author)) {
       parser.parseDraftComments();
 
-      buildCommentTable(draftBaseComments, parser.draftBaseComments);
+      comments = ImmutableListMultimap.copyOf(parser.comments);
-      buildCommentTable(draftPsComments, parser.draftPsComments);
       noteMap = parser.noteMap;
     }
   }
@@ -152,20 +130,11 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
 
   @Override
   protected void loadDefaults() {
-    // Do nothing; tables are final and initialized in constructor.
+    comments = ImmutableListMultimap.of();
   }
 
   @Override
   protected Project.NameKey getProjectName() {
     return draftsProject;
   }
-
-  private void buildCommentTable(
-      Table<PatchSet.Id, String, PatchLineComment> commentTable,
-      Multimap<PatchSet.Id, PatchLineComment> allComments) {
-    for (PatchLineComment c : allComments.values()) {
-      commentTable.put(getCommentPsId(c), c.getKey().get(), c);
-    }
-  }
-
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/DraftCommentNotesParser.java

@@ -19,8 +19,8 @@ import com.google.common.collect.Multimap;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
-import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.RefNames;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.config.AllUsersName;
 import com.google.gerrit.server.git.GitRepositoryManager;
 
@@ -34,8 +34,7 @@ import org.eclipse.jgit.revwalk.RevWalk;
 import java.io.IOException;
 
 class DraftCommentNotesParser implements AutoCloseable {
-  final Multimap<PatchSet.Id, PatchLineComment> draftBaseComments;
+  final Multimap<RevId, PatchLineComment> comments;
-  final Multimap<PatchSet.Id, PatchLineComment> draftPsComments;
   NoteMap noteMap;
 
   private final Change.Id changeId;
@@ -53,8 +52,7 @@ class DraftCommentNotesParser implements AutoCloseable {
     this.repo = repoManager.openMetadataRepository(draftsProject);
     this.author = author;
 
-    draftBaseComments = ArrayListMultimap.create();
+    comments = ArrayListMultimap.create();
-    draftPsComments = ArrayListMultimap.create();
   }
 
   @Override
@@ -66,7 +64,6 @@ class DraftCommentNotesParser implements AutoCloseable {
     walk.markStart(walk.parseCommit(tip));
     noteMap = CommentsInNotesUtil.parseCommentsFromNotes(repo,
         RefNames.refsDraftComments(author, changeId),
-        walk, changeId, draftBaseComments,
+        walk, changeId, comments, PatchLineComment.Status.DRAFT);
-        draftPsComments, PatchLineComment.Status.DRAFT);
   }
 }

gerrit-server/src/test/java/com/google/gerrit/server/change/CommentsTest.java

@@ -234,23 +234,23 @@ public class CommentsTest  {
     plc1 = newPatchLineComment(psId1, "Comment1", null,
         "FileOne.txt", Side.REVISION, 3, ownerId, timeBase,
         "First Comment", new CommentRange(1, 2, 3, 4));
-    plc1.setRevId(new RevId("ABCDABCDABCDABCDABCDABCDABCDABCDABCDABCD"));
+    plc1.setRevId(new RevId("abcdabcdabcdabcdabcdabcdabcdabcdabcdabcd"));
     plc2 = newPatchLineComment(psId1, "Comment2", "Comment1",
         "FileOne.txt", Side.REVISION, 3, otherUserId, timeBase + 1000,
         "Reply to First Comment",  new CommentRange(1, 2, 3, 4));
-    plc2.setRevId(new RevId("ABCDABCDABCDABCDABCDABCDABCDABCDABCDABCD"));
+    plc2.setRevId(new RevId("abcdabcdabcdabcdabcdabcdabcdabcdabcdabcd"));
     plc3 = newPatchLineComment(psId1, "Comment3", "Comment1",
         "FileOne.txt", Side.PARENT, 3, ownerId, timeBase + 2000,
         "First Parent Comment",  new CommentRange(1, 2, 3, 4));
-    plc3.setRevId(new RevId("CDEFCDEFCDEFCDEFCDEFCDEFCDEFCDEFCDEFCDEF"));
+    plc3.setRevId(new RevId("cdefcdefcdefcdefcdefcdefcdefcdefcdefcdef"));
     plc4 = newPatchLineComment(psId2, "Comment4", null, "FileOne.txt",
         Side.REVISION, 3, ownerId, timeBase + 3000, "Second Comment",
         new CommentRange(1, 2, 3, 4), Status.DRAFT);
-    plc4.setRevId(new RevId("BCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDE"));
+    plc4.setRevId(new RevId("bcdebcdebcdebcdebcdebcdebcdebcdebcdebcde"));
     plc5 = newPatchLineComment(psId2, "Comment5", null, "FileOne.txt",
         Side.REVISION, 5, ownerId, timeBase + 4000, "Third Comment",
         new CommentRange(3, 4, 5, 6), Status.DRAFT);
-    plc5.setRevId(new RevId("BCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDEBCDE"));
+    plc5.setRevId(new RevId("bcdebcdebcdebcdebcdebcdebcdebcdebcdebcde"));
 
     List<PatchLineComment> commentsByOwner = Lists.newArrayList();
     commentsByOwner.add(plc1);

gerrit-server/src/test/java/com/google/gerrit/server/notedb/AbstractChangeNotesTest.java

@@ -203,16 +203,16 @@ public class AbstractChangeNotesTest {
     return label;
   }
 
-  protected PatchLineComment newPublishedPatchLineComment(PatchSet.Id psId,
+  protected PatchLineComment newPublishedComment(PatchSet.Id psId,
       String filename, String UUID, CommentRange range, int line,
       IdentifiedUser commenter, String parentUUID, Timestamp t,
       String message, short side, String commitSHA1) {
-    return newPatchLineComment(psId, filename, UUID, range, line, commenter,
+    return newComment(psId, filename, UUID, range, line, commenter,
         parentUUID, t, message, side, commitSHA1,
         PatchLineComment.Status.PUBLISHED);
   }
 
-  protected PatchLineComment newPatchLineComment(PatchSet.Id psId,
+  protected PatchLineComment newComment(PatchSet.Id psId,
       String filename, String UUID, CommentRange range, int line,
       IdentifiedUser commenter, String parentUUID, Timestamp t,
       String message, short side, String commitSHA1,

gerrit-server/src/test/java/com/google/gerrit/server/notedb/ChangeNotesTest.java

@@ -22,12 +22,11 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.LinkedListMultimap;
 import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Multimap;
 import com.google.gerrit.common.TimeUtil;
 import com.google.gerrit.common.data.SubmitRecord;
 import com.google.gerrit.reviewdb.client.Account;
@@ -38,8 +37,8 @@ import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchLineComment.Status;
 import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.PatchSetApproval;
+import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gerrit.server.git.VersionedMetaData.BatchMetaDataUpdate;
-import com.google.gwtorm.server.OrmException;
 
 import org.eclipse.jgit.lib.BatchRefUpdate;
 import org.eclipse.jgit.lib.CommitBuilder;
@@ -47,12 +46,12 @@ import org.eclipse.jgit.lib.Constants;
 import org.eclipse.jgit.lib.NullProgressMonitor;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.notes.Note;
+import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.RevCommit;
 import org.eclipse.jgit.revwalk.RevWalk;
 import org.eclipse.jgit.transport.ReceiveCommand;
 import org.junit.Test;
 
-import java.io.IOException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
 import java.util.LinkedHashSet;
@@ -424,7 +423,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     PatchSet.Id psId = c.currentPatchSetId();
     BatchRefUpdate bru = repo.getRefDatabase().newBatchUpdate();
     BatchMetaDataUpdate batch = update1.openUpdateInBatch(bru);
-    PatchLineComment comment1 = newPublishedPatchLineComment(psId, "file1",
+    PatchLineComment comment1 = newPublishedComment(psId, "file1",
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update1.setPatchSetId(psId);
@@ -451,7 +450,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
       ImmutableListMultimap<PatchSet.Id, PatchSetApproval> approvals1 =
           notesWithComments.buildApprovals();
       assertThat(approvals1).isEmpty();
-      assertThat(notesWithComments.commentsForBase).hasSize(1);
+      assertThat(notesWithComments.comments).hasSize(1);
       notesWithComments.close();
 
       ChangeNotesParser notesWithApprovals =
@@ -460,7 +459,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
       ImmutableListMultimap<PatchSet.Id, PatchSetApproval> approvals2 =
           notesWithApprovals.buildApprovals();
       assertThat(approvals2).hasSize(1);
-      assertThat(notesWithApprovals.commentsForBase).hasSize(1);
+      assertThat(notesWithApprovals.comments).hasSize(1);
       notesWithApprovals.close();
     } finally {
       batch.close();
@@ -674,7 +673,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     Timestamp time3 = TimeUtil.nowTs();
     PatchSet.Id psId = c.currentPatchSetId();
 
-    PatchLineComment comment1 = newPublishedPatchLineComment(psId, "file1",
+    PatchLineComment comment1 = newPublishedComment(psId, "file1",
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
@@ -683,7 +682,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     update = newUpdate(c, otherUser);
     CommentRange range2 = new CommentRange(2, 1, 3, 1);
-    PatchLineComment comment2 = newPublishedPatchLineComment(psId, "file1",
+    PatchLineComment comment2 = newPublishedComment(psId, "file1",
         uuid2, range2, range2.getEndLine(), otherUser, null, time2, message2,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
@@ -692,7 +691,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     update = newUpdate(c, otherUser);
     CommentRange range3 = new CommentRange(3, 1, 4, 1);
-    PatchLineComment comment3 = newPublishedPatchLineComment(psId, "file2",
+    PatchLineComment comment3 = newPublishedComment(psId, "file2",
         uuid3, range3, range3.getEndLine(), otherUser, null, time3, message3,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
@@ -753,7 +752,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     Timestamp time2 = TimeUtil.nowTs();
     PatchSet.Id psId = c.currentPatchSetId();
 
-    PatchLineComment comment1 = newPublishedPatchLineComment(psId, "file1",
+    PatchLineComment comment1 = newPublishedComment(psId, "file1",
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
@@ -762,7 +761,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     update = newUpdate(c, otherUser);
     CommentRange range2 = new CommentRange(2, 1, 3, 1);
-    PatchLineComment comment2 = newPublishedPatchLineComment(psId, "file1",
+    PatchLineComment comment2 = newPublishedComment(psId, "file1",
         uuid2, range2, range2.getEndLine(), otherUser, null, time2, message2,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
@@ -808,6 +807,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     ChangeUpdate update = newUpdate(c, otherUser);
     String uuid1 = "uuid1";
     String uuid2 = "uuid2";
+    String rev1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
+    String rev2 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     String messageForBase = "comment for base";
     String messageForPS = "comment for ps";
     CommentRange range = new CommentRange(1, 1, 2, 1);
@@ -815,32 +816,26 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     PatchSet.Id psId = c.currentPatchSetId();
 
     PatchLineComment commentForBase =
-        newPublishedPatchLineComment(psId, "filename", uuid1,
+        newPublishedComment(psId, "filename", uuid1,
         range, range.getEndLine(), otherUser, null, now, messageForBase,
-        (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        (short) 0, rev1);
     update.setPatchSetId(psId);
     update.upsertComment(commentForBase);
     update.commit();
 
     update = newUpdate(c, otherUser);
     PatchLineComment commentForPS =
-        newPublishedPatchLineComment(psId, "filename", uuid2,
+        newPublishedComment(psId, "filename", uuid2,
         range, range.getEndLine(), otherUser, null, now, messageForPS,
-        (short) 1, "abcd4567abcd4567abcd4567abcd4567abcd4567");
+        (short) 1, rev2);
     update.setPatchSetId(psId);
     update.upsertComment(commentForPS);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    Multimap<PatchSet.Id, PatchLineComment> commentsForBase =
+        ImmutableMultimap.of(
-        notes.getBaseComments();
+            new RevId(rev1), commentForBase,
-    Multimap<PatchSet.Id, PatchLineComment> commentsForPS =
+            new RevId(rev2), commentForPS));
-        notes.getPatchSetComments();
-    assertThat(commentsForBase).hasSize(1);
-    assertThat(commentsForPS).hasSize(1);
-
-    assertThat(commentsForBase.get(psId)).containsExactly(commentForBase);
-    assertThat(commentsForPS.get(psId)).containsExactly(commentForPS);
   }
 
   @Test
@@ -848,6 +843,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     Change c = newChange();
     String uuid1 = "uuid1";
     String uuid2 = "uuid2";
+    String rev = "abcd1234abcd1234abcd1234abcd1234abcd1234";
     CommentRange range = new CommentRange(1, 1, 2, 1);
     PatchSet.Id psId = c.currentPatchSetId();
     String filename = "filename";
@@ -856,31 +852,25 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     ChangeUpdate update = newUpdate(c, otherUser);
     Timestamp timeForComment1 = TimeUtil.nowTs();
     Timestamp timeForComment2 = TimeUtil.nowTs();
-    PatchLineComment comment1 = newPublishedPatchLineComment(psId, filename,
+    PatchLineComment comment1 = newPublishedComment(psId, filename,
         uuid1, range, range.getEndLine(), otherUser, null, timeForComment1,
-        "comment 1", side, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        "comment 1", side, rev);
     update.setPatchSetId(psId);
     update.upsertComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
-    PatchLineComment comment2 = newPublishedPatchLineComment(psId, filename,
+    PatchLineComment comment2 = newPublishedComment(psId, filename,
         uuid2, range, range.getEndLine(), otherUser, null, timeForComment2,
-        "comment 2", side, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        "comment 2", side, rev);
     update.setPatchSetId(psId);
     update.upsertComment(comment2);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    Multimap<PatchSet.Id, PatchLineComment> commentsForBase =
+        ImmutableMultimap.of(
-        notes.getBaseComments();
+          new RevId(rev), comment1,
-    Multimap<PatchSet.Id, PatchLineComment> commentsForPS =
+          new RevId(rev), comment2)).inOrder();
-        notes.getPatchSetComments();
-    assertThat(commentsForBase).isEmpty();
-    assertThat(commentsForPS).hasSize(2);
-
-    assertThat(commentsForPS.get(psId))
-        .containsExactly(comment1, comment2).inOrder();
   }
 
   @Test
@@ -888,6 +878,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
       throws Exception {
     Change c = newChange();
     String uuid = "uuid";
+    String rev = "abcd1234abcd1234abcd1234abcd1234abcd1234";
     CommentRange range = new CommentRange(1, 1, 2, 1);
     PatchSet.Id psId = c.currentPatchSetId();
     String filename1 = "filename1";
@@ -896,37 +887,33 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     ChangeUpdate update = newUpdate(c, otherUser);
     Timestamp now = TimeUtil.nowTs();
-    PatchLineComment comment1 = newPublishedPatchLineComment(psId, filename1,
+    PatchLineComment comment1 = newPublishedComment(psId, filename1,
         uuid, range, range.getEndLine(), otherUser, null, now, "comment 1",
-        side, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        side, rev);
     update.setPatchSetId(psId);
     update.upsertComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
-    PatchLineComment comment2 = newPublishedPatchLineComment(psId, filename2,
+    PatchLineComment comment2 = newPublishedComment(psId, filename2,
         uuid, range, range.getEndLine(), otherUser, null, now, "comment 2",
-        side, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        side, rev);
     update.setPatchSetId(psId);
     update.upsertComment(comment2);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    ListMultimap<PatchSet.Id, PatchLineComment> commentsForBase =
+        ImmutableMultimap.of(
-        notes.getBaseComments();
+          new RevId(rev), comment1,
-    ListMultimap<PatchSet.Id, PatchLineComment> commentsForPS =
+          new RevId(rev), comment2)).inOrder();
-        notes.getPatchSetComments();
-    assertThat(commentsForBase).isEmpty();
-    assertThat(commentsForPS).hasSize(2);
-
-    assertThat(commentsForPS.get(psId))
-        .containsExactly(comment1, comment2).inOrder();
   }
 
   @Test
   public void patchLineCommentMultiplePatchsets() throws Exception {
     Change c = newChange();
     String uuid = "uuid";
+    String rev1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
+    String rev2 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     CommentRange range = new CommentRange(1, 1, 2, 1);
     PatchSet.Id ps1 = c.currentPatchSetId();
     String filename = "filename1";
@@ -934,9 +921,9 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     ChangeUpdate update = newUpdate(c, otherUser);
     Timestamp now = TimeUtil.nowTs();
-    PatchLineComment comment1 = newPublishedPatchLineComment(ps1, filename,
+    PatchLineComment comment1 = newPublishedComment(ps1, filename,
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps1",
-        side, "abcd1234abcd1234abcd1234abcd1234abcd1234");
+        side, rev1);
     update.setPatchSetId(ps1);
     update.upsertComment(comment1);
     update.commit();
@@ -946,31 +933,24 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     update = newUpdate(c, otherUser);
     now = TimeUtil.nowTs();
-    PatchLineComment comment2 = newPublishedPatchLineComment(ps2, filename,
+    PatchLineComment comment2 = newPublishedComment(ps2, filename,
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps2",
-        side, "abcd4567abcd4567abcd4567abcd4567abcd4567");
+        side, rev2);
     update.setPatchSetId(ps2);
     update.upsertComment(comment2);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    LinkedListMultimap<PatchSet.Id, PatchLineComment> commentsForBase =
+        ImmutableMultimap.of(
-        LinkedListMultimap.create(notes.getBaseComments());
+          new RevId(rev1), comment1,
-    LinkedListMultimap<PatchSet.Id, PatchLineComment> commentsForPS =
+          new RevId(rev2), comment2));
-        LinkedListMultimap.create(notes.getPatchSetComments());
-    assertThat(commentsForBase).isEmpty();
-    assertThat(commentsForPS).hasSize(2);
-
-    assertThat(commentsForPS).containsExactly(
-        ImmutableListMultimap.of(
-          ps1, comment1,
-          ps2, comment2));
   }
 
   @Test
   public void patchLineCommentSingleDraftToPublished() throws Exception {
     Change c = newChange();
     String uuid = "uuid";
+    String rev = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     CommentRange range = new CommentRange(1, 1, 2, 1);
     PatchSet.Id ps1 = c.currentPatchSetId();
     String filename = "filename1";
@@ -978,16 +958,17 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     ChangeUpdate update = newUpdate(c, otherUser);
     Timestamp now = TimeUtil.nowTs();
-    PatchLineComment comment1 = newPatchLineComment(ps1, filename, uuid,
+    PatchLineComment comment1 = newComment(ps1, filename, uuid, range,
-        range, range.getEndLine(), otherUser, null, now, "comment on ps1", side,
+        range.getEndLine(), otherUser, null, now, "comment on ps1", side,
-        "abcd4567abcd4567abcd4567abcd4567abcd4567", Status.DRAFT);
+        rev, Status.DRAFT);
     update.setPatchSetId(ps1);
     update.insertComment(comment1);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
-    assertThat(notes.getDraftPsComments(otherUserId)).hasSize(1);
+    assertThat(notes.getDraftComments(otherUserId)).containsExactly(
-    assertThat(notes.getDraftBaseComments(otherUserId)).isEmpty();
+        ImmutableMultimap.of(new RevId(rev), comment1));
+    assertThat(notes.getComments()).isEmpty();
 
     comment1.setStatus(Status.PUBLISHED);
     update = newUpdate(c, otherUser);
@@ -996,46 +977,44 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.commit();
 
     notes = newNotes(c);
-
+    assertThat(notes.getDraftComments(otherUserId)).isEmpty();
-    assertThat(notes.getDraftPsComments(otherUserId).values()).isEmpty();
+    assertThat(notes.getComments()).containsExactly(
-    assertThat(notes.getDraftBaseComments(otherUserId).values()).isEmpty();
+        ImmutableMultimap.of(new RevId(rev), comment1));
-
-    assertThat(notes.getBaseComments()).isEmpty();
-    assertThat(notes.getPatchSetComments().values()).containsExactly(comment1);
   }
 
   @Test
   public void patchLineCommentMultipleDraftsSameSidePublishOne()
-      throws OrmException, IOException {
+      throws Exception {
     Change c = newChange();
     String uuid1 = "uuid1";
     String uuid2 = "uuid2";
+    String rev = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     CommentRange range1 = new CommentRange(1, 1, 2, 2);
     CommentRange range2 = new CommentRange(2, 2, 3, 3);
     String filename = "filename1";
     short side = (short) 1;
     Timestamp now = TimeUtil.nowTs();
-    String commitSHA1 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     PatchSet.Id psId = c.currentPatchSetId();
 
     // Write two drafts on the same side of one patch set.
     ChangeUpdate update = newUpdate(c, otherUser);
     update.setPatchSetId(psId);
-    PatchLineComment comment1 = newPatchLineComment(psId, filename, uuid1,
+    PatchLineComment comment1 = newComment(psId, filename, uuid1,
         range1, range1.getEndLine(), otherUser, null, now, "comment on ps1",
-        side, commitSHA1, Status.DRAFT);
+        side, rev, Status.DRAFT);
-    PatchLineComment comment2 = newPatchLineComment(psId, filename, uuid2,
+    PatchLineComment comment2 = newComment(psId, filename, uuid2,
         range2, range2.getEndLine(), otherUser, null, now, "other on ps1",
-        side, commitSHA1, Status.DRAFT);
+        side, rev, Status.DRAFT);
     update.insertComment(comment1);
     update.insertComment(comment2);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
-    assertThat(notes.getDraftBaseComments(otherUserId)).isEmpty();
+    assertThat(notes.getDraftComments(otherUserId)).containsExactly(
-
+        ImmutableMultimap.of(
-    assertThat(notes.getDraftPsComments(otherUserId).values())
+          new RevId(rev), comment1,
-        .containsExactly(comment1, comment2);
+          new RevId(rev), comment2)).inOrder();
+    assertThat(notes.getComments()).isEmpty();
 
     // Publish first draft.
     update = newUpdate(c, otherUser);
@@ -1045,47 +1024,46 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.commit();
 
     notes = newNotes(c);
-    assertThat(notes.getPatchSetComments().get(psId)).containsExactly(comment1);
+    assertThat(notes.getDraftComments(otherUserId)).containsExactly(
-    assertThat(notes.getDraftPsComments(otherUserId).values())
+        ImmutableMultimap.of(new RevId(rev), comment2));
-        .containsExactly(comment2);
+    assertThat(notes.getComments()).containsExactly(
-
+        ImmutableMultimap.of(new RevId(rev), comment1));
-    assertThat(notes.getBaseComments()).isEmpty();
-    assertThat(notes.getDraftBaseComments(otherUserId)).isEmpty();
   }
 
   @Test
   public void patchLineCommentsMultipleDraftsBothSidesPublishAll()
-      throws OrmException, IOException {
+      throws Exception {
     Change c = newChange();
     String uuid1 = "uuid1";
     String uuid2 = "uuid2";
+    String rev1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
+    String rev2 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
     CommentRange range1 = new CommentRange(1, 1, 2, 2);
     CommentRange range2 = new CommentRange(2, 2, 3, 3);
     String filename = "filename1";
     Timestamp now = TimeUtil.nowTs();
-    String commitSHA1 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
-    String baseSHA1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
     PatchSet.Id psId = c.currentPatchSetId();
 
     // Write two drafts, one on each side of the patchset.
     ChangeUpdate update = newUpdate(c, otherUser);
     update.setPatchSetId(psId);
-    PatchLineComment baseComment = newPatchLineComment(psId, filename, uuid1,
+    PatchLineComment baseComment = newComment(psId, filename, uuid1,
         range1, range1.getEndLine(), otherUser, null, now, "comment on base",
-        (short) 0, baseSHA1, Status.DRAFT);
+        (short) 0, rev1, Status.DRAFT);
-    PatchLineComment psComment = newPatchLineComment(psId, filename, uuid2,
+    PatchLineComment psComment = newComment(psId, filename, uuid2,
         range2, range2.getEndLine(), otherUser, null, now, "comment on ps",
-        (short) 1, commitSHA1, Status.DRAFT);
+        (short) 1, rev2, Status.DRAFT);
 
     update.insertComment(baseComment);
     update.insertComment(psComment);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
-    assertThat(notes.getDraftBaseComments(otherUserId).values())
+    assertThat(notes.getDraftComments(otherUserId)).containsExactly(
-        .containsExactly(baseComment);
+        ImmutableMultimap.of(
-    assertThat(notes.getDraftPsComments(otherUserId).values())
+            new RevId(rev1), baseComment,
-        .containsExactly(psComment);
+            new RevId(rev2), psComment));
+    assertThat(notes.getComments()).isEmpty();
 
     // Publish both comments.
     update = newUpdate(c, otherUser);
@@ -1098,13 +1076,98 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.commit();
 
     notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId)).isEmpty();
+    assertThat(notes.getComments()).containsExactly(
+        ImmutableMultimap.of(
+            new RevId(rev1), baseComment,
+            new RevId(rev2), psComment));
+  }
 
-    assertThat(notes.getBaseComments().get(psId)).containsExactly(baseComment);
+  @Test
-    assertThat(notes.getPatchSetComments().get(psId))
+  public void patchLineCommentsDeleteAllDrafts() throws Exception {
-        .containsExactly(psComment);
+    Change c = newChange();
+    String uuid = "uuid";
+    String rev = "abcd1234abcd1234abcd1234abcd1234abcd1234";
+    ObjectId objId = ObjectId.fromString(rev);
+    CommentRange range = new CommentRange(1, 1, 2, 1);
+    PatchSet.Id psId = c.currentPatchSetId();
+    String filename = "filename";
+    short side = (short) 1;
 
-    assertThat(notes.getDraftBaseComments(otherUserId)).isEmpty();
+    ChangeUpdate update = newUpdate(c, otherUser);
-    assertThat(notes.getDraftPsComments(otherUserId)).isEmpty();
+    Timestamp now = TimeUtil.nowTs();
+    PatchLineComment comment = newComment(psId, filename, uuid, range,
+        range.getEndLine(), otherUser, null, now, "comment on ps1", side,
+        rev, Status.DRAFT);
+    update.setPatchSetId(psId);
+    update.upsertComment(comment);
+    update.commit();
+
+    ChangeNotes notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId)).hasSize(1);
+    assertThat(notes.getDraftCommentNotes().getNoteMap().contains(objId))
+      .isTrue();
+
+    update = newUpdate(c, otherUser);
+    now = TimeUtil.nowTs();
+    update.setPatchSetId(psId);
+    update.deleteComment(comment);
+    update.commit();
+
+    notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId)).isEmpty();
+    assertThat(notes.getDraftCommentNotes().getNoteMap()).isNull();
+  }
+
+  @Test
+  public void patchLineCommentsDeleteAllDraftsForOneRevision()
+      throws Exception {
+    Change c = newChange();
+    String uuid = "uuid";
+    String rev1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
+    String rev2 = "abcd4567abcd4567abcd4567abcd4567abcd4567";
+    ObjectId objId1 = ObjectId.fromString(rev1);
+    ObjectId objId2 = ObjectId.fromString(rev2);
+    CommentRange range = new CommentRange(1, 1, 2, 1);
+    PatchSet.Id ps1 = c.currentPatchSetId();
+    String filename = "filename1";
+    short side = (short) 1;
+
+    ChangeUpdate update = newUpdate(c, otherUser);
+    Timestamp now = TimeUtil.nowTs();
+    PatchLineComment comment1 = newComment(ps1, filename,
+        uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps1",
+        side, rev1, Status.DRAFT);
+    update.setPatchSetId(ps1);
+    update.upsertComment(comment1);
+    update.commit();
+
+    incrementPatchSet(c);
+    PatchSet.Id ps2 = c.currentPatchSetId();
+
+    update = newUpdate(c, otherUser);
+    now = TimeUtil.nowTs();
+    PatchLineComment comment2 = newComment(ps2, filename,
+        uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps2",
+        side, rev2, Status.DRAFT);
+    update.setPatchSetId(ps2);
+    update.upsertComment(comment2);
+    update.commit();
+
+    ChangeNotes notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId)).hasSize(2);
+
+    update = newUpdate(c, otherUser);
+    now = TimeUtil.nowTs();
+    update.setPatchSetId(ps2);
+    update.deleteComment(comment2);
+    update.commit();
+
+    notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId)).hasSize(1);
+    NoteMap noteMap = notes.getDraftCommentNotes().getNoteMap();
+    assertThat(noteMap.contains(objId1)).isTrue();
+    assertThat(noteMap.contains(objId2)).isFalse();
   }
 
   @Test
@@ -1112,22 +1175,20 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     Change c = newChange();
     ChangeUpdate update = newUpdate(c, otherUser);
     String uuid = "uuid";
+    String rev = "abcd1234abcd1234abcd1234abcd1234abcd1234";
     String messageForBase = "comment for base";
     Timestamp now = TimeUtil.nowTs();
     PatchSet.Id psId = c.currentPatchSetId();
 
-    PatchLineComment commentForBase =
+    PatchLineComment comment = newPublishedComment(
-        newPublishedPatchLineComment(psId, "filename", uuid,
+        psId, "filename", uuid, null, 0, otherUser, null, now, messageForBase,
-        null, 0, otherUser, null, now, messageForBase,
+        (short) 0, rev);
-        (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(commentForBase);
+    update.upsertComment(comment);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    assertThat(notes.getPatchSetComments()).isEmpty();
+        ImmutableMultimap.of(new RevId(rev), comment));
-    assertThat(notes.getBaseComments().get(psId))
-        .containsExactly(commentForBase);
   }
 
   @Test
@@ -1135,21 +1196,19 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     Change c = newChange();
     ChangeUpdate update = newUpdate(c, otherUser);
     String uuid = "uuid";
+    String rev = "abcd1234abcd1234abcd1234abcd1234abcd1234";
     String messageForBase = "comment for base";
     Timestamp now = TimeUtil.nowTs();
     PatchSet.Id psId = c.currentPatchSetId();
 
-    PatchLineComment commentForBase =
+    PatchLineComment comment = newPublishedComment(
-        newPublishedPatchLineComment(psId, "filename", uuid,
+        psId, "filename", uuid, null, 1, otherUser, null, now, messageForBase,
-        null, 1, otherUser, null, now, messageForBase,
+        (short) 0, rev);
-        (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(commentForBase);
+    update.upsertComment(comment);
     update.commit();
 
-    ChangeNotes notes = newNotes(c);
+    assertThat(newNotes(c).getComments()).containsExactly(
-    assertThat(notes.getPatchSetComments()).isEmpty();
+        ImmutableMultimap.of(new RevId(rev), comment));
-    assertThat(notes.getBaseComments().get(psId))
-        .containsExactly(commentForBase);
   }
 }