[clangd] Introduce MemoryTrees

A structure that can be used to represent memory usage of a nested
set of systems.

Differential Revision: https://reviews.llvm.org/D88411

clang-tools-extra/clangd/support/CMakeLists.txt

@@ -21,6 +21,7 @@ add_clang_library(clangdSupport
   Context.cpp
   Logger.cpp
   Markup.cpp
+  MemoryTree.cpp
   Shutdown.cpp
   Threading.cpp
   ThreadsafeFS.cpp

clang-tools-extra/clangd/support/MemoryTree.cpp

@@ -0,0 +1,26 @@
+#include "support/MemoryTree.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include <cstddef>
+
+namespace clang {
+namespace clangd {
+
+MemoryTree &MemoryTree::createChild(llvm::StringRef Name) {
+  auto &Child = Children.try_emplace(Name, DetailAlloc).first->getSecond();
+  return Child;
+}
+
+const llvm::DenseMap<llvm::StringRef, MemoryTree> &
+MemoryTree::children() const {
+  return Children;
+}
+
+size_t MemoryTree::total() const {
+  size_t Total = Size;
+  for (const auto &Entry : Children)
+    Total += Entry.getSecond().total();
+  return Total;
+}
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/support/MemoryTree.h

@@ -0,0 +1,86 @@
+//===--- MemoryTree.h - A special tree for components and sizes -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_SUPPORT_MEMORYTREE_H_
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_SUPPORT_MEMORYTREE_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/StringSaver.h"
+#include <cstddef>
+#include <string>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+
+/// A tree that can be used to represent memory usage of nested components while
+/// preserving the hierarchy.
+/// Edges have associated names. An edge that might not be interesting to all
+/// traversers or costly to copy (e.g. file names) can be marked as "detail".
+/// Tree construction allows chosing between a detailed and brief mode, in brief
+/// mode all "detail" edges are ignored and tree is constructed without any
+/// string copies.
+struct MemoryTree {
+public:
+  /// If Alloc is nullptr, tree is in brief mode and will ignore detail edges.
+  MemoryTree(llvm::BumpPtrAllocator *DetailAlloc = nullptr)
+      : DetailAlloc(DetailAlloc) {}
+
+  /// No copy of the \p Name.
+  /// Note that returned pointers are invalidated with subsequent calls to
+  /// child/detail.
+  MemoryTree &child(llvm::StringLiteral Name) { return createChild(Name); }
+
+  MemoryTree(const MemoryTree &) = delete;
+  MemoryTree &operator=(const MemoryTree &) = delete;
+
+  MemoryTree(MemoryTree &&) = default;
+  MemoryTree &operator=(MemoryTree &&) = default;
+
+  /// Makes a copy of the \p Name in detailed mode, returns current node
+  /// otherwise.
+  /// Note that returned pointers are invalidated with subsequent calls to
+  /// child/detail.
+  MemoryTree &detail(llvm::StringRef Name) {
+    return DetailAlloc ? createChild(Name.copy(*DetailAlloc)) : *this;
+  }
+
+  /// Increases size of current node by \p Increment.
+  void addUsage(size_t Increment) { Size += Increment; }
+
+  /// Returns edges to direct children of this node.
+  const llvm::DenseMap<llvm::StringRef, MemoryTree> &children() const;
+
+  /// Returns total number of bytes used by this sub-tree. Performs a traversal.
+  size_t total() const;
+
+  /// Returns total number of bytes used by this node only.
+  size_t self() const { return Size; }
+
+private:
+  /// Adds a child with an edge labeled as \p Name. Multiple calls to this
+  /// function returns the same node.
+  MemoryTree &createChild(llvm::StringRef Name);
+
+  /// Allocator to use for detailed edge names.
+  llvm::BumpPtrAllocator *DetailAlloc = nullptr;
+
+  /// Bytes owned by this component specifically.
+  size_t Size = 0;
+
+  /// Edges from current node to its children. Keys are the labels for edges.
+  llvm::DenseMap<llvm::StringRef, MemoryTree> Children;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clang-tools-extra/clangd/unittests/CMakeLists.txt

@@ -99,6 +99,7 @@ add_unittest(ClangdUnitTests ClangdTests
   support/ContextTests.cpp
   support/FunctionTests.cpp
   support/MarkupTests.cpp
+  support/MemoryTreeTests.cpp
   support/ThreadingTests.cpp
   support/TestTracer.cpp
   support/TraceTests.cpp

clang-tools-extra/clangd/unittests/support/MemoryTreeTests.cpp

@@ -0,0 +1,78 @@
+//===-- MemoryTreeTests.cpp -------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "support/MemoryTree.h"
+#include "llvm/Support/Allocator.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include <ostream>
+
+namespace clang {
+namespace clangd {
+namespace {
+using testing::Contains;
+using testing::IsEmpty;
+using testing::UnorderedElementsAre;
+
+MATCHER_P2(WithNameAndSize, Name, Size, "") {
+  return arg.first == Name &&
+         arg.getSecond().total() == static_cast<size_t>(Size);
+}
+
+TEST(MemoryTree, Basics) {
+  MemoryTree MT;
+  EXPECT_EQ(MT.total(), 0U);
+  EXPECT_THAT(MT.children(), IsEmpty());
+
+  MT.addUsage(42);
+  EXPECT_EQ(MT.total(), 42U);
+  EXPECT_THAT(MT.children(), IsEmpty());
+
+  MT.child("leaf").addUsage(1);
+  EXPECT_EQ(MT.total(), 43U);
+  EXPECT_THAT(MT.children(), UnorderedElementsAre(WithNameAndSize("leaf", 1)));
+
+  // child should be idempotent.
+  MT.child("leaf").addUsage(1);
+  EXPECT_EQ(MT.total(), 44U);
+  EXPECT_THAT(MT.children(), UnorderedElementsAre(WithNameAndSize("leaf", 2)));
+}
+
+TEST(MemoryTree, DetailedNodesWithoutDetails) {
+  MemoryTree MT;
+  MT.detail("should_be_ignored").addUsage(2);
+  EXPECT_THAT(MT.children(), IsEmpty());
+  EXPECT_EQ(MT.total(), 2U);
+
+  // Make sure children from details are merged.
+  MT.detail("first_detail").child("leaf").addUsage(1);
+  MT.detail("second_detail").child("leaf").addUsage(1);
+  EXPECT_THAT(MT.children(), Contains(WithNameAndSize("leaf", 2)));
+}
+
+TEST(MemoryTree, DetailedNodesWithDetails) {
+  llvm::BumpPtrAllocator Alloc;
+  MemoryTree MT(&Alloc);
+
+  {
+    auto &Detail = MT.detail("first_detail");
+    Detail.child("leaf").addUsage(1);
+    EXPECT_THAT(MT.children(), Contains(WithNameAndSize("first_detail", 1)));
+    EXPECT_THAT(Detail.children(), Contains(WithNameAndSize("leaf", 1)));
+  }
+
+  {
+    auto &Detail = MT.detail("second_detail");
+    Detail.child("leaf").addUsage(1);
+    EXPECT_THAT(MT.children(), Contains(WithNameAndSize("second_detail", 1)));
+    EXPECT_THAT(Detail.children(), Contains(WithNameAndSize("leaf", 1)));
+  }
+}
+} // namespace
+} // namespace clangd
+} // namespace clang