[WGSL] Only visit functions and structures once

https://bugs.webkit.org/show_bug.cgi?id=260144
rdar://113854125

Reviewed by Dan Glastonbury.

Originally, the type checker and name mangling pass did two passes over global
declarations. That was necessary since declarations can refer to other declarations
that only appear later in the program. However, in 266862@main we introduced a new
pass that sorts all declarations, guaranteeing that by the time we visit a given
declaration, any other declaration it depends on will have already been visited.
This allows simplifying the code, and doing a single pass over all declarations.

* Source/WebGPU/WGSL/MangleNames.cpp:
(WGSL::NameManglerVisitor::run):
(WGSL::NameManglerVisitor::visit):
(WGSL::NameManglerVisitor::visitFunctionBody): Deleted.
* Source/WebGPU/WGSL/TypeCheck.cpp:
(WGSL::TypeChecker::check):
(WGSL::TypeChecker::visit):
(WGSL::TypeChecker::visitStructMembers): Deleted.
(WGSL::TypeChecker::visitFunctionBody): Deleted.
* Source/WebGPU/WGSL/TypeStore.cpp:
(WGSL::TypeStore::structType):
* Source/WebGPU/WGSL/TypeStore.h:
(WGSL::TypeStore::structType):

Canonical link: https://commits.webkit.org/266900@main

Source/WebGPU/WGSL/MangleNames.cpp

@@ -93,7 +93,6 @@ class NameManglerVisitor : public AST::Visitor, public ContextProvider<MangledNa
     MangledName makeMangledName(const String&, MangledName::Kind);
 
     void visitVariableDeclaration(AST::Variable&, MangledName::Kind);
-    void visitFunctionBody(AST::Function&);
 
     const CallGraph& m_callGraph;
     PrepareResult& m_result;
@@ -104,30 +103,23 @@ class NameManglerVisitor : public AST::Visitor, public ContextProvider<MangledNa
 void NameManglerVisitor::run()
 {
     auto& module = m_callGraph.ast();
+    for (auto& structure : module.structures())
+        visit(structure);
+
+    for (auto& variable : module.variables())
+        visit(variable);
+
     for (auto& function : module.functions()) {
         String originalName = function.name();
         introduceVariable(function.name(), MangledName::Function);
         auto it = m_result.entryPoints.find(originalName);
         if (it != m_result.entryPoints.end())
             it->value.mangledName = function.name();
+        visit(function);
     }
-
-    for (auto& structure : module.structures())
-        visit(structure);
-
-    for (auto& variable : module.variables())
-        visit(variable);
-
-    for (auto& function : module.functions())
-        visitFunctionBody(function);
 }
 
 void NameManglerVisitor::visit(AST::Function& function)
