tint/uniformity: Retain control flow graphs

Move the control flow graphs to `FunctionInfo` so that they are not
destroyed when we finish processing the function.

This will make it easier to produce more detailed diagnostics for
uniformity issues, which will be coming in subsequent CLs.

There is no functional change in this CL.

Bug: tint:880
Change-Id: I1322fb54b16bd1c660799a62435fbdcd7fb39cb2
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/89822
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
diff --git a/src/tint/resolver/uniformity.cc b/src/tint/resolver/uniformity.cc
index 381a6da..cf3ba17 100644
--- a/src/tint/resolver/uniformity.cc
+++ b/src/tint/resolver/uniformity.cc
@@ -67,46 +67,16 @@
     ParameterNoRestriction,
 };
 
-/// ParameterInfo holds information about the uniformity requirements and effects for a particular
-/// function parameter.
-struct ParameterInfo {
-    /// The parameter's uniformity requirements.
-    ParameterTag tag = ParameterNoRestriction;
-    /// Will be `true` if this function may cause the contents of this pointer parameter to become
-    /// non-uniform.
-    bool pointer_may_become_non_uniform = false;
-    /// The parameters that are required to be uniform for the contents of this pointer parameter to
-    /// be uniform at function exit.
-    std::vector<const sem::Parameter*> pointer_param_output_sources;
-};
-
-/// FunctionInfo holds information about the uniformity requirements and effects for a particular
-/// function.
-struct FunctionInfo {
-    /// The call site uniformity requirements.
-    CallSiteTag callsite_tag;
-    /// The function's uniformity effects.
-    FunctionTag function_tag;
-    /// The uniformity requirements of the function's parameters.
-    std::vector<ParameterInfo> parameters;
-};
-
 /// Node represents a node in the graph of control flow and value nodes within the analysis of a
 /// single function.
 struct Node {
     /// Constructor
-    /// @param t the node tag (used for debugging)
-    Node([[maybe_unused]] std::string t, const ast::Node* a)
-        :
-#if TINT_DUMP_UNIFORMITY_GRAPH
-          tag(t),
-#endif
-          ast(a) {
-    }
+    /// @param a the corresponding AST node
+    explicit Node(const ast::Node* a) : ast(a) {}
 
 #if TINT_DUMP_UNIFORMITY_GRAPH
     /// The node tag.
-    const std::string tag;
+    std::string tag;
 #endif
 
     /// The corresponding AST node, or nullptr.
@@ -121,11 +91,149 @@
     /// The node that this node was visited from, or nullptr if not visited.
     Node* visited_from = nullptr;
 
-    /// Add a edge to the `to` node.
+    /// Add an edge to the `to` node.
     /// @param to the destination node
     void AddEdge(Node* to) { edges.add(to); }
 };
 
