CVE-2024-40776

diff --git a/Source/WebGPU/WGSL/AST/ASTCallExpression.h b/Source/WebGPU/WGSL/AST/ASTCallExpression.h
index 7abadd2d3217..0a85ce277959 100644
--- a/Source/WebGPU/WGSL/AST/ASTCallExpression.h
+++ b/Source/WebGPU/WGSL/AST/ASTCallExpression.h
@@ -28,6 +28,7 @@
 #include "ASTExpression.h"
 
 namespace WGSL {
+class RewriteGlobalVariables;
 class TypeChecker;
 
 namespace AST {
@@ -39,6 +40,7 @@ namespace AST {
 class CallExpression final : public Expression {
     WGSL_AST_BUILDER_NODE(CallExpression);
 
+    friend RewriteGlobalVariables;
     friend TypeChecker;
 
 public:
diff --git a/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp b/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp
index edf73a1f6ce9..4dd9180b9f04 100644
--- a/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp
+++ b/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp
@@ -114,13 +114,12 @@ class RewriteGlobalVariables : public AST::Visitor {
     AST::Expression& bufferLengthType();
     AST::Expression& bufferLengthReferenceType();
 
-    // atomics
-    void initializeAtomics(AST::Function&, const UsedPrivateGlobals&, size_t);
+    // zero initialization
+    void initializeVariables(AST::Function&, const UsedPrivateGlobals&, size_t);
     void insertWorkgroupBarrier(AST::Function&, size_t);
     AST::Identifier& findOrInsertLocalInvocationIndex(AST::Function&);
-    AST::Statement::List atomicStoreInitialValue(const UsedPrivateGlobals&);
-    void atomicStoreInitialValue(AST::Expression&, AST::Statement::List&, unsigned);
-    bool containsAtomic(const Type*);
+    AST::Statement::List storeInitialValue(const UsedPrivateGlobals&);
+    void storeInitialValue(AST::Expression&, AST::Statement::List&, unsigned, bool isNested);
 
     void packResource(AST::Variable&);
     void packArrayResource(AST::Variable&, const Types::Array*);
@@ -1415,12 +1414,12 @@ void RewriteGlobalVariables::insertLocalDefinitions(AST::Function& function, con
     }
 
     auto offset = function.body().statements().size() - initialBodySize;
-    initializeAtomics(function, usedPrivateGlobals, offset);
+    initializeVariables(function, usedPrivateGlobals, offset);
 }
 
-void RewriteGlobalVariables::initializeAtomics(AST::Function& function, const UsedPrivateGlobals& globals, size_t offset)
+void RewriteGlobalVariables::initializeVariables(AST::Function& function, const UsedPrivateGlobals& globals, size_t offset)
 {
-    auto initializations = atomicStoreInitialValue(globals);
+    auto initializations = storeInitialValue(globals);
     if (initializations.isEmpty())
         return;
 
@@ -1511,43 +1510,53 @@ AST::Identifier& RewriteGlobalVariables::findOrInsertLocalInvocationIndex(AST::F
     return parameter.name();
 }
 
-AST::Statement::List RewriteGlobalVariables::atomicStoreInitialValue(const UsedPrivateGlobals& globals)
+AST::Statement::List RewriteGlobalVariables::storeInitialValue(const UsedPrivateGlobals& globals)
 {
     AST::Statement::List statements;
     for (auto* global : globals) {
         auto& variable = *global->declaration;
-        auto* type = variable.storeType();
-        if (!containsAtomic(type))
+
+        if (auto addressSpace = variable.addressSpace(); !addressSpace.has_value() || *addressSpace != AddressSpace::Workgroup)
             continue;
 
+        auto* type = variable.storeType();
         auto& target = m_callGraph.ast().astBuilder().construct<AST::IdentifierExpression>(
             SourceSpan::empty(),
             AST::Identifier::make(variable.name().id())
         );
         target.m_inferredType = type;
-        atomicStoreInitialValue(target, statements, 0);
+        storeInitialValue(target, statements, 0, false);
     }
     return statements;
 }
 
-bool RewriteGlobalVariables::containsAtomic(const Type* type)
+void RewriteGlobalVariables::storeInitialValue(AST::Expression& target, AST::Statement::List& statements, unsigned arrayDepth, bool isNested)
 {
-    if (std::holds_alternative<Types::Atomic>(*type))
-        return true;
-    if (auto* arrayType = std::get_if<Types::Array>(type))
-        return containsAtomic(arrayType->element);
-    if (auto* structType = std::get_if<Types::Struct>(type)) {
-        for (const auto& [_, fieldType] : structType->fields) {
-            if (containsAtomic(fieldType))
-                return true;
-        }
-        return false;
-    }