Generate bazel trees in LocalCasReader

...and use this functionality in ExecutionServer

4 files changed, 81 insertions, 120 deletions

diff --git a/src/buildtool/execution_api/execution_service/TARGETS b/src/buildtool/execution_api/execution_service/TARGETS
index 7b3ecdb6..8d96545f 100644
--- a/src/buildtool/execution_api/execution_service/TARGETS
+++ b/src/buildtool/execution_api/execution_service/TARGETS
@@ -21,6 +21,7 @@
     , ["src/buildtool/logging", "log_level"]
     , "operation_cache"
     , ["src/utils/cpp", "verify_hash"]
+    , ["src/buildtool/execution_api/local", "local"]
     ]
   , "private-ldflags":
     ["-pthread", "-Wl,--whole-archive,-lpthread,--no-whole-archive"]
diff --git a/src/buildtool/execution_api/execution_service/execution_server.cpp b/src/buildtool/execution_api/execution_service/execution_server.cpp
index c99cfa04..84a5e53f 100644
--- a/src/buildtool/execution_api/execution_service/execution_server.cpp
+++ b/src/buildtool/execution_api/execution_service/execution_server.cpp
@@ -17,12 +17,12 @@
 #include <algorithm>
 #include <fstream>
 #include <string>
-#include <unordered_map>
 #include <utility>
 
 #include "execution_server.hpp"
 #include "fmt/core.h"
 #include "src/buildtool/execution_api/execution_service/operation_cache.hpp"
+#include "src/buildtool/execution_api/local/local_cas_reader.hpp"
 #include "src/buildtool/file_system/file_system_manager.hpp"
 #include "src/buildtool/logging/log_level.hpp"
 #include "src/buildtool/storage/garbage_collector.hpp"
@@ -153,129 +153,16 @@ auto ExecutionServiceImpl::GetIExecutionAction(
     return {std::move(i_execution_action), std::nullopt};
 }
 