-{
-    introduceVariable(function.name(), MangledName::Function);
-}
-
-void NameManglerVisitor::visitFunctionBody(AST::Function& function)
 {
     ContextScope functionScope(this);
 

Source/WebGPU/WGSL/TypeCheck.cpp

@@ -111,8 +111,6 @@ class TypeChecker : public AST::Visitor, public ContextProvider<Binding> {
         Global,
     };
 
-    void visitFunctionBody(AST::Function&);
-    void visitStructMembers(AST::Structure&);
     void visitVariable(AST::Variable&, VariableKind);
     const Type* vectorFieldAccess(const Types::Vector&, AST::FieldAccessExpression&);
     void visitAttributes(AST::Attribute::List&);
@@ -166,18 +164,12 @@ std::optional<FailedCheck> TypeChecker::check()
     for (auto& structure : m_shaderModule.structures())
         visit(structure);
 
-    for (auto& structure : m_shaderModule.structures())
-        visitStructMembers(structure);
-
     for (auto& variable : m_shaderModule.variables())
         visitVariable(variable, VariableKind::Global);
 
     for (auto& function : m_shaderModule.functions())
         visit(function);
 
-    for (auto& function : m_shaderModule.functions())
-        visitFunctionBody(function);
-
     if (shouldDumpInferredTypes) {
         for (auto& error : m_errors)
             dataLogLn(error);
@@ -186,7 +178,6 @@ std::optional<FailedCheck> TypeChecker::check()
     if (m_errors.isEmpty())
         return std::nullopt;
 
-
     // FIXME: add support for warnings
     Vector<Warning> warnings { };
     return FailedCheck { WTFMove(m_errors), WTFMove(warnings) };
@@ -195,33 +186,15 @@ std::optional<FailedCheck> TypeChecker::check()
 // Declarations
 void TypeChecker::visit(AST::Structure& structure)
 {
-    const Type* structType = m_types.structType(structure);
-    introduceType(structure.name(), structType);
-}
-
-void TypeChecker::visitStructMembers(AST::Structure& structure)
-{
-    auto* binding = readVariable(structure.name());
-    ASSERT(binding && binding->kind == Binding::Type);
-    auto* type = binding->type;
-    ASSERT(std::holds_alternative<Types::Struct>(*type));
-
-    // This is the only place we need to modify a type.
-    // Since struct fields can reference other structs declared later in the
-    // program, the creation of struct types is a 2-step process:
-    // - First, we create an empty struct type for all structs in the program,
-    //   and expose them in the global context
-    // - Then, in a second pass, we populate the structs' fields.
-    // This way, the type of all structs will be available at the time we populate
-    // struct fields, even the ones defiend later in the program.
-    auto& structType = std::get<Types::Struct>(*type);
-    auto& fields = const_cast<HashMap<String, const Type*>&>(structType.fields);
+    HashMap<String, const Type*> fields;
     for (auto& member : structure.members()) {
         visitAttributes(member.attributes());
         auto* memberType = resolve(member.type());
         auto result = fields.add(member.name().id(), memberType);
         ASSERT_UNUSED(result, result.isNewEntry);
     }
+    const Type* structType = m_types.structType(structure, WTFMove(fields));
+    introduceType(structure.name(), structType);
 }
 
 void TypeChecker::visit(AST::Variable& variable)
@@ -289,19 +262,15 @@ void TypeChecker::visit(AST::Function& function)
         result = resolve(*function.maybeReturnType());
     else
         result = m_types.voidType();
+
     const Type* functionType = m_types.functionType(WTFMove(parameters), result);
     introduceValue(function.name(), functionType);
-}
 
-void TypeChecker::visitFunctionBody(AST::Function& function)
-{
     ContextScope functionContext(this);
-
     for (auto& parameter : function.parameters()) {
         auto* parameterType = resolve(parameter.typeName());
         introduceValue(parameter.name(), parameterType);
     }
-
     AST::Visitor::visit(function.body());
 }
 

Source/WebGPU/WGSL/TypeStore.cpp

@@ -154,9 +154,9 @@ TypeStore::TypeStore()
     allocateConstructor(&TypeStore::textureType, AST::ParameterizedTypeName::Base::TextureStorage3d, Texture::Kind::TextureStorage3d);
 }
 
-const Type* TypeStore::structType(AST::Structure& structure)
+const Type* TypeStore::structType(AST::Structure& structure, HashMap<String, const Type*>&& fields)
 {
-    return allocateType<Struct>(structure);
+    return allocateType<Struct>(structure, fields);
 }
 
 const Type* TypeStore::constructType(AST::ParameterizedTypeName::Base base, const Type* elementType)

Source/WebGPU/WGSL/TypeStore.h

@@ -30,6 +30,7 @@
 #include <wtf/FixedVector.h>
 #include <wtf/HashMap.h>
 #include <wtf/Vector.h>
+#include <wtf/text/StringHash.h>
 
 namespace WGSL {
 
@@ -56,7 +57,7 @@ class TypeStore {
     const Type* samplerType() const { return m_sampler; }
     const Type* textureExternalType() const { return m_textureExternal; }
 
-    const Type* structType(AST::Structure&);
+    const Type* structType(AST::Structure&, HashMap<String, const Type*>&& = { });
     const Type* arrayType(const Type*, std::optional<unsigned>);
     const Type* vectorType(const Type*, uint8_t);
     const Type* matrixType(const Type*, uint8_t columns, uint8_t rows);