From af47388345cc9460a7fab4d1f412a3148409f021 Mon Sep 17 00:00:00 2001
From: Ben Clayton <bclayton@google.com>
Date: Thu, 21 Jul 2022 23:46:15 +0000
Subject: [PATCH] tint/resolver: Optimize node marking

Use the new utils::Bitset to replace the much slower unordered_map.

Bug: tint:1613
Change-Id: I80503d2b897f754408c4ce14a221210d6abc18f7
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/96403
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: dan sinclair <dsinclair@google.com>
---
 src/tint/resolver/resolver.cc | 9 +++++++--
 src/tint/resolver/resolver.h  | 3 ++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/tint/resolver/resolver.cc b/src/tint/resolver/resolver.cc
index db9851cf52..b418deb485 100644
--- a/src/tint/resolver/resolver.cc
+++ b/src/tint/resolver/resolver.cc
@@ -105,6 +105,9 @@ bool Resolver::Resolve() {
 
     builder_->Sem().Reserve(builder_->LastAllocatedNodeID());
 
+    // Pre-allocate the marked bitset with the total number of AST nodes.
+    marked_.Resize(builder_->ASTNodes().Count());
+
     if (!DependencyGraph::Build(builder_->AST(), builder_->Symbols(), builder_->Diagnostics(),
                                 dependencies_)) {
         return false;
@@ -161,7 +164,7 @@ bool Resolver::ResolveInternal() {
 
     bool result = true;
     for (auto* node : builder_->ASTNodes().Objects()) {
-        if (marked_.count(node) == 0) {
+        if (!marked_[node->node_id.value]) {
             TINT_ICE(Resolver, diagnostics_)
                 << "AST node '" << node->TypeInfo().name << "' was not reached by the resolver\n"
                 << "At: " << node->source << "\n"
@@ -2842,7 +2845,9 @@ bool Resolver::Mark(const ast::Node* node) {
         TINT_ICE(Resolver, diagnostics_) << "Resolver::Mark() called with nullptr";
         return false;
     }
-    if (marked_.emplace(node).second) {
+    auto marked_bit_ref = marked_[node->node_id.value];
+    if (!marked_bit_ref) {
+        marked_bit_ref = true;
         return true;
     }
     TINT_ICE(Resolver, diagnostics_) << "AST node '" << node->TypeInfo().name
diff --git a/src/tint/resolver/resolver.h b/src/tint/resolver/resolver.h
index f3aa38b145..4a1c250289 100644
--- a/src/tint/resolver/resolver.h
+++ b/src/tint/resolver/resolver.h
@@ -35,6 +35,7 @@
 #include "src/tint/sem/constant.h"
 #include "src/tint/sem/function.h"
 #include "src/tint/sem/struct.h"
+#include "src/tint/utils/bitset.h"
 #include "src/tint/utils/unique_vector.h"
 
 // Forward declarations
@@ -420,7 +421,7 @@ class Resolver {
     ast::Extensions enabled_extensions_;
     std::vector<sem::Function*> entry_points_;
     std::unordered_map<const sem::Type*, const Source&> atomic_composite_info_;
-    std::unordered_set<const ast::Node*> marked_;
+    utils::Bitset<0> marked_;
     std::unordered_map<uint32_t, const sem::Variable*> constant_ids_;
     std::unordered_map<ArrayConstructorSig, sem::CallTarget*> array_ctors_;
     std::unordered_map<StructConstructorSig, sem::CallTarget*> struct_ctors_;