Allow GraphRegistered::GetFromGraph without a parameter

R=pmonette

Bug: 40755583
Change-Id: I817fdb8fe479bee421a70ce7509d29e36145c55b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6335121
Commit-Queue: Joe Mason <joenotcharles@google.com>
Auto-Submit: Joe Mason <joenotcharles@google.com>
Reviewed-by: Patrick Monette <pmonette@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1430552}

components/performance_manager/graph/graph_registered_unittest.cc

@@ -12,7 +12,7 @@
 #include "components/performance_manager/public/graph/graph.h"
 #include "components/performance_manager/public/performance_manager.h"
 #include "components/performance_manager/test_support/graph_test_harness.h"
-#include "components/performance_manager/test_support/performance_manager_test_harness.h"
+#include "components/performance_manager/test_support/test_harness_helper.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -172,4 +172,71 @@ TEST_F(GraphRegisteredTest, GraphOwnedAndRegistered) {
   TearDownAndDestroyGraph();
 }
 
+TEST_F(GraphRegisteredTest, GetFromGraph_Default) {
+  PerformanceManagerTestHarnessHelper pm_helper;
+
+  // Before PerformanceManager is available, GetFromGraph() should safely return
+  // nullptr.
+  EXPECT_FALSE(PerformanceManager::IsAvailable());
+  EXPECT_EQ(Foo::GetFromGraph(), nullptr);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), nullptr);
+
+  pm_helper.SetUp();
+
+  // Objects not registered.
+  EXPECT_TRUE(PerformanceManager::IsAvailable());
+  EXPECT_EQ(Foo::GetFromGraph(), nullptr);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), nullptr);
+
+  // Register objects.
+  Graph* graph = PerformanceManager::GetGraph();
+  ASSERT_TRUE(graph);
+  Foo foo;
+  graph->RegisterObject(&foo);
+  OwnedFoo* owned_foo = graph->PassToGraph(std::make_unique<OwnedFoo>());
+  EXPECT_EQ(Foo::GetFromGraph(), &foo);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), owned_foo);
+
+  // Remove objects.
+  graph->UnregisterObject(&foo);
+  graph->TakeFromGraph(owned_foo);
+  EXPECT_EQ(Foo::GetFromGraph(), nullptr);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), nullptr);
+
+  pm_helper.TearDown();
+
+  // After PerformanceManager is gone, GetFromGraph() should safely return
+  // nullptr again.
+  EXPECT_FALSE(PerformanceManager::IsAvailable());
+  EXPECT_EQ(Foo::GetFromGraph(), nullptr);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), nullptr);
+}
+
+TEST_F(GraphRegisteredTest, GetFromGraph_WithParam) {
+  PerformanceManagerTestHarnessHelper pm_helper;
+  pm_helper.SetUp();
+
+  Graph* graph = PerformanceManager::GetGraph();
+  TestGraphImpl graph2;
+  graph2.SetUp();
+
+  OwnedFoo* foo = graph->PassToGraph(std::make_unique<OwnedFoo>());
+  OwnedBar* bar = graph2.PassToGraph(std::make_unique<OwnedBar>());
+
+  EXPECT_EQ(OwnedFoo::GetFromGraph(), foo);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(graph), foo);
+  EXPECT_EQ(OwnedFoo::GetFromGraph(&graph2), nullptr);
+
+  EXPECT_EQ(OwnedBar::GetFromGraph(), nullptr);
+  EXPECT_EQ(OwnedBar::GetFromGraph(graph), nullptr);
+  EXPECT_EQ(OwnedBar::GetFromGraph(&graph2), bar);
+
+  // Passing a parameter should still work when PM is not available.
+  pm_helper.TearDown();
+  EXPECT_FALSE(PerformanceManager::IsAvailable());
+  EXPECT_EQ(OwnedBar::GetFromGraph(&graph2), bar);
+
+  graph2.TearDown();
+}
+
 }  // namespace performance_manager

components/performance_manager/public/graph/graph_registered.h

@@ -5,7 +5,9 @@
 #ifndef COMPONENTS_PERFORMANCE_MANAGER_PUBLIC_GRAPH_GRAPH_REGISTERED_H_
 #define COMPONENTS_PERFORMANCE_MANAGER_PUBLIC_GRAPH_GRAPH_REGISTERED_H_
 
+#include "base/check.h"
 #include "components/performance_manager/public/graph/graph.h"
+#include "components/performance_manager/public/performance_manager.h"
 
 namespace performance_manager {
 
@@ -32,6 +34,9 @@ namespace performance_manager {
 //   foo->DoSomething();
 // }
 //
+// GraphRegisteredImpl also provides a GetFromGraph() convenience helper that
+// wraps `graph->GetRegisteredObjectAs<...>()`.
+//
 // This may easily (and commonly) be combined with GraphOwned, allowing the
 // registered object to be owned by the graph as well. The
 // GraphOwnedAndRegistered helper will register the object automatically when
@@ -104,8 +109,17 @@ class GraphRegisteredImpl : public GraphRegistered {
   uintptr_t GetTypeId() const override { return TypeId(); }
 
   // Helper function for looking up the registered object of this type from the
-  // provided graph. Syntactic sugar for "Graph::GetRegisteredObjectAs".
-  static SelfType* GetFromGraph(Graph* graph) {
+  // provided `graph`. If none is provided, looks up the object in the default
+  // graph returned by PerformanceManager::GetGraph(), returning nullptr if
+  // the default graph is not available.
+  static SelfType* GetFromGraph(Graph* graph = nullptr) {
+    if (!graph) {
+      if (!PerformanceManager::IsAvailable()) {
+        return nullptr;
+      }
+      graph = PerformanceManager::GetGraph();
+      CHECK(graph);
+    }
     return graph->GetRegisteredObjectAs<SelfType>();
   }
 
@@ -146,4 +160,4 @@ class GraphOwnedAndRegistered : public GraphRegisteredImpl<SelfType>,
 
 }  // namespace performance_manager
 
-#endif  // COMPONENTS_PERFORMANCE_MANAGER_PUBLIC_GRAPH_GRAPH_REGISTERED_H_
+#endif  // COMPONENTS_PERFORMANCE_MANAGER_PUBLIC_GRAPH_GRAPH_REGISTERED_H_