-static auto GetDirectoryFromDigest(::bazel_re::Digest const& digest,
-                                   Storage const& storage) noexcept
-    -> std::optional<::bazel_re::Directory> {
-    // determine directory path from digest
-    auto const& path = storage.CAS().BlobPath(digest, /*is_executable=*/false);
-    if (not path) {
-        return std::nullopt;
-    }
-
-    // read directory content from path
-    auto const& content = FileSystemManager::ReadFile(*path);
-    if (not content) {
-        return std::nullopt;
-    }
-
-    // parse directory content
-    ::bazel_re::Directory dir{};
-    if (not dir.ParseFromString(*content)) {
-        return std::nullopt;
-    }
-    return dir;
-}
-
-// NOLINTNEXTLINE(misc-no-recursion)
-static auto CollectChildDirectoriesRecursively(
-    ::bazel_re::Directory const& root,
-    Storage const& storage,
-    gsl::not_null<std::unordered_map<::bazel_re::Digest,
-                                     ::bazel_re::Directory>*> map) noexcept
-    -> bool {
-    return std::all_of(root.directories().begin(),
-                       root.directories().end(),
-                       // NOLINTNEXTLINE(misc-no-recursion)
-                       [&map, &storage](auto const& node) {
-                           if (map->find(node.digest()) != map->end()) {
-                               return true;
-                           }
-                           auto tmp_root =
-                               GetDirectoryFromDigest(node.digest(), storage);
-                           if (not tmp_root) {
-                               return false;
-                           }
-                           if (not CollectChildDirectoriesRecursively(
-                                   *tmp_root, storage, map)) {
-                               return false;
-                           }
-                           try {
-                               map->emplace(node.digest(), *tmp_root);
-                           } catch (...) {
-                               return false;
-                           }
-                           return true;
-                       });
-}
-
-static auto GetChildrenFromDirectory(::bazel_re::Directory const& root,
-                                     Storage const& storage) noexcept
-    -> std::optional<std::vector<::bazel_re::Directory>> {
-    // determine child directories
-    std::unordered_map<::bazel_re::Digest, ::bazel_re::Directory> map{};
-    if (not CollectChildDirectoriesRecursively(root, storage, &map)) {
-        return std::nullopt;
-    }
-
-    // extract digests from child directories
-    std::vector<::bazel_re::Digest> digests{};
-    digests.reserve(map.size());
-    std::transform(map.begin(),
-                   map.end(),
-                   std::back_inserter(digests),
-                   [](auto const& pair) { return pair.first; });
-
-    // sort digests
-    std::sort(digests.begin(),
-              digests.end(),
-              [](auto const& left, auto const& right) {
-                  return left.hash() < right.hash();
-              });
-
-    // extract directory messages
-    std::vector<::bazel_re::Directory> children{};
-    children.reserve(digests.size());
-    std::transform(digests.begin(),
-                   digests.end(),
-                   std::back_inserter(children),
-                   [&map](auto const& digest) { return map[digest]; });
-
-    return children;
-}
-
 static auto CreateTreeDigestFromDirectoryDigest(
     ::bazel_re::Digest const& dir_digest,
     Storage const& storage) noexcept -> std::optional<::bazel_re::Digest> {
-    // determine root directory message
-    auto root = GetDirectoryFromDigest(dir_digest, storage);
-    if (not root) {
-        return std::nullopt;
+    LocalCasReader reader(&storage.CAS());
+    if (auto tree = reader.MakeTree(ArtifactDigest{dir_digest})) {
+        // serialize and store tree message
+        auto content = tree->SerializeAsString();
+        return storage.CAS().StoreBlob(content, /*is_executable=*/false);
     }
-
-    // determine child directory messages
-    auto children = GetChildrenFromDirectory(*root, storage);
-    if (not children) {
-        return std::nullopt;
-    }
-
-    // create tree message
-    ::bazel_re::Tree tree{};
-    tree.set_allocated_root(
-        gsl::owner<::bazel_re::Directory*>{new ::bazel_re::Directory{*root}});
-    tree.mutable_children()->Reserve(gsl::narrow<int>((*children).size()));
-    std::copy((*children).begin(),
-              (*children).end(),
-              ::pb::back_inserter(tree.mutable_children()));
-
-    // serialize and store tree message
-    auto content = tree.SerializeAsString();
-    auto tree_digest =
-        storage.CAS().StoreBlob(content, /*is_executable=*/false);
-    if (not tree_digest) {
-        return std::nullopt;
-    }
-
-    return tree_digest;
+    return std::nullopt;
 }
 
 static auto AddOutputPaths(::bazel_re::ExecuteResponse* response,
diff --git a/src/buildtool/execution_api/local/local_cas_reader.cpp b/src/buildtool/execution_api/local/local_cas_reader.cpp
index fd7ede7f..56915fa9 100644
--- a/src/buildtool/execution_api/local/local_cas_reader.cpp
+++ b/src/buildtool/execution_api/local/local_cas_reader.cpp
@@ -14,8 +14,13 @@
 
 #include "src/buildtool/execution_api/local/local_cas_reader.hpp"
 
+#include <algorithm>
 #include <cstdio>
 #include <memory>
+#include <stack>
+#include <unordered_map>
+#include <utility>
+#include <vector>
 
 #include "gsl/gsl"
 #include "src/buildtool/crypto/hash_function.hpp"
@@ -25,6 +30,13 @@
 #include "src/buildtool/logging/logger.hpp"
 #include "src/utils/cpp/path.hpp"
 
+namespace {
+[[nodiscard]] auto AssembleTree(
+    bazel_re::Directory root,
+    std::unordered_map<ArtifactDigest, bazel_re::Directory> directories)
+    -> bazel_re::Tree;
+}  // namespace
+
 auto LocalCasReader::ReadDirectory(ArtifactDigest const& digest) const noexcept
     -> std::optional<bazel_re::Directory> {
     if (auto const path = cas_.TreePath(digest)) {
@@ -38,6 +50,38 @@ auto LocalCasReader::ReadDirectory(ArtifactDigest const& digest) const noexcept
     return std::nullopt;
 }
 
+auto LocalCasReader::MakeTree(ArtifactDigest const& root) const noexcept
+    -> std::optional<bazel_re::Tree> {
+    try {
+        std::unordered_map<ArtifactDigest, bazel_re::Directory> directories;
+
+        std::stack<ArtifactDigest> to_check;
+        to_check.push(root);
+        while (not to_check.empty()) {
+            auto current = to_check.top();
+            to_check.pop();
+
+            if (directories.contains(current)) {
+                continue;
+            }
+
+            auto read_dir = ReadDirectory(current);
+            if (not read_dir) {
+                return std::nullopt;
+            }
+            for (auto const& node : read_dir->directories()) {
+                to_check.push(ArtifactDigest{node.digest()});
+            }
+            directories.insert_or_assign(std::move(current),
+                                         *std::move(read_dir));
+        }
+        auto root_directory = directories.extract(root).mapped();
+        return AssembleTree(std::move(root_directory), std::move(directories));
+    } catch (...) {
+        return std::nullopt;
+    }
+}
+
 auto LocalCasReader::ReadGitTree(ArtifactDigest const& digest) const noexcept
     -> std::optional<GitRepo::tree_entries_t> {
     if (auto const path = cas_.TreePath(digest)) {
@@ -112,3 +156,29 @@ auto LocalCasReader::DumpRaw(std::filesystem::path const& path,
     }
     return true;
 }
+
+namespace {
+[[nodiscard]] auto AssembleTree(
+    bazel_re::Directory root,
+    std::unordered_map<ArtifactDigest, bazel_re::Directory> directories)
+    -> bazel_re::Tree {
+    using Pair = std::pair<ArtifactDigest const, bazel_re::Directory>;
+    std::vector<Pair*> sorted;
+    sorted.reserve(directories.size());
+    for (auto& p : directories) {
+        sorted.push_back(&p);
+    }
+    std::sort(sorted.begin(), sorted.end(), [](Pair const* l, Pair const* r) {
+        return l->first.hash() < r->first.hash();
+    });
+
+    ::bazel_re::Tree result{};
+    (*result.mutable_root()) = std::move(root);
+    result.mutable_children()->Reserve(gsl::narrow<int>(sorted.size()));
+    std::transform(sorted.begin(),
+                   sorted.end(),
+                   ::pb::back_inserter(result.mutable_children()),
+                   [](Pair* p) { return std::move(p->second); });
+    return result;
+}
+}  // namespace
diff --git a/src/buildtool/execution_api/local/local_cas_reader.hpp b/src/buildtool/execution_api/local/local_cas_reader.hpp
index 64a1a51e..b989ec09 100644
--- a/src/buildtool/execution_api/local/local_cas_reader.hpp
+++ b/src/buildtool/execution_api/local/local_cas_reader.hpp
@@ -38,6 +38,9 @@ class LocalCasReader final {
     [[nodiscard]] auto ReadDirectory(ArtifactDigest const& digest)
         const noexcept -> std::optional<bazel_re::Directory>;
 
+    [[nodiscard]] auto MakeTree(ArtifactDigest const& root) const noexcept
+        -> std::optional<bazel_re::Tree>;
+
     [[nodiscard]] auto ReadGitTree(ArtifactDigest const& digest) const noexcept
         -> std::optional<GitRepo::tree_entries_t>;