From ba0778a236a34f7689137d7fb9dd6b1c1790c162 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <Jakob.botsch.nielsen@gmail.com>
Date: Fri, 23 Sep 2022 18:32:58 +0200
Subject: [PATCH] JIT: Update gtHasLocalsWithAddrOp to use GenTreeVisitor
 (#75672)

Also change it to check for any local, not just GT_LCL_VAR.
---
 src/coreclr/jit/compiler.h  |  2 --
 src/coreclr/jit/gentree.cpp | 69 +++++++++++++++++++------------------
 2 files changed, 35 insertions(+), 36 deletions(-)

diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
index 383b0a7de59..8504476d2cd 100644
--- a/src/coreclr/jit/compiler.h
+++ b/src/coreclr/jit/compiler.h
@@ -5900,8 +5900,6 @@ private:
     // the variable is not enregistered, and is therefore not promoted independently.
     void fgLclFldAssign(unsigned lclNum);
 
-    static fgWalkPreFn gtHasLocalsWithAddrOpCB;
-
     enum TypeProducerKind
     {
         TPK_Unknown = 0, // May not be a RuntimeType
diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp
index 5549a5166b3..1e97a840a32 100644
--- a/src/coreclr/jit/gentree.cpp
+++ b/src/coreclr/jit/gentree.cpp
@@ -2724,48 +2724,49 @@ AGAIN:
     return result;
 }
 
-struct AddrTakenDsc
-{
-    Compiler* comp;
-    bool      hasAddrTakenLcl;
-};
-
-/* static */
-Compiler::fgWalkResult Compiler::gtHasLocalsWithAddrOpCB(GenTree** pTree, fgWalkData* data)
+//------------------------------------------------------------------------------
+// gtHasLocalsWithAddrOp:
+//   Check if this tree contains locals with lvHasLdAddrOp or
+//   IsAddressExposed() flags set. Does a full tree walk.
+//
+// Paramters:
+//   tree - the tree
+//
+// Return Value:
+//    True if any sub tree is such a local.
+//
+bool Compiler::gtHasLocalsWithAddrOp(GenTree* tree)
 {
-    GenTree*  tree = *pTree;
-    Compiler* comp = data->compiler;
-
-    if (tree->gtOper == GT_LCL_VAR)
+    struct LocalsWithAddrOpVisitor : GenTreeVisitor<LocalsWithAddrOpVisitor>
     {
-        const LclVarDsc* varDsc = comp->lvaGetDesc(tree->AsLclVarCommon());
-
-        if (varDsc->lvHasLdAddrOp || varDsc->IsAddressExposed())
+        enum
         {
-            ((AddrTakenDsc*)data->pCallbackData)->hasAddrTakenLcl = true;
-            return WALK_ABORT;
-        }
-    }
-
-    return WALK_CONTINUE;
-}
+            DoPreOrder    = true,
+            DoLclVarsOnly = true,
+        };
 
-/*****************************************************************************
- *
- *  Return true if this tree contains locals with lvHasLdAddrOp or IsAddressExposed()
- *  flag(s) set.
- */
+        bool HasAddrTakenLocal = false;
 
-bool Compiler::gtHasLocalsWithAddrOp(GenTree* tree)
-{
-    AddrTakenDsc desc;
+        LocalsWithAddrOpVisitor(Compiler* comp) : GenTreeVisitor(comp)
+        {
+        }
 
-    desc.comp            = this;
-    desc.hasAddrTakenLcl = false;
+        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
+        {
+            LclVarDsc* varDsc = m_compiler->lvaGetDesc((*use)->AsLclVarCommon());
+            if (varDsc->lvHasLdAddrOp || varDsc->IsAddressExposed())
+            {
+                HasAddrTakenLocal = true;
+                return WALK_ABORT;
+            }
 
-    fgWalkTreePre(&tree, gtHasLocalsWithAddrOpCB, &desc);
+            return WALK_CONTINUE;
+        }
+    };
 
-    return desc.hasAddrTakenLcl;
+    LocalsWithAddrOpVisitor visitor(this);
+    visitor.WalkTree(&tree, nullptr);
+    return visitor.HasAddrTakenLocal;
 }
 
 #ifdef DEBUG
-- 
GitLab