+/// ParameterInfo holds information about the uniformity requirements and effects for a particular
+/// function parameter.
+struct ParameterInfo {
+    /// The semantic node in corresponds to this parameter.
+    const sem::Parameter* sem;
+    /// The parameter's uniformity requirements.
+    ParameterTag tag = ParameterNoRestriction;
+    /// Will be `true` if this function may cause the contents of this pointer parameter to become
+    /// non-uniform.
+    bool pointer_may_become_non_uniform = false;
+    /// The parameters that are required to be uniform for the contents of this pointer parameter to
+    /// be uniform at function exit.
+    std::vector<const sem::Parameter*> pointer_param_output_sources;
+    /// The node in the graph that corresponds to this parameter's initial value.
+    Node* init_value;
+    /// The node in the graph that corresponds to this parameter's output value (or nullptr).
+    Node* pointer_return_value = nullptr;
+};
+
+/// FunctionInfo holds information about the uniformity requirements and effects for a particular
+/// function, as well as the control flow graph.
+struct FunctionInfo {
+    /// Constructor
+    /// @param func the AST function
+    /// @param builder the program builder
+    FunctionInfo(const ast::Function* func, const ProgramBuilder* builder) {
+        name = builder->Symbols().NameFor(func->symbol);
+        callsite_tag = CallSiteNoRestriction;
+        function_tag = NoRestriction;
+
+        // Create special nodes.
+        required_to_be_uniform = CreateNode("RequiredToBeUniform");
+        may_be_non_uniform = CreateNode("MayBeNonUniform");
+        cf_start = CreateNode("CF_start");
+        cf_return = CreateNode("CF_return");
+        if (func->return_type) {
+            value_return = CreateNode("Value_return");
+        }
+
+        // Create nodes for parameters.
+        parameters.resize(func->params.size());
+        for (size_t i = 0; i < func->params.size(); i++) {
+            auto* param = func->params[i];
+            auto param_name = builder->Symbols().NameFor(param->symbol);
+            auto* sem = builder->Sem().Get<sem::Parameter>(param);
+            parameters[i].sem = sem;
+
+            Node* node_init;
+            if (sem->Type()->Is<sem::Pointer>()) {
+                node_init = CreateNode("ptrparam_" + name + "_init");
+                parameters[i].pointer_return_value = CreateNode("ptrparam_" + name + "_return");
+                local_var_decls.insert(sem);
+            } else {
+                node_init = CreateNode("param_" + name);
+            }
+            parameters[i].init_value = node_init;
+            variables.Set(sem, node_init);
+        }
+    }
+
+    /// The name of the function.
+    std::string name;
+
+    /// The call site uniformity requirements.
+    CallSiteTag callsite_tag;
+    /// The function's uniformity effects.
+    FunctionTag function_tag;
+    /// The uniformity requirements of the function's parameters.
+    std::vector<ParameterInfo> parameters;
+
+    /// The control flow graph.
+    utils::BlockAllocator<Node> nodes;
+
+    /// Special `RequiredToBeUniform` node.
+    Node* required_to_be_uniform;
+    /// Special `MayBeNonUniform` node.
+    Node* may_be_non_uniform;
+    /// Special `CF_start` node.
+    Node* cf_start;
+    /// Special `CF_return` node.
+    Node* cf_return;
+    /// Special `Value_return` node.
+    Node* value_return;
+
+    /// Map from variables to their value nodes in the graph, scoped with respect to control flow.
+    ScopeStack<const sem::Variable*, Node*> variables;
+
+    /// The set of a local read-write vars that are in scope at any given point in the process.
+    /// Includes pointer parameters.
+    std::unordered_set<const sem::Variable*> local_var_decls;
+
+    /// LoopSwitchInfo tracks information about the value of variables for a control flow construct.
+    struct LoopSwitchInfo {
+        /// The type of this control flow construct.
+        std::string type;
+        /// The input values for local variables at the start of this construct.
+        std::unordered_map<const sem::Variable*, Node*> var_in_nodes;
+        /// The exit values for local variables at the end of this construct.
+        std::unordered_map<const sem::Variable*, Node*> var_exit_nodes;
+    };
+
+    /// Map from control flow statements to the corresponding LoopSwitchInfo structure.
+    std::unordered_map<const sem::Statement*, LoopSwitchInfo> loop_switch_infos;
+
+    /// Create a new node.
+    /// @param tag a tag used to identify the node for debugging purposes
+    /// @param ast the optional AST node that this node corresponds to
+    /// @returns the new node
+    Node* CreateNode([[maybe_unused]] std::string tag, const ast::Node* ast = nullptr) {
+        auto* node = nodes.Create(ast);
+
+#if TINT_DUMP_UNIFORMITY_GRAPH
+        // Make the tag unique and set it.
+        // This only matters if we're dumping the graph.
+        std::string unique_tag = tag;
+        int suffix = 0;
+        while (tags_.count(unique_tag)) {
+            unique_tag = tag + "_$" + std::to_string(++suffix);
+        }
+        tags_.insert(unique_tag);
+        node->tag = name + "." + unique_tag;
+#endif
+
+        return node;
+    }
+
+    /// Reset the visited status of every node in the graph.
+    void ResetVisited() {
+        for (auto* node : nodes.Objects()) {
+            node->visited_from = nullptr;
+        }
+    }
+
+  private:
+    /// A list of tags that have already been used within the current function.
+    std::unordered_set<std::string> tags_;
+};
+
 /// UniformityGraph is used to analyze the uniformity requirements and effects of functions in a
 /// module.
 class UniformityGraph {
@@ -174,115 +282,34 @@
     /// Map of analyzed function results.
     std::unordered_map<const ast::Function*, FunctionInfo> functions_;
 
-    /// Allocator for nodes.
-    utils::BlockAllocator<Node> nodes_;
-
-    /// Name of the function currently being analyzed.
-    std::string current_function_;
-
-    /// Special `RequiredToBeUniform` node.
-    Node* required_to_be_uniform_;
-    /// Special `MayBeNonUniform` node.
-    Node* may_be_non_uniform_;
-    /// Special `CF_return` node.
-    Node* cf_return_;
-    /// Special `Value_return` node.
-    Node* value_return_;
-    /// Special `{param}_return` nodes for pointer parameters.
-    std::unordered_map<const sem::Parameter*, Node*> pointer_param_returns_;
-
-    /// Map from variables to their value nodes in the graph, scoped with respect to control flow.
-    ScopeStack<const sem::Variable*, Node*> variables_;
-
-    /// The set of a local read-write vars that are in scope at any given point in the process.
-    /// Includes pointer parameters.
-    std::unordered_set<const sem::Variable*> local_var_decls_;
-
-    /// LoopSwitchInfo tracks information about the value of variables for a control flow construct.
-    struct LoopSwitchInfo {
-        /// The type of this control flow construct.
-        std::string type;
-        /// The input values for local variables at the start of this construct.
-        std::unordered_map<const sem::Variable*, Node*> var_in_nodes;
-        /// The exit values for local variables at the end of this construct.
-        std::unordered_map<const sem::Variable*, Node*> var_exit_nodes;
-    };
-
-    /// Map from control flow statements to the corresponding LoopSwitchInfo structure.
-    std::unordered_map<const sem::Statement*, LoopSwitchInfo> loop_switch_infos_;
-
-    /// A list of tags that have already been used within the current function.
-    std::unordered_set<std::string> tags_;
+    /// The function currently being analyzed.
+    FunctionInfo* current_function_;
 
     /// Create a new node.
     /// @param tag a tag used to identify the node for debugging purposes.
     /// @param ast the optional AST node that this node corresponds to
     /// @returns the new node
     Node* CreateNode(std::string tag, const ast::Node* ast = nullptr) {
-        std::string unique_tag = tag;
-
-#if TINT_DUMP_UNIFORMITY_GRAPH
-        // Make the tag unique.
-        // This only matters if we're dumping the graph.
-        int suffix = 0;
-        while (tags_.count(unique_tag)) {
-            unique_tag = tag + "_$" + std::to_string(++suffix);
-        }
-        tags_.insert(unique_tag);
-#endif
-
-        return nodes_.Create(current_function_ + "." + unique_tag, ast);
+        return current_function_->CreateNode(tag, ast);
     }
 
     /// Process a function.
     /// @param func the function to process
     /// @returns true if there are no uniformity issues, false otherwise
     bool ProcessFunction(const ast::Function* func) {
-        nodes_.Reset();
-        variables_.Clear();
-        pointer_param_returns_.clear();
-        tags_.clear();
-
-        current_function_ = builder_->Symbols().NameFor(func->symbol);
-
-        // Create special nodes.
-        Node* cf_start = CreateNode("CF_start");
-        required_to_be_uniform_ = CreateNode("RequiredToBeUniform");
-        may_be_non_uniform_ = CreateNode("MayBeNonUniform");
-        cf_return_ = CreateNode("CF_return");
-        if (func->return_type) {
-            value_return_ = CreateNode("Value_return");
-        }
-
-        // Create nodes for parameters.
-        std::vector<Node*> param_init_values(func->params.size(), nullptr);
-        for (size_t i = 0; i < func->params.size(); i++) {
-            auto* param = func->params[i];
-            auto name = builder_->Symbols().NameFor(param->symbol);
-            auto sem = sem_.Get<sem::Parameter>(param);
-
-            Node* node_init;
-            if (sem->Type()->Is<sem::Pointer>()) {
-                node_init = CreateNode("ptrparam_" + name + "_init");
-                pointer_param_returns_[sem] = CreateNode("ptrparam_" + name + "_return");
-                local_var_decls_.insert(sem);
-            } else {
-                node_init = CreateNode("param_" + name);
-            }
-            param_init_values[i] = node_init;
-            variables_.Set(sem, node_init);
-        }
+        functions_.emplace(func, FunctionInfo(func, builder_));
+        current_function_ = &functions_.at(func);
 
         // Process function body.
         if (func->body) {
-            ProcessStatement(cf_start, func->body);
+            ProcessStatement(current_function_->cf_start, func->body);
         }
 
 #if TINT_DUMP_UNIFORMITY_GRAPH
         // Dump the graph for this function as a subgraph.
-        std::cout << "\nsubgraph cluster_" << current_function_ << " {\n";
-        std::cout << "  label=" << current_function_ << ";";
-        for (auto* node : nodes_.Objects()) {
+        std::cout << "\nsubgraph cluster_" << current_function_->name << " {\n";
+        std::cout << "  label=" << current_function_->name << ";";
+        for (auto* node : current_function_->nodes.Objects()) {
             std::cout << "\n  \"" << node->tag << "\";";
             for (auto* edge : node->edges) {
                 std::cout << "\n  \"" << node->tag << "\" -> \"" << edge->tag << "\";";
@@ -291,29 +318,24 @@
         std::cout << "\n}\n";
 #endif
 
-        FunctionInfo& info = functions_[func];
-        info.callsite_tag = CallSiteNoRestriction;
-        info.function_tag = NoRestriction;
-        info.parameters.resize(func->params.size());
-
         // Look at which nodes are reachable from "RequiredToBeUniform".
         {
             utils::UniqueVector<Node*> reachable;
-            Traverse(required_to_be_uniform_, reachable);
-            if (reachable.contains(may_be_non_uniform_)) {
+            Traverse(current_function_->required_to_be_uniform, reachable);
+            if (reachable.contains(current_function_->may_be_non_uniform)) {
                 MakeError();
                 return false;
             }
-            if (reachable.contains(cf_start)) {
-                info.callsite_tag = CallSiteRequiredToBeUniform;
+            if (reachable.contains(current_function_->cf_start)) {
+                current_function_->callsite_tag = CallSiteRequiredToBeUniform;
             }
 
             // Set the parameter tag to ParameterRequiredToBeUniform for each parameter node that
             // was reachable.
             for (size_t i = 0; i < func->params.size(); i++) {
                 auto* param = func->params[i];
-                if (reachable.contains(variables_.Get(sem_.Get(param)))) {
-                    info.parameters[i].tag = ParameterRequiredToBeUniform;
+                if (reachable.contains(current_function_->variables.Get(sem_.Get(param)))) {
+                    current_function_->parameters[i].tag = ParameterRequiredToBeUniform;
                 }
             }
         }
@@ -321,62 +343,62 @@
         // Look at which nodes are reachable from "CF_return"
         {
             utils::UniqueVector<Node*> reachable;
-            Traverse(cf_return_, reachable);
-            if (reachable.contains(may_be_non_uniform_)) {
-                info.function_tag = SubsequentControlFlowMayBeNonUniform;
+            Traverse(current_function_->cf_return, reachable);
+            if (reachable.contains(current_function_->may_be_non_uniform)) {
+                current_function_->function_tag = SubsequentControlFlowMayBeNonUniform;
             }
 
             // Set the parameter tag to ParameterRequiredToBeUniformForSubsequentControlFlow for
             // each parameter node that was reachable.
             for (size_t i = 0; i < func->params.size(); i++) {
                 auto* param = func->params[i];
-                if (reachable.contains(variables_.Get(sem_.Get(param)))) {
-                    info.parameters[i].tag = ParameterRequiredToBeUniformForSubsequentControlFlow;
+                if (reachable.contains(current_function_->variables.Get(sem_.Get(param)))) {
+                    current_function_->parameters[i].tag =
+                        ParameterRequiredToBeUniformForSubsequentControlFlow;
                 }
             }
         }
 
         // If "Value_return" exists, look at which nodes are reachable from it
-        if (value_return_) {
+        if (current_function_->value_return) {
             utils::UniqueVector<Node*> reachable;
-            Traverse(value_return_, reachable);
-            if (reachable.contains(may_be_non_uniform_)) {
-                info.function_tag = ReturnValueMayBeNonUniform;
+            Traverse(current_function_->value_return, reachable);
+            if (reachable.contains(current_function_->may_be_non_uniform)) {
+                current_function_->function_tag = ReturnValueMayBeNonUniform;
             }
 
             // Set the parameter tag to ParameterRequiredToBeUniformForReturnValue for each
             // parameter node that was reachable.
             for (size_t i = 0; i < func->params.size(); i++) {
                 auto* param = func->params[i];
-                if (reachable.contains(variables_.Get(sem_.Get(param)))) {
-                    info.parameters[i].tag = ParameterRequiredToBeUniformForReturnValue;
+                if (reachable.contains(current_function_->variables.Get(sem_.Get(param)))) {
+                    current_function_->parameters[i].tag =
+                        ParameterRequiredToBeUniformForReturnValue;
                 }
             }
         }
 
         // Traverse the graph for each pointer parameter.
         for (size_t i = 0; i < func->params.size(); i++) {
-            auto* param_dest = sem_.Get<sem::Parameter>(func->params[i]);
-            if (!param_dest->Type()->Is<sem::Pointer>()) {
+            if (current_function_->parameters[i].pointer_return_value == nullptr) {
                 continue;
             }
 
             // Reset "visited" state for all nodes.
-            for (auto* node : nodes_.Objects()) {
-                node->visited_from = nullptr;
-            }
+            current_function_->ResetVisited();
 
             utils::UniqueVector<Node*> reachable;
-            Traverse(pointer_param_returns_[param_dest], reachable);
-            if (reachable.contains(may_be_non_uniform_)) {
-                info.parameters[i].pointer_may_become_non_uniform = true;
+            Traverse(current_function_->parameters[i].pointer_return_value, reachable);
+            if (reachable.contains(current_function_->may_be_non_uniform)) {
+                current_function_->parameters[i].pointer_may_become_non_uniform = true;
             }
 
             // Check every other parameter to see if they feed into this parameter's final value.
             for (size_t j = 0; j < func->params.size(); j++) {
                 auto* param_source = sem_.Get<sem::Parameter>(func->params[j]);
-                if (reachable.contains(param_init_values[j])) {
-                    info.parameters[i].pointer_param_output_sources.push_back(param_source);
+                if (reachable.contains(current_function_->parameters[j].init_value)) {
+                    current_function_->parameters[i].pointer_param_output_sources.push_back(
+                        param_source);
                 }
             }
         }
@@ -407,8 +429,8 @@
                 std::unordered_map<const sem::Variable*, Node*> scoped_assignments;
                 {
                     // Push a new scope for variable assignments in the block.
-                    variables_.Push();
-                    TINT_DEFER(variables_.Pop());
+                    current_function_->variables.Push();
+                    TINT_DEFER(current_function_->variables.Pop());
 
                     for (auto* s : b->statements) {
                         cf = ProcessStatement(cf, s);
@@ -420,12 +442,15 @@
                     if (sem_.Get<sem::FunctionBlockStatement>(b)) {
                         // We've reached the end of the function body.
                         // Add edges from pointer parameter outputs to their current value.
-                        for (auto param : pointer_param_returns_) {
-                            param.second->AddEdge(variables_.Get(param.first));
+                        for (auto param : current_function_->parameters) {
+                            if (param.pointer_return_value) {
+                                param.pointer_return_value->AddEdge(
+                                    current_function_->variables.Get(param.sem));
+                            }
                         }
                     }
 
-                    scoped_assignments = std::move(variables_.Top());
+                    scoped_assignments = std::move(current_function_->variables.Top());
                 }
 
                 // Propagate all variables assignments to the containing scope if the behavior is
@@ -434,13 +459,13 @@
                 if (behaviors.Contains(sem::Behavior::kNext) ||
                     behaviors.Contains(sem::Behavior::kFallthrough)) {
                     for (auto var : scoped_assignments) {
-                        variables_.Set(var.first, var.second);
+                        current_function_->variables.Set(var.first, var.second);
                     }
                 }
 
                 // Remove any variables declared in this scope from the set of in-scope variables.
                 for (auto* d : sem_.Get<sem::BlockStatement>(b)->Decls()) {
-                    local_var_decls_.erase(sem_.Get<sem::LocalVariable>(d));
+                    current_function_->local_var_decls.erase(sem_.Get<sem::LocalVariable>(d));
                 }
 
                 return cf;
@@ -451,11 +476,11 @@
                 auto* parent = sem_.Get(b)
                                    ->FindFirstParent<sem::SwitchStatement, sem::LoopStatement,
                                                      sem::ForLoopStatement>();
-                TINT_ASSERT(Resolver, loop_switch_infos_.count(parent));
-                auto& info = loop_switch_infos_.at(parent);
+                TINT_ASSERT(Resolver, current_function_->loop_switch_infos.count(parent));
+                auto& info = current_function_->loop_switch_infos.at(parent);
 
                 // Propagate variable values to the loop/switch exit nodes.
-                for (auto* var : local_var_decls_) {
+                for (auto* var : current_function_->local_var_decls) {
                     // Skip variables that were declared inside this loop/switch.
                     if (auto* lv = var->As<sem::LocalVariable>();
                         lv &&
@@ -468,7 +493,7 @@
                         auto name = builder_->Symbols().NameFor(var->Declaration()->symbol);
                         return CreateNode(name + "_value_" + info.type + "_exit");
                     });
-                    exit_node->AddEdge(variables_.Get(var));
+                    exit_node->AddEdge(current_function_->variables.Get(var));
                 }
 
                 return cf;
@@ -496,11 +521,11 @@
                 // Find the loop statement that we are in.
                 auto* parent =
                     sem_.Get(c)->FindFirstParent<sem::LoopStatement, sem::ForLoopStatement>();
-                TINT_ASSERT(Resolver, loop_switch_infos_.count(parent));
-                auto& info = loop_switch_infos_.at(parent);
+                TINT_ASSERT(Resolver, current_function_->loop_switch_infos.count(parent));
+                auto& info = current_function_->loop_switch_infos.at(parent);
 
                 // Propagate assignments to the loop input nodes.
-                for (auto* var : local_var_decls_) {
+                for (auto* var : current_function_->local_var_decls) {
                     // Skip variables that were declared inside this loop.
                     if (auto* lv = var->As<sem::LocalVariable>();
                         lv &&
@@ -511,7 +536,7 @@
                     // Add an edge from the variable's loop input node to its value at this point.
                     TINT_ASSERT(Resolver, info.var_in_nodes.count(var));
                     auto* in_node = info.var_in_nodes.at(var);
-                    auto* out_node = variables_.Get(var);
+                    auto* out_node = current_function_->variables.Get(var);
                     if (out_node != in_node) {
                         in_node->AddEdge(out_node);
                     }
@@ -520,7 +545,7 @@
             },
 
             [&](const ast::DiscardStatement*) {
-                cf_return_->AddEdge(cf);
+                current_function_->cf_return->AddEdge(cf);
                 return cf;
             },
 
@@ -537,16 +562,16 @@
                 }
                 auto cf_start = cf_init;
 
-                auto& info = loop_switch_infos_[sem_loop];
+                auto& info = current_function_->loop_switch_infos[sem_loop];
                 info.type = "forloop";
 
                 // Create input nodes for any variables declared before this loop.
-                for (auto* v : local_var_decls_) {
+                for (auto* v : current_function_->local_var_decls) {
                     auto name = builder_->Symbols().NameFor(v->Declaration()->symbol);
                     auto* in_node = CreateNode(name + "_value_forloop_in");
-                    in_node->AddEdge(variables_.Get(v));
+                    in_node->AddEdge(current_function_->variables.Get(v));
                     info.var_in_nodes[v] = in_node;
-                    variables_.Set(v, in_node);
+                    current_function_->variables.Set(v, in_node);
                 }
 
                 // Insert the condition at the start of the loop body.
@@ -557,12 +582,12 @@
                     cf_start = cf_condition_end;
 
                     // Propagate assignments to the loop exit nodes.
-                    for (auto* var : local_var_decls_) {
+                    for (auto* var : current_function_->local_var_decls) {
                         auto* exit_node = utils::GetOrCreate(info.var_exit_nodes, var, [&]() {
                             auto name = builder_->Symbols().NameFor(var->Declaration()->symbol);
                             return CreateNode(name + "_value_" + info.type + "_exit");
                         });
-                        exit_node->AddEdge(variables_.Get(var));
+                        exit_node->AddEdge(current_function_->variables.Get(var));
                     }
                 }
                 auto cf1 = ProcessStatement(cf_start, f->body);
@@ -579,7 +604,7 @@
                 // Add edges from variable loop input nodes to their values at the end of the loop.
                 for (auto v : info.var_in_nodes) {
                     auto* in_node = v.second;
-                    auto* out_node = variables_.Get(v.first);
+                    auto* out_node = current_function_->variables.Get(v.first);
                     if (out_node != in_node) {
                         in_node->AddEdge(out_node);
                     }
@@ -587,10 +612,10 @@
 
                 // Set each variable's exit node as its value in the outer scope.
                 for (auto v : info.var_exit_nodes) {
-                    variables_.Set(v.first, v.second);
+                    current_function_->variables.Set(v.first, v.second);
                 }
 
-                loop_switch_infos_.erase(sem_loop);
+                current_function_->loop_switch_infos.erase(sem_loop);
 
                 if (sem_loop->Behaviors() == sem::Behaviors{sem::Behavior::kNext}) {
                     return cf;
@@ -613,15 +638,15 @@
                     [&](Node* cf_in, const ast::Statement* s,
                         std::unordered_map<const sem::Variable*, Node*>& assigned_vars) {
                         // Push a new scope for variable assignments.
-                        variables_.Push();
+                        current_function_->variables.Push();
 
                         // Process the statement.
                         auto cf_out = ProcessStatement(cf_in, s);
 
-                        assigned_vars = variables_.Top();
+                        assigned_vars = current_function_->variables.Top();
 
                         // Pop the scope and return.
-                        variables_.Pop();
+                        current_function_->variables.Pop();
                         return cf_out;
                     };
 
@@ -639,7 +664,7 @@
                 }
 
                 // Update values for any variables assigned in the if or else blocks.
-                for (auto var : local_var_decls_) {
+                for (auto var : current_function_->local_var_decls) {
                     // Skip variables not assigned in either block.
                     if (true_vars.count(var) == 0 && false_vars.count(var) == 0) {
                         continue;
@@ -655,18 +680,18 @@
                         if (true_vars.count(var)) {
                             out_node->AddEdge(true_vars.at(var));
                         } else {
-                            out_node->AddEdge(variables_.Get(var));
+                            out_node->AddEdge(current_function_->variables.Get(var));
                         }
                     }
                     if (false_has_next) {
                         if (false_vars.count(var)) {
                             out_node->AddEdge(false_vars.at(var));
                         } else {
-                            out_node->AddEdge(variables_.Get(var));
+                            out_node->AddEdge(current_function_->variables.Get(var));
                         }
                     }
 
-                    variables_.Set(var, out_node);
+                    current_function_->variables.Set(var, out_node);
                 }
 
                 if (sem_if->Behaviors() != sem::Behaviors{sem::Behavior::kNext}) {
@@ -696,16 +721,16 @@
                 auto* sem_loop = sem_.Get(l);
                 auto cfx = CreateNode("loop_start");
 
-                auto& info = loop_switch_infos_[sem_loop];
+                auto& info = current_function_->loop_switch_infos[sem_loop];
                 info.type = "loop";
 
                 // Create input nodes for any variables declared before this loop.
-                for (auto* v : local_var_decls_) {
+                for (auto* v : current_function_->local_var_decls) {
                     auto name = builder_->Symbols().NameFor(v->Declaration()->symbol);
                     auto* in_node = CreateNode(name + "_value_loop_in");
-                    in_node->AddEdge(variables_.Get(v));
+                    in_node->AddEdge(current_function_->variables.Get(v));
                     info.var_in_nodes[v] = in_node;
-                    variables_.Set(v, in_node);
+                    current_function_->variables.Set(v, in_node);
                 }
 
                 auto cf1 = ProcessStatement(cfx, l->body);
@@ -720,7 +745,7 @@
                 // Add edges from variable loop input nodes to their values at the end of the loop.
                 for (auto v : info.var_in_nodes) {
                     auto* in_node = v.second;
-                    auto* out_node = variables_.Get(v.first);
+                    auto* out_node = current_function_->variables.Get(v.first);
                     if (out_node != in_node) {
                         in_node->AddEdge(out_node);
                     }
@@ -728,10 +753,10 @@
 
                 // Set each variable's exit node as its value in the outer scope.
                 for (auto v : info.var_exit_nodes) {
-                    variables_.Set(v.first, v.second);
+                    current_function_->variables.Set(v.first, v.second);
                 }
 
-                loop_switch_infos_.erase(sem_loop);
+                current_function_->loop_switch_infos.erase(sem_loop);
 
                 if (sem_loop->Behaviors() == sem::Behaviors{sem::Behavior::kNext}) {
                     return cf;
@@ -743,18 +768,21 @@
                 Node* cf_ret;
                 if (r->value) {
                     auto [cf1, v] = ProcessExpression(cf, r->value);
-                    cf_return_->AddEdge(cf1);
-                    value_return_->AddEdge(v);
+                    current_function_->cf_return->AddEdge(cf1);
+                    current_function_->value_return->AddEdge(v);
                     cf_ret = cf1;
                 } else {
                     TINT_ASSERT(Resolver, cf != nullptr);
-                    cf_return_->AddEdge(cf);
+                    current_function_->cf_return->AddEdge(cf);
                     cf_ret = cf;
                 }
 
                 // Add edges from each pointer parameter output to its current value.
-                for (auto param : pointer_param_returns_) {
-                    param.second->AddEdge(variables_.Get(param.first));
+                for (auto param : current_function_->parameters) {
+                    if (param.pointer_return_value) {
+                        param.pointer_return_value->AddEdge(
+                            current_function_->variables.Get(param.sem));
+                    }
                 }
 
                 return cf_ret;
@@ -768,7 +796,7 @@
                     cf_end = CreateNode("switch_CFend");
                 }
 
-                auto& info = loop_switch_infos_[sem_switch];
+                auto& info = current_function_->loop_switch_infos[sem_switch];
                 info.type = "switch";
 
                 auto cf_n = v;
@@ -779,7 +807,7 @@
                     if (previous_case_has_fallthrough) {
                         cf_n = ProcessStatement(cf_n, c->body);
                     } else {
-                        variables_.Push();
+                        current_function_->variables.Push();
                         cf_n = ProcessStatement(v, c->body);
                     }
 
@@ -792,7 +820,7 @@
                     if (!has_fallthrough) {
                         if (sem_case->Behaviors().Contains(sem::Behavior::kNext)) {
                             // Propagate variable values to the switch exit nodes.
-                            for (auto* var : local_var_decls_) {
+                            for (auto* var : current_function_->local_var_decls) {
                                 // Skip variables that were declared inside the switch.
                                 if (auto* lv = var->As<sem::LocalVariable>();
                                     lv && lv->Statement()->FindFirstParent(
@@ -807,17 +835,17 @@
                                             builder_->Symbols().NameFor(var->Declaration()->symbol);
                                         return CreateNode(name + "_value_" + info.type + "_exit");
                                     });
-                                exit_node->AddEdge(variables_.Get(var));
+                                exit_node->AddEdge(current_function_->variables.Get(var));
                             }
                         }
-                        variables_.Pop();
+                        current_function_->variables.Pop();
                     }
                     previous_case_has_fallthrough = has_fallthrough;
                 }
 
                 // Update nodes for any variables assigned in the switch statement.
                 for (auto var : info.var_exit_nodes) {
-                    variables_.Set(var.first, var.second);
+                    current_function_->variables.Set(var.first, var.second);
                 }
 
                 return cf_end ? cf_end : cf;
@@ -831,10 +859,11 @@
                 } else {
                     node = cf;
                 }
-                variables_.Set(sem_.Get(decl->variable), node);
+                current_function_->variables.Set(sem_.Get(decl->variable), node);
 
                 if (!decl->variable->is_const) {
-                    local_var_decls_.insert(sem_.Get<sem::LocalVariable>(decl->variable));
+                    current_function_->local_var_decls.insert(
+                        sem_.Get<sem::LocalVariable>(decl->variable));
                 }
 
                 return cf;
@@ -881,16 +910,17 @@
                                 uniform = false;
                             }
                         }
-                        return std::make_pair(cf, uniform ? cf : may_be_non_uniform_);
+                        return std::make_pair(cf,
+                                              uniform ? cf : current_function_->may_be_non_uniform);
                     } else {
                         if (has_nonuniform_entry_point_attribute(param->Declaration())) {
-                            return std::make_pair(cf, may_be_non_uniform_);
+                            return std::make_pair(cf, current_function_->may_be_non_uniform);
                         }
                         return std::make_pair(cf, cf);
                     }
                 } else {
                     auto* result = CreateNode(name + "_result");
-                    auto* x = variables_.Get(param);
+                    auto* x = current_function_->variables.Get(param);
                     result->AddEdge(cf);
                     result->AddEdge(x);
                     return std::make_pair(cf, result);
@@ -901,14 +931,14 @@
                 if (global->Declaration()->is_const || global->Access() == ast::Access::kRead) {
                     return std::make_pair(cf, cf);
                 } else {
-                    return std::make_pair(cf, may_be_non_uniform_);
+                    return std::make_pair(cf, current_function_->may_be_non_uniform);
                 }
             },
 
             [&](const sem::LocalVariable* local) {
                 auto* result = CreateNode(name + "_result");
                 result->AddEdge(cf);
-                if (auto* x = variables_.Get(local)) {
+                if (auto* x = current_function_->variables.Get(local)) {
                     result->AddEdge(x);
                 }
                 return std::make_pair(cf, result);
@@ -971,7 +1001,7 @@
                     // Cut the analysis short, since we only need to know the originating variable
                     // which is being accessed.
                     auto* source_var = sem_.Get(u)->SourceVariable();
-                    auto* value = variables_.Get(source_var);
+                    auto* value = current_function_->variables.Get(source_var);
                     if (!value) {
                         value = cf;
                     }
@@ -999,11 +1029,11 @@
                 auto name = builder_->Symbols().NameFor(i->symbol);
                 auto* sem = sem_.Get<sem::VariableUser>(i);
                 if (sem->Variable()->Is<sem::GlobalVariable>()) {
-                    return std::make_pair(cf, may_be_non_uniform_);
+                    return std::make_pair(cf, current_function_->may_be_non_uniform);
                 } else if (auto* local = sem->Variable()->As<sem::LocalVariable>()) {
                     // Create a new value node for this variable.
                     auto* value = CreateNode(name + "_lvalue");
-                    auto* old_value = variables_.Set(local, value);
+                    auto* old_value = current_function_->variables.Set(local, value);
 
                     // Aggregate values link back to their previous value, as they can never become
                     // uniform again.
@@ -1038,7 +1068,7 @@
                     auto* source_var = sem_.Get(u)->SourceVariable();
                     auto name = builder_->Symbols().NameFor(source_var->Declaration()->symbol);
                     auto* deref = CreateNode(name + "_deref");
-                    auto* old_value = variables_.Set(source_var, deref);
+                    auto* old_value = current_function_->variables.Set(source_var, deref);
 
                     // Aggregate values link back to their previous value, as they can never become
                     // uniform again.
@@ -1138,14 +1168,14 @@
             // Note: This deviates from the rules in the specification, which would add the edge
             // directly to the incoming CF node. Going through CF_after instead makes it easier to
             // produce diagnostics that can identify the function being called.
-            required_to_be_uniform_->AddEdge(cf_after);
+            current_function_->required_to_be_uniform->AddEdge(cf_after);
         }
         cf_after->AddEdge(cf_last_arg);
 
         if (function_tag == SubsequentControlFlowMayBeNonUniform) {
-            cf_after->AddEdge(may_be_non_uniform_);
+            cf_after->AddEdge(current_function_->may_be_non_uniform);
         } else if (function_tag == ReturnValueMayBeNonUniform) {
-            result->AddEdge(may_be_non_uniform_);
+            result->AddEdge(current_function_->may_be_non_uniform);
         }
 
         result->AddEdge(cf_after);
@@ -1155,7 +1185,7 @@
             if (func_info) {
                 switch (func_info->parameters[i].tag) {
                     case ParameterRequiredToBeUniform:
-                        required_to_be_uniform_->AddEdge(args[i]);
+                        current_function_->required_to_be_uniform->AddEdge(args[i]);
                         break;
                     case ParameterRequiredToBeUniformForSubsequentControlFlow:
                         cf_after->AddEdge(args[i]);
@@ -1172,7 +1202,7 @@
                     auto* ptr_result =
                         CreateNode(name + "_ptrarg_" + std::to_string(i) + "_result");
                     if (func_info->parameters[i].pointer_may_become_non_uniform) {
-                        ptr_result->AddEdge(may_be_non_uniform_);
+                        ptr_result->AddEdge(current_function_->may_be_non_uniform);
                     } else {
                         // Add edges from the resulting pointer value to any other arguments that
                         // feed it.
@@ -1184,7 +1214,7 @@
                     // Update the current stored value for this pointer argument.
                     auto* source_var = sem_arg->SourceVariable();
                     TINT_ASSERT(Resolver, source_var);
-                    variables_.Set(source_var, ptr_result);
+                    current_function_->variables.Set(source_var, ptr_result);
                 }
             } else {
                 // All builtin function parameters are RequiredToBeUniformForReturnValue, as are
@@ -1219,10 +1249,10 @@
     void MakeError() {
         // Trace back to find a node that is required to be uniform that was reachable from a
         // non-uniform value or control flow node.
-        Node* current = may_be_non_uniform_;
+        Node* current = current_function_->may_be_non_uniform;
         while (current) {
             TINT_ASSERT(Resolver, current->visited_from);
-            if (current->visited_from == required_to_be_uniform_) {
+            if (current->visited_from == current_function_->required_to_be_uniform) {
                 break;
             }
             current = current->visited_from;