[llvm] [orc-rt] Return ref from Session::addService, add createService. (PR #186640)

[Lang Hames via llvm-commits]

https://github.com/lhames created https://github.com/llvm/llvm-project/pull/186640

Session::addService now returns a reference to the added Service. This allows clients to hold a reference for further direct interaction with the Service object.

This commit also introduces a new Session::createService convenience method that creates the service and returns a reference to it.

>From a8f7438bf738903a1d1e707debc296e2bd82ba1f Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Sun, 15 Mar 2026 14:39:32 +1100
Subject: [PATCH] [orc-rt] Return ref from Session::addService, add
 createService.

Session::addService now returns a reference to the added Service. This allows
clients to hold a reference for further direct interaction with the Service
object.

This commit also introduces a new Session::createService convenience method
that creates the service and returns a reference to it.
---
 orc-rt/include/orc-rt/Session.h  | 17 ++++++++++++++++-
 orc-rt/lib/executor/Session.cpp  |  6 ------
 orc-rt/unittests/SessionTest.cpp | 27 +++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/orc-rt/include/orc-rt/Session.h b/orc-rt/include/orc-rt/Session.h
index 258a63945c98c..f6848dae09c1a 100644
--- a/orc-rt/include/orc-rt/Session.h
+++ b/orc-rt/include/orc-rt/Session.h
@@ -143,7 +143,22 @@ class Session {
   void waitForShutdown();
 
   /// Add a Service to the session.
-  void addService(std::unique_ptr<Service> Srv);
+  template <typename ServiceT>
+  ServiceT &addService(std::unique_ptr<ServiceT> Srv) {
+    assert(Srv && "addService called with null value");
+    ServiceT &Ref = *Srv;
+    std::scoped_lock<std::mutex> Lock(M);
+    assert(!SI && "addService called after shutdown");
+    Services.push_back(std::move(Srv));
+    return Ref;
+  }
+
+  /// Construct an instance of ServiceT from the given arguments and add it to
+  /// the Session.
+  template <typename ServiceT, typename... ArgTs>
+  ServiceT &createService(ArgTs &&...Args) {
+    return addService(std::make_unique<ServiceT>(std::forward<ArgTs>(Args)...));
+  }
 
   /// Set the ControllerAccess object.
   void setController(std::shared_ptr<ControllerAccess> CA);
diff --git a/orc-rt/lib/executor/Session.cpp b/orc-rt/lib/executor/Session.cpp
index fe023064073a1..f8b66c1224245 100644
--- a/orc-rt/lib/executor/Session.cpp
+++ b/orc-rt/lib/executor/Session.cpp
@@ -63,12 +63,6 @@ void Session::waitForShutdown() {
   F.get();
 }
 
-void Session::addService(std::unique_ptr<Service> Srv) {
-  std::scoped_lock<std::mutex> Lock(M);
-  assert(!SI && "addService called after shutdown");
-  Services.push_back(std::move(Srv));
-}
-
 void Session::setController(std::shared_ptr<ControllerAccess> CA) {
   assert(CA && "Cannot attach null controller");
   std::scoped_lock<std::mutex> Lock(M);
diff --git a/orc-rt/unittests/SessionTest.cpp b/orc-rt/unittests/SessionTest.cpp
index ca1a25b0928d3..ecf632d4eb4e4 100644
--- a/orc-rt/unittests/SessionTest.cpp
+++ b/orc-rt/unittests/SessionTest.cpp
@@ -55,6 +55,21 @@ class MockService : public Service {
   move_only_function<Error(Op)> GenResult;
 };
 
+class ConfigurableService : public Service {
+public:
+  ConfigurableService(int ConstructorOption) {}
+
+  void onDetach(OnCompleteFn OnComplete) override {
+    OnComplete(Error::success());
+  }
+
+  void onShutdown(OnCompleteFn OnComplete) override {
+    OnComplete(Error::success());
+  }
+
+  void doMoreConfig(int) noexcept {}
+};
+
 class NoDispatcher : public TaskDispatcher {
 public:
   void dispatch(std::unique_ptr<Task> T) override {
@@ -362,6 +377,18 @@ TEST(SessionTest, ExpectedShutdownSequence) {
   EXPECT_TRUE(SessionShutdownComplete);
 }
 
+TEST(SessionTest, AddServiceAndUseRef) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  auto &CS = S.addService(std::make_unique<ConfigurableService>(42));
+  CS.doMoreConfig(1);
+}
+
+TEST(SessionTest, CreateServiceAndUseRef) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  auto &CS = S.createService<ConfigurableService>(42);
+  CS.doMoreConfig(1);
+}
+
 TEST(ControllerAccessTest, Basics) {
   // Test that we can set the ControllerAccess implementation and still shut
   // down as expected.