tint/sem: Move variable decls to CompoundStatement

Variables can be declared in more than just BlockStatement.
For example, for-loops can declare a variable.

Change this to be a map instead of a vector. This helps with lookups.

Change-Id: Ic9429425af70e9535c21cc0875b875f145724266
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/104040
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/tint/resolver/resolver.cc

@@ -2268,21 +2268,16 @@ sem::Expression* Resolver::Identifier(const ast::IdentifierExpression* expr) {
                     current_statement_->FindFirstParent<sem::LoopContinuingBlockStatement>()) {
                 auto* loop_block = continuing_block->FindFirstParent<sem::LoopBlockStatement>();
                 if (loop_block->FirstContinue()) {
-                    auto& decls = loop_block->Decls();
                     // If our identifier is in loop_block->decls, make sure its index is
                     // less than first_continue
-                    auto iter = std::find_if(decls.begin(), decls.end(),
-                                             [&symbol](auto* v) { return v->symbol == symbol; });
-                    if (iter != decls.end()) {
-                        auto var_decl_index =
-                            static_cast<size_t>(std::distance(decls.begin(), iter));
-                        if (var_decl_index >= loop_block->NumDeclsAtFirstContinue()) {
+                    if (auto* decl = loop_block->Decls().Find(symbol)) {
+                        if (decl->order >= loop_block->NumDeclsAtFirstContinue()) {
                             AddError("continue statement bypasses declaration of '" +
                                          builder_->Symbols().NameFor(symbol) + "'",
                                      loop_block->FirstContinue()->source);
                             AddNote("identifier '" + builder_->Symbols().NameFor(symbol) +
                                         "' declared here",
-                                    (*iter)->source);
+                                    decl->variable->Declaration()->source);
                             AddNote("identifier '" + builder_->Symbols().NameFor(symbol) +
                                         "' referenced in continuing block here",
                                     expr->source);
@@ -3128,9 +3123,7 @@ sem::Statement* Resolver::VariableDeclStatement(const ast::VariableDeclStatement
             }
         }
 
-        if (current_block_) {  // Not all statements are inside a block
-            current_block_->AddDecl(stmt->variable);
-        }
+        current_compound_statement_->AddDecl(variable->As<sem::LocalVariable>());
 
         if (auto* ctor = variable->Constructor()) {
             sem->Behaviors() = ctor->Behaviors();
@@ -3232,7 +3225,7 @@ sem::Statement* Resolver::ContinueStatement(const ast::ContinueStatement* stmt)
         if (auto* block = sem->FindFirstParent<sem::LoopBlockStatement>()) {
             if (!block->FirstContinue()) {
                 const_cast<sem::LoopBlockStatement*>(block)->SetFirstContinue(
-                    stmt, block->Decls().size());
+                    stmt, block->Decls().Count());
             }
         }
 
@@ -3337,12 +3330,10 @@ SEM* Resolver::StatementScope(const ast::Statement* ast, SEM* sem, F&& callback)
     builder_->Sem().Add(ast, sem);
 
     auto* as_compound = As<sem::CompoundStatement, CastFlags::kDontErrorOnImpossibleCast>(sem);
-    auto* as_block = As<sem::BlockStatement, CastFlags::kDontErrorOnImpossibleCast>(sem);
 
     TINT_SCOPED_ASSIGNMENT(current_statement_, sem);
     TINT_SCOPED_ASSIGNMENT(current_compound_statement_,
                            as_compound ? as_compound : current_compound_statement_);
-    TINT_SCOPED_ASSIGNMENT(current_block_, as_block ? as_block : current_block_);
 
     if (!callback()) {
         return nullptr;

src/tint/resolver/resolver.h

@@ -375,11 +375,8 @@ class Resolver {
     /// * Assigns `sem` to #current_statement_
     /// * Assigns `sem` to #current_compound_statement_ if `sem` derives from
     ///   sem::CompoundStatement.
-    /// * Assigns `sem` to #current_block_ if `sem` derives from
-    ///   sem::BlockStatement.
     /// * Then calls `callback`.
-    /// * Before returning #current_statement_, #current_compound_statement_, and
-    ///   #current_block_ are restored to their original values.
+    /// * Before returning #current_statement_ and #current_compound_statement_ are restored to their original values.
     /// @returns `sem` if `callback` returns true, otherwise `nullptr`.
     template <typename SEM, typename F>
     SEM* StatementScope(const ast::Statement* ast, SEM* sem, F&& callback);
@@ -441,7 +438,6 @@ class Resolver {
     sem::Function* current_function_ = nullptr;
     sem::Statement* current_statement_ = nullptr;
     sem::CompoundStatement* current_compound_statement_ = nullptr;
-    sem::BlockStatement* current_block_ = nullptr;
 };
 
 }  // namespace tint::resolver

src/tint/resolver/uniformity.cc

@@ -481,8 +481,8 @@ class UniformityGraph {
                 }
 
                 // Remove any variables declared in this scope from the set of in-scope variables.
-                for (auto* d : sem_.Get<sem::BlockStatement>(b)->Decls()) {
-                    current_function_->local_var_decls.erase(sem_.Get<sem::LocalVariable>(d));
+                for (auto decl : sem_.Get<sem::BlockStatement>(b)->Decls()) {
+                    current_function_->local_var_decls.erase(decl.value.variable);
                 }
 
                 return cf;

src/tint/sem/block_statement.cc

@@ -35,10 +35,6 @@ const ast::BlockStatement* BlockStatement::Declaration() const {
     return Base::Declaration()->As<ast::BlockStatement>();
 }
 
-void BlockStatement::AddDecl(const ast::Variable* var) {
-    decls_.push_back(var);
-}
-
 FunctionBlockStatement::FunctionBlockStatement(const sem::Function* function)
     : Base(function->Declaration()->body, nullptr, function) {
     TINT_ASSERT(Semantic, function);

src/tint/sem/block_statement.h

@@ -47,16 +47,6 @@ class BlockStatement : public Castable<BlockStatement, CompoundStatement> {
     /// @returns the AST block statement associated with this semantic block
     /// statement
     const ast::BlockStatement* Declaration() const;
-
-    /// @returns the declarations associated with this block
-    const std::vector<const ast::Variable*>& Decls() const { return decls_; }
-
-    /// Associates a declaration with this block.
-    /// @param var a variable declaration to be added to the block
-    void AddDecl(const ast::Variable* var);
-
-  private:
-    std::vector<const ast::Variable*> decls_;
 };
 
 /// The root block statement for a function

src/tint/sem/statement.cc

@@ -17,8 +17,10 @@
 #include "src/tint/ast/block_statement.h"
 #include "src/tint/ast/loop_statement.h"
 #include "src/tint/ast/statement.h"
+#include "src/tint/ast/variable.h"
 #include "src/tint/sem/block_statement.h"
 #include "src/tint/sem/statement.h"
+#include "src/tint/sem/variable.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::sem::Statement);
 TINT_INSTANTIATE_TYPEINFO(tint::sem::CompoundStatement);
@@ -43,4 +45,8 @@ CompoundStatement::CompoundStatement(const ast::Statement* declaration,
 
 CompoundStatement::~CompoundStatement() = default;
 
+void CompoundStatement::AddDecl(const sem::LocalVariable* var) {
+    decls_.Add(var->Declaration()->symbol, OrderedLocalVariable{decls_.Count(), var});
+}
+
 }  // namespace tint::sem

src/tint/sem/statement.h

@@ -17,6 +17,8 @@
 
 #include "src/tint/sem/behavior.h"
 #include "src/tint/sem/node.h"
+#include "src/tint/symbol.h"
+#include "src/tint/utils/hashmap.h"
 
 // Forward declarations
 namespace tint::ast {
@@ -26,6 +28,7 @@ namespace tint::sem {
 class BlockStatement;
 class CompoundStatement;
 class Function;
+class LocalVariable;
 }  // namespace tint::sem
 
 namespace tint::sem {
@@ -128,6 +131,25 @@ class CompoundStatement : public Castable<Statement, Statement> {
 
     /// Destructor
     ~CompoundStatement() override;
+
+    /// OrderedLocalVariable describes a local variable declaration, and order of declaration.
+    struct OrderedLocalVariable {
+        /// The 0-based declaration order index of the variable
+        size_t order;
+        /// The variable
+        const LocalVariable* variable;
+    };
+
+    /// @returns a map of variable name to variable declarations associated with this block
+    const utils::Hashmap<Symbol, OrderedLocalVariable, 4>& Decls() const { return decls_; }
+
+    /// Associates a declaration with this block.
+    /// @note this method must be called in variable declaration order
+    /// @param var a variable declaration to be added to the block
+    void AddDecl(const LocalVariable* var);
+
+  private:
+    utils::Hashmap<Symbol, OrderedLocalVariable, 4> decls_;
 };
 
 template <typename Pred>