JIT: Add a gtFindNodeInTree helper (dotnet#115975)

Add a gtFindNodeInTree helper to find a node in a tree matching a
predicate, and use it in a couple of places

Author: jakobbotsch

src/coreclr/jit/compiler.h

@@ -6161,8 +6161,6 @@ class Compiler
     PhaseStatus fgHeadTailMerge(bool early);
     bool fgHeadMerge(BasicBlock* block, bool early);
     bool fgTryOneHeadMerge(BasicBlock* block, bool early);
-    template<typename Predicate>
-    bool gtTreeContainsCall(GenTree* tree, Predicate pred);
     bool gtTreeContainsTailCall(GenTree* tree);
     bool gtTreeContainsAsyncCall(GenTree* tree);
     bool fgCanMoveFirstStatementIntoPred(bool early, Statement* firstStmt, BasicBlock* pred);
@@ -6914,6 +6912,9 @@ class Compiler
     bool gtIsTypeHandleToRuntimeTypeHelper(GenTreeCall* call);
     bool gtIsTypeHandleToRuntimeTypeHandleHelper(GenTreeCall* call, CorInfoHelpFunc* pHelper = nullptr);
 
+    template<GenTreeFlags RequiredFlagsToDescendIntoTree, typename Predicate>
+    GenTree* gtFindNodeInTree(GenTree* tree, Predicate predicate);
+
     bool gtTreeContainsOper(GenTree* tree, genTreeOps op);
     ExceptionSetFlags gtCollectExceptions(GenTree* tree);
 

src/coreclr/jit/compiler.hpp

@@ -2199,6 +2199,74 @@ inline bool GenTree::gtOverflowEx() const
     return OperMayOverflow() && gtOverflow();
 }
 
