[llvm][mustache] Use BumpPtrAllocator to save ASTNodes (#159194)

We make the Mustache ASTNodes usable in the arena by first removing all
of the memory owning data structures, like std::vector, std::unique_ptr,
and SmallVector. We use standard LLVM list types to hold this data
instead, and make use of a UniqueStringSaver to hold the various
templates strings.

Additionally, update clang-doc APIs to use the new interfaces.

Future work can make better use of Twine interfaces to help avoid any
intermediate copies or allocations.

Author: ilovepi

clang-tools-extra/clang-doc/HTMLMustacheGenerator.cpp

@@ -46,18 +46,22 @@ class MustacheHTMLGenerator : public Generator {
                            const ClangDocContext &CDCtx) override;
 };
 
-class MustacheTemplateFile : public Template {
+class MustacheTemplateFile {
+  BumpPtrAllocator Allocator;
+  StringSaver Saver;
+  MustacheContext Ctx;
+  Template T;
+  std::unique_ptr<MemoryBuffer> Buffer;
+
 public:
   static Expected<std::unique_ptr<MustacheTemplateFile>>
   createMustacheFile(StringRef FileName) {
     ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrError =
         MemoryBuffer::getFile(FileName);
     if (auto EC = BufferOrError.getError())
       return createFileOpenError(FileName, EC);
-
-    std::unique_ptr<MemoryBuffer> Buffer = std::move(BufferOrError.get());
-    StringRef FileContent = Buffer->getBuffer();
-    return std::make_unique<MustacheTemplateFile>(FileContent);
+    return std::make_unique<MustacheTemplateFile>(
+        std::move(BufferOrError.get()));
   }
 
   Error registerPartialFile(StringRef Name, StringRef FileName) {
@@ -68,11 +72,15 @@ class MustacheTemplateFile : public Template {
 
     std::unique_ptr<MemoryBuffer> Buffer = std::move(BufferOrError.get());
     StringRef FileContent = Buffer->getBuffer();
-    registerPartial(Name.str(), FileContent.str());
+    T.registerPartial(Name.str(), FileContent.str());
     return Error::success();
   }
 
-  MustacheTemplateFile(StringRef TemplateStr) : Template(TemplateStr) {}
+  void render(json::Value &V, raw_ostream &OS) { T.render(V, OS); }
+
+  MustacheTemplateFile(std::unique_ptr<MemoryBuffer> &&B)
+      : Saver(Allocator), Ctx(Allocator, Saver), T(B->getBuffer(), Ctx),
+        Buffer(std::move(B)) {}
 };
 
 static std::unique_ptr<MustacheTemplateFile> NamespaceTemplate = nullptr;

llvm/benchmarks/Mustache.cpp

@@ -8,7 +8,7 @@
 static const std::string LongHtmlString = [] {
   std::string S;
   S.reserve(500000);
-  for (int i = 0; i < 50000; ++i) {
+  for (int Idx = 0; Idx < 50000; ++Idx) {
     S += "<script>alert('xss');</script>";
   }
   return S;
@@ -153,7 +153,11 @@ static const std::string LargeOutputStringTemplate = "{{long_string}}";
 // syntaxes.
 static void BM_Mustache_StringRendering(benchmark::State &state,
                                         const std::string &TplStr) {
-  llvm::mustache::Template Tpl(TplStr);
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(TplStr, Ctx);
   llvm::json::Value Data =
       llvm::json::Object({{"content", llvm::json::Value(LongHtmlString)}});
   for (auto _ : state) {
@@ -172,7 +176,11 @@ BENCHMARK_CAPTURE(BM_Mustache_StringRendering, Unescaped_Ampersand,
 // Tests the "hot render" cost of repeatedly traversing a deep and wide
 // JSON object.
 static void BM_Mustache_DeepTraversal(benchmark::State &state) {
-  llvm::mustache::Template Tpl(DeepTraversalTemplate);
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(DeepTraversalTemplate, Ctx);
   for (auto _ : state) {
     std::string Result;
     llvm::raw_string_ostream OS(Result);
@@ -184,7 +192,12 @@ BENCHMARK(BM_Mustache_DeepTraversal);
 
 // Tests the "hot render" cost of pushing and popping a deep context stack.
 static void BM_Mustache_DeeplyNestedRendering(benchmark::State &state) {
-  llvm::mustache::Template Tpl(DeeplyNestedRenderingTemplate);
+
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(DeeplyNestedRenderingTemplate, Ctx);
   for (auto _ : state) {
     std::string Result;
     llvm::raw_string_ostream OS(Result);
@@ -197,7 +210,11 @@ BENCHMARK(BM_Mustache_DeeplyNestedRendering);
 // Tests the performance of the loop logic when iterating over a huge number of
 // items.
 static void BM_Mustache_HugeArrayIteration(benchmark::State &state) {
-  llvm::mustache::Template Tpl(HugeArrayIterationTemplate);
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(HugeArrayIterationTemplate, Ctx);
   for (auto _ : state) {
     std::string Result;
     llvm::raw_string_ostream OS(Result);
@@ -209,25 +226,37 @@ BENCHMARK(BM_Mustache_HugeArrayIteration);
 
 // Tests the performance of the parser on a large, "wide" template.
 static void BM_Mustache_ComplexTemplateParsing(benchmark::State &state) {
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
   for (auto _ : state) {
-    llvm::mustache::Template Tpl(ComplexTemplateParsingTemplate);
+    llvm::mustache::Template Tpl(ComplexTemplateParsingTemplate, Ctx);
     benchmark::DoNotOptimize(Tpl);
   }
 }
 BENCHMARK(BM_Mustache_ComplexTemplateParsing);
 
 // Tests the performance of the parser on a small, "deep" template.
 static void BM_Mustache_SmallTemplateParsing(benchmark::State &state) {
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
   for (auto _ : state) {
-    llvm::mustache::Template Tpl(SmallTemplateParsingTemplate);
+    llvm::mustache::Template Tpl(SmallTemplateParsingTemplate, Ctx);
     benchmark::DoNotOptimize(Tpl);
   }
 }
 BENCHMARK(BM_Mustache_SmallTemplateParsing);
 
 // Tests the performance of rendering a template that includes a partial.
 static void BM_Mustache_PartialsRendering(benchmark::State &state) {
-  llvm::mustache::Template Tpl(ComplexPartialTemplate);
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(ComplexPartialTemplate, Ctx);
   Tpl.registerPartial("item_partial", ItemPartialTemplate);
   llvm::json::Value Data = HugeArrayData;
 
@@ -243,7 +272,11 @@ BENCHMARK(BM_Mustache_PartialsRendering);
 // Tests the performance of the underlying buffer management when generating a
 // very large output.
 static void BM_Mustache_LargeOutputString(benchmark::State &state) {
-  llvm::mustache::Template Tpl(LargeOutputStringTemplate);
+  llvm::BumpPtrAllocator Allocator;
+  llvm::StringSaver Saver(Allocator);
+  llvm::mustache::MustacheContext Ctx(Allocator, Saver);
+
+  llvm::mustache::Template Tpl(LargeOutputStringTemplate, Ctx);
   for (auto _ : state) {
     std::string Result;
     llvm::raw_string_ostream OS(Result);

llvm/include/llvm/Support/Mustache.h

@@ -71,6 +71,8 @@
 
 #include "Error.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/ilist.h"
+#include "llvm/ADT/ilist_node.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/JSON.h"
@@ -84,10 +86,15 @@ using Lambda = std::function<llvm::json::Value()>;
 using SectionLambda = std::function<llvm::json::Value(std::string)>;
 
 class ASTNode;
-using AstPtr = std::unique_ptr<ASTNode>;
+using AstPtr = ASTNode *;
 using EscapeMap = DenseMap<char, std::string>;
+using ASTNodeList = iplist<ASTNode>;
 
 struct MustacheContext {
+  MustacheContext(BumpPtrAllocator &Allocator, StringSaver &Saver)
+      : Allocator(Allocator), Saver(Saver) {}
+  BumpPtrAllocator &Allocator;
+  StringSaver &Saver;
   StringMap<AstPtr> Partials;
   StringMap<Lambda> Lambdas;
   StringMap<SectionLambda> SectionLambdas;
@@ -98,7 +105,7 @@ struct MustacheContext {
 // and Lambdas that are registered with it.
 class Template {
 public:
-  LLVM_ABI Template(StringRef TemplateStr);
+  LLVM_ABI Template(StringRef TemplateStr, MustacheContext &Ctx);
 
   Template(const Template &) = delete;
 
@@ -110,7 +117,7 @@ class Template {
   // type.
   LLVM_ABI ~Template();
 
-  LLVM_ABI Template &operator=(Template &&Other) noexcept;
+  Template &operator=(Template &&) = delete;
 
   LLVM_ABI void render(const llvm::json::Value &Data, llvm::raw_ostream &OS);
 
@@ -126,7 +133,7 @@ class Template {
   LLVM_ABI void overrideEscapeCharacters(DenseMap<char, std::string> Escapes);
 
 private:
-  MustacheContext Ctx;
+  MustacheContext &Ctx;
   AstPtr Tree;
 };
 } // namespace llvm::mustache