From f68743f20ef83ddfc2d7b24cc6749c6c9d67cf2e Mon Sep 17 00:00:00 2001
From: Tianyu Liu <kingcrimsontianyu@gmail.com>
Date: Wed, 9 Oct 2024 01:31:52 -0400
Subject: [PATCH] Improve use of map

---
 cpp/include/kvikio/posix_io.hpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/cpp/include/kvikio/posix_io.hpp b/cpp/include/kvikio/posix_io.hpp
index e044ab0bca..c62d43b688 100644
--- a/cpp/include/kvikio/posix_io.hpp
+++ b/cpp/include/kvikio/posix_io.hpp
@@ -60,12 +60,14 @@ class StreamsByThread {
     auto key = std::make_pair(ctx, thd_id);
 
     // Create a new stream if `ctx` doesn't have one.
-    if (_instance._streams.find(key) == _instance._streams.end()) {
+    if (auto search = _instance._streams.find(key); search == _instance._streams.end()) {
       CUstream stream{};
       CUDA_DRIVER_TRY(cudaAPI::instance().StreamCreate(&stream, CU_STREAM_DEFAULT));
       _instance._streams[key] = stream;
+      return stream;
+    } else {
+      return search->second;
     }
-    return _instance._streams.at(key);
   }
 
   static CUstream get()