+//------------------------------------------------------------------------
+// gtFindNodeInTree:
+//   Check if a tree contains a node matching the specified predicate. Descend
+//   only into subtrees with the specified flags set on them (can be GTF_EMPTY
+//   to descend into all nodes).
+//
+// Type parameters:
+//   RequiredFlagsToDescendIntoNode - Flags that must be set on the node to
+//                                    descend into it (GTF_EMPTY to descend into all nodes)
+//   Predicate - Type of the predicate (GenTree* -> bool)
+//
+// Parameters:
+//   tree - The tree
+//   pred - Predicate that the call must match
+//
+// Returns:
+//   Node matching the predicate, or nullptr if no such node was found.
+//
+template <GenTreeFlags RequiredFlagsToDescendIntoNode, typename Predicate>
+GenTree* Compiler::gtFindNodeInTree(GenTree* tree, Predicate pred)
+{
+    struct FindNodeVisitor : GenTreeVisitor<FindNodeVisitor>
+    {
+    private:
+        Predicate& m_pred;
+
+    public:
+        GenTree* Result = nullptr;
+
+        enum
+        {
+            DoPreOrder = true
+        };
+
+        FindNodeVisitor(Compiler* comp, Predicate& pred)
+            : GenTreeVisitor<FindNodeVisitor>(comp)
+            , m_pred(pred)
+        {
+        }
+
+        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
+        {
+            GenTree* node = *use;
+            if ((node->gtFlags & RequiredFlagsToDescendIntoNode) != RequiredFlagsToDescendIntoNode)
+            {
+                return WALK_SKIP_SUBTREES;
+            }
+
+            if (m_pred(node))
+            {
+                Result = node;
+                return WALK_ABORT;
+            }
+
+            return WALK_CONTINUE;
+        }
+    };
+
+    if ((tree->gtFlags & RequiredFlagsToDescendIntoNode) != RequiredFlagsToDescendIntoNode)
+    {
+        return nullptr;
+    }
+
+    FindNodeVisitor findNode(this, pred);
+    findNode.WalkTree(&tree, nullptr);
+    return findNode.Result;
+}
+
 /*
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

src/coreclr/jit/fgopt.cpp

@@ -5604,63 +5604,6 @@ bool Compiler::fgHeadMerge(BasicBlock* block, bool early)
     return madeChanges;
 }
 
-//------------------------------------------------------------------------
-// gtTreeContainsCall:
-//   Check if a tree contains a call node matching the given predicate.
-//
-// Parameters:
-//   tree - The tree
-//   pred - Predicate that the call must match
-//
-// Returns:
-//   True if a call node matching the predicate was found, false otherwise.
-//
-template <typename Predicate>
-bool Compiler::gtTreeContainsCall(GenTree* tree, Predicate pred)
-{
-    struct HasCallVisitor : GenTreeVisitor<HasCallVisitor>
-    {
-    private:
-        Predicate& m_pred;
-
-    public:
-        enum
-        {
-            DoPreOrder = true
-        };
-
-        HasCallVisitor(Compiler* comp, Predicate& pred)
-            : GenTreeVisitor<HasCallVisitor>(comp)
-            , m_pred(pred)
-        {
-        }
-
-        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
-        {
-            GenTree* node = *use;
-            if ((node->gtFlags & GTF_CALL) == 0)
-            {
-                return WALK_SKIP_SUBTREES;
-            }
-
-            if (node->IsCall() && m_pred(node->AsCall()))
-            {
-                return WALK_ABORT;
-            }
-
-            return WALK_CONTINUE;
-        }
-    };
-
-    if ((tree->gtFlags & GTF_CALL) == 0)
-    {
-        return false;
-    }
-
-    HasCallVisitor hasCall(this, pred);
-    return hasCall.WalkTree(&tree, nullptr) == WALK_ABORT;
-}
-
 //------------------------------------------------------------------------
 // gtTreeContainsTailCall: Check if a tree contains any tail call or tail call
 // candidate.
@@ -5676,9 +5619,11 @@ bool Compiler::gtTreeContainsCall(GenTree* tree, Predicate pred)
 //
 bool Compiler::gtTreeContainsTailCall(GenTree* tree)
 {
-    return gtTreeContainsCall(tree, [](GenTreeCall* call) {
-        return call->CanTailCall() || call->IsTailCall();
-    });
+    auto isTailCall = [](GenTree* tree) {
+        return tree->IsCall() && (tree->AsCall()->CanTailCall() || tree->AsCall()->IsTailCall());
+    };
+
+    return gtFindNodeInTree<GTF_CALL>(tree, isTailCall) != nullptr;
 }
 
 //------------------------------------------------------------------------
@@ -5697,9 +5642,11 @@ bool Compiler::gtTreeContainsAsyncCall(GenTree* tree)
         return false;
     }
 
-    return gtTreeContainsCall(tree, [](GenTreeCall* call) {
-        return call->IsAsync();
-    });
+    auto isAsyncCall = [](GenTree* tree) {
+        return tree->IsCall() && tree->AsCall()->IsAsync();
+    };
+
+    return gtFindNodeInTree<GTF_CALL>(tree, isAsyncCall) != nullptr;
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/gentree.cpp

@@ -17716,35 +17716,10 @@ bool Compiler::gtIsTypeHandleToRuntimeTypeHandleHelper(GenTreeCall* call, CorInf
 //
 bool Compiler::gtTreeContainsOper(GenTree* tree, genTreeOps oper)
 {
-    class Visitor final : public GenTreeVisitor<Visitor>
-    {
-        genTreeOps m_oper;
-
-    public:
-        Visitor(Compiler* comp, genTreeOps oper)
-            : GenTreeVisitor(comp)
-            , m_oper(oper)
-        {
-        }
-
-        enum
-        {
-            DoPreOrder = true,
-        };
-
-        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
-        {
-            if ((*use)->OperIs(m_oper))
-            {
-                return WALK_ABORT;
-            }
-
-            return WALK_CONTINUE;
-        }
+    auto hasOper = [oper](GenTree* tree) {
+        return tree->OperGet() == oper;
     };
-
-    Visitor visitor(this, oper);
-    return visitor.WalkTree(&tree, nullptr) == WALK_ABORT;
+    return gtFindNodeInTree<GTF_EMPTY>(tree, hasOper) != nullptr;
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/gschecks.cpp

@@ -537,45 +537,19 @@ void Compiler::gsParamsToShadows()
                 // inserting reverse pinvoke transitions way too early in the
                 // JIT.
 
-                struct HasReversePInvokeEnterVisitor : GenTreeVisitor<HasReversePInvokeEnterVisitor>
-                {
-                    enum
-                    {
-                        DoPreOrder = true,
-                    };
-
-                    HasReversePInvokeEnterVisitor(Compiler* comp)
-                        : GenTreeVisitor(comp)
-                    {
-                    }
-
-                    fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
-                    {
-                        if (((*use)->gtFlags & GTF_CALL) == 0)
-                        {
-                            return fgWalkResult::WALK_SKIP_SUBTREES;
-                        }
-
-                        if ((*use)->IsHelperCall(m_compiler, CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER) ||
-                            (*use)->IsHelperCall(m_compiler, CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER_TRACK_TRANSITIONS))
-                        {
-                            return fgWalkResult::WALK_ABORT;
-                        }
-
-                        return fgWalkResult::WALK_CONTINUE;
-                    }
+                auto isReversePInvoke = [=](GenTree* tree) {
+                    return tree->IsHelperCall(this, CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER) ||
+                           tree->IsHelperCall(this, CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER_TRACK_TRANSITIONS);
                 };
 
-                HasReversePInvokeEnterVisitor checker(this);
-
                 Statement* reversePInvokeStmt = nullptr;
                 for (Statement* const stmt : fgFirstBB->Statements())
                 {
                     // assert that we don't have any uses of the local variable
                     // at the point before we insert the shadow copy statement.
                     assert(!gtHasRef(stmt->GetRootNode(), lclNum));
 
-                    if (checker.WalkTree(stmt->GetRootNodePointer(), nullptr) == fgWalkResult::WALK_ABORT)
+                    if (gtFindNodeInTree<GTF_CALL>(stmt->GetRootNode(), isReversePInvoke) != nullptr)
                     {
                         reversePInvokeStmt = stmt;
                         break;