Avoid unnecessary copying when building Merkle trees.

tjgq · copybara-github · commit f63ce7973b8f · 2023-03-30T02:56:10.000-07:00
Instead of accumulating a single set of children in DirectoryTreeBuilder and later splitting it up into file, symlink and subdirectory sets, we can accumulate the latter directly. Closes bazelbuild#17912. PiperOrigin-RevId: 520585350 Change-Id: I02b26825976c72d59462a66ffd9afaec3d7c4176
diff --git a/src/main/java/com/google/devtools/build/lib/remote/merkletree/DirectoryTree.java b/src/main/java/com/google/devtools/build/lib/remote/merkletree/DirectoryTree.java
@@ -16,16 +16,15 @@
 import build.bazel.remote.execution.v2.Digest;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.protobuf.ByteString;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.SortedSet;
+import java.util.TreeSet;
 
 /**
  * Intermediate tree representation of a list of lexicographically sorted list of files. Each node
@@ -34,11 +33,12 @@
 final class DirectoryTree {
 
   interface Visitor {
+
     void visitDirectory(
         PathFragment dirname,
-        List<FileNode> files,
-        List<SymlinkNode> symlinks,
-        List<DirectoryNode> dirs);
+        SortedSet<FileNode> files,
+        SortedSet<SymlinkNode> symlinks,
+        SortedSet<DirectoryNode> dirs);
   }
 
   abstract static class Node implements Comparable<Node> {
@@ -204,26 +204,44 @@ public String toString() {
   }
 
   static class DirectoryNode extends Node {
-    private final SortedSet<Node> children = Sets.newTreeSet();
+
+    private final SortedSet<FileNode> files = new TreeSet<>();
+    private final SortedSet<SymlinkNode> symlinks = new TreeSet<>();
+    private final SortedSet<DirectoryNode> subdirs = new TreeSet<>();
 
     DirectoryNode(String pathSegment) {
       super(pathSegment);
     }
 
-    boolean addChild(Node child) {
-      return children.add(Preconditions.checkNotNull(child, "child"));
+    @CanIgnoreReturnValue
+    boolean addChild(FileNode file) {
+      return files.add(Preconditions.checkNotNull(file, "file"));
+    }
+
+    @CanIgnoreReturnValue
+    boolean addChild(SymlinkNode symlink) {
+      return symlinks.add(Preconditions.checkNotNull(symlink, "symlink"));
+    }
+
+    @CanIgnoreReturnValue
+    boolean addChild(DirectoryNode subdir) {
+      return subdirs.add(Preconditions.checkNotNull(subdir, "subdir"));
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(super.hashCode(), children.hashCode());
+      return Objects.hash(
+          super.hashCode(), files.hashCode(), symlinks.hashCode(), subdirs.hashCode());
     }
 
     @Override
     public boolean equals(Object o) {
       if (o instanceof DirectoryNode) {
         DirectoryNode other = (DirectoryNode) o;
-        return super.equals(other) && Objects.equals(children, other.children);
+        return super.equals(other)
+            && Objects.equals(files, other.files)
+            && Objects.equals(symlinks, other.symlinks)
+            && Objects.equals(subdirs, other.subdirs);
       }
       return false;
     }
@@ -265,23 +283,10 @@ private void visit(Visitor visitor, PathFragment dirname) {
       return;
     }
 
-    List<FileNode> files = new ArrayList<>(dir.children.size());
-    List<SymlinkNode> symlinks = new ArrayList<>();
-    List<DirectoryNode> dirs = new ArrayList<>();
-    for (Node child : dir.children) {
-      if (child instanceof FileNode) {
-        files.add((FileNode) child);
-      } else if (child instanceof SymlinkNode) {
-        symlinks.add((SymlinkNode) child);
-      } else if (child instanceof DirectoryNode) {
-        dirs.add((DirectoryNode) child);
-        visit(visitor, dirname.getRelative(child.pathSegment));
-      } else {
-        throw new IllegalStateException(
-            String.format("Node type '%s' is not supported", child.getClass().getSimpleName()));
-      }
+    for (DirectoryNode subdir : dir.subdirs) {
+      visit(visitor, dirname.getRelative(subdir.getPathSegment()));
     }
-    visitor.visitDirectory(dirname, files, symlinks, dirs);
+    visitor.visitDirectory(dirname, dir.files, dir.symlinks, dir.subdirs);
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTree.java b/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTree.java
@@ -28,7 +28,6 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
-import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.MetadataProvider;
@@ -299,8 +298,7 @@ private static MerkleTree build(DirectoryTree tree, DigestUtil digestUtil) {
                     m.remove(subDirname), "subMerkleTree at '%s' was null", subDirname);
             subDirs.put(dir.getPathSegment(), subMerkleTree);
           }
-          MerkleTree mt =
-              buildMerkleTree(new TreeSet<>(files), new TreeSet<>(symlinks), subDirs, digestUtil);
+          MerkleTree mt = buildMerkleTree(files, symlinks, subDirs, digestUtil);
           m.put(dirname, mt);
         });
     MerkleTree rootMerkleTree = m.get(PathFragment.EMPTY_FRAGMENT);
@@ -326,11 +324,11 @@ public static MerkleTree merge(Collection<MerkleTree> merkleTrees, DigestUtil di
     }
 
     // Some differ, do a full merge.
-    SortedSet<DirectoryTree.FileNode> files = Sets.newTreeSet();
+    SortedSet<DirectoryTree.FileNode> files = new TreeSet<>();
     for (MerkleTree merkleTree : merkleTrees) {
       files.addAll(merkleTree.getFiles());
     }
-    SortedSet<DirectoryTree.SymlinkNode> symlinks = Sets.newTreeSet();
+    SortedSet<DirectoryTree.SymlinkNode> symlinks = new TreeSet<>();
     for (MerkleTree merkleTree : merkleTrees) {
       symlinks.addAll(merkleTree.getSymlinks());
     }
diff --git a/src/test/java/com/google/devtools/build/lib/remote/merkletree/DirectoryTreeTest.java b/src/test/java/com/google/devtools/build/lib/remote/merkletree/DirectoryTreeTest.java
@@ -33,6 +33,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.SortedSet;
 import java.util.stream.Collectors;
 import org.junit.Before;
 import org.junit.Test;
@@ -121,9 +122,9 @@ static void assertLexicographicalOrder(DirectoryTree tree) {
     // Assert the lexicographical order as defined by the remote execution protocol
     tree.visit(
         (PathFragment dirname,
-            List<FileNode> files,
-            List<SymlinkNode> symlinks,
-            List<DirectoryNode> dirs) -> {
+            SortedSet<FileNode> files,
+            SortedSet<SymlinkNode> symlinks,
+            SortedSet<DirectoryNode> dirs) -> {
           assertThat(files).isInStrictOrder();
           assertThat(dirs).isInStrictOrder();
         });
@@ -141,9 +142,9 @@ private static List<DirectoryNode> directoryNodesAtDepth(DirectoryTree tree, int
     List<DirectoryNode> directoryNodes = new ArrayList<>();
     tree.visit(
         (PathFragment dirname,
-            List<FileNode> files,
-            List<SymlinkNode> symlinks,
-            List<DirectoryNode> dirs) -> {
+            SortedSet<FileNode> files,
+            SortedSet<SymlinkNode> symlinks,
+            SortedSet<DirectoryNode> dirs) -> {
           int currDepth = dirname.segmentCount();
           if (currDepth == depth) {
             directoryNodes.addAll(dirs);
@@ -156,9 +157,9 @@ static List<FileNode> fileNodesAtDepth(DirectoryTree tree, int depth) {
     List<FileNode> fileNodes = new ArrayList<>();
     tree.visit(
         (PathFragment dirname,
-            List<FileNode> files,
-            List<SymlinkNode> symlinks,
-            List<DirectoryNode> dirs) -> {
+            SortedSet<FileNode> files,
+            SortedSet<SymlinkNode> symlinks,
+            SortedSet<DirectoryNode> dirs) -> {
           int currDepth = dirname.segmentCount();
           if (currDepth == depth) {
             fileNodes.addAll(files);
