From 3de8bac6d773197df6ed3a29faa0531c3c1fff91 Mon Sep 17 00:00:00 2001
From: Rafal Rudnicki <rafal.rudnicki@intel.com>
Date: Fri, 17 Jan 2025 18:48:53 +0100
Subject: [PATCH] fix parseDisjointPoolConfig and add tests

---
 .../umf_pools/disjoint_pool_config_parser.cpp | 57 +++++++------------
 test/usm/usmPoolManager.cpp                   | 47 +++++++++++++++
 2 files changed, 68 insertions(+), 36 deletions(-)

diff --git a/source/common/umf_pools/disjoint_pool_config_parser.cpp b/source/common/umf_pools/disjoint_pool_config_parser.cpp
index 5a5329ff17..1bc2c37d7d 100644
--- a/source/common/umf_pools/disjoint_pool_config_parser.cpp
+++ b/source/common/umf_pools/disjoint_pool_config_parser.cpp
@@ -177,46 +177,31 @@ DisjointPoolAllConfigs parseDisjointPoolConfig(const std::string &config,
     };
 
     size_t MaxSize = (std::numeric_limits<size_t>::max)();
+    size_t EnableBuffers = 1;
 
     // Update pool settings if specified in environment.
-    size_t EnableBuffers = 1;
-    if (config != "") {
-        std::string Params = config;
-        size_t Pos = Params.find(';');
-        if (Pos != std::string::npos) {
-            if (Pos > 0) {
-                GetValue(Params, Pos, EnableBuffers);
-            }
-            Params.erase(0, Pos + 1);
-            size_t Pos = Params.find(';');
-            if (Pos != std::string::npos) {
-                if (Pos > 0) {
-                    GetValue(Params, Pos, MaxSize);
-                }
-                Params.erase(0, Pos + 1);
-                do {
-                    size_t Pos = Params.find(';');
-                    if (Pos != std::string::npos) {
-                        if (Pos > 0) {
-                            std::string MemParams = Params.substr(0, Pos);
-                            MemTypeParser(MemParams);
-                        }
-                        Params.erase(0, Pos + 1);
-                        if (Params.size() == 0) {
-                            break;
-                        }
-                    } else {
-                        MemTypeParser(Params);
-                        break;
-                    }
-                } while (true);
-            } else {
-                // set MaxPoolSize for all configs
-                GetValue(Params, Params.size(), MaxSize);
-            }
+    bool EnableBuffersSet = false;
+    bool MaxSizeSet = false;
+    size_t Start = 0;
+    size_t End = config.find(';');
+    while (true) {
+        std::string Param = config.substr(Start, End - Start);
+        if (!EnableBuffersSet && isdigit(Param[0])) {
+            GetValue(Param, Param.size(), EnableBuffers);
+            EnableBuffersSet = true;
+        } else if (!MaxSizeSet && isdigit(Param[0])) {
+            GetValue(Param, Param.size(), MaxSize);
+            MaxSizeSet = true;
         } else {
-            GetValue(Params, Params.size(), EnableBuffers);
+            MemTypeParser(Param);
         }
+
+        if (End == std::string::npos) {
+            break;
+        }
+
+        Start = End + 1;
+        End = config.find(';', Start);
     }
 
     AllConfigs.EnableBuffers = EnableBuffers;
diff --git a/test/usm/usmPoolManager.cpp b/test/usm/usmPoolManager.cpp
index ad1ca6e093..b6c4b46bb4 100644
--- a/test/usm/usmPoolManager.cpp
+++ b/test/usm/usmPoolManager.cpp
@@ -3,6 +3,7 @@
 // See LICENSE.TXT
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "umf_pools/disjoint_pool_config_parser.hpp"
 #include "ur_pool_manager.hpp"
 
 #include <uur/fixtures.h>
@@ -18,6 +19,27 @@ auto createMockPoolHandle() {
         [](umf_memory_pool_t *) {});
 }
 
+bool compareConfig(const usm::umf_disjoint_pool_config_t &left,
+                   usm::umf_disjoint_pool_config_t &right) {
+    return left.MaxPoolableSize == right.MaxPoolableSize &&
+           left.Capacity == right.Capacity &&
+           left.SlabMinSize == right.SlabMinSize;
+}
+
+bool compareConfigs(const usm::DisjointPoolAllConfigs &left,
+                    usm::DisjointPoolAllConfigs &right) {
+    return left.EnableBuffers == right.EnableBuffers &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Host],
+                         right.Configs[usm::DisjointPoolMemType::Host]) &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Device],
+                         right.Configs[usm::DisjointPoolMemType::Device]) &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Shared],
+                         right.Configs[usm::DisjointPoolMemType::Shared]) &&
+           compareConfig(
+               left.Configs[usm::DisjointPoolMemType::SharedReadOnly],
+               right.Configs[usm::DisjointPoolMemType::SharedReadOnly]);
+}
+
 TEST_P(urUsmPoolDescriptorTest, poolIsPerContextTypeAndDevice) {
     auto &devices = uur::DevicesEnvironment::instance->devices;
 
@@ -111,4 +133,29 @@ TEST_P(urUsmPoolManagerTest, poolManagerGetNonexistant) {
     }
 }
 
+TEST_P(urUsmPoolManagerTest, config) {
+    // Check default config
+    usm::DisjointPoolAllConfigs def;
+    usm::DisjointPoolAllConfigs parsed1 =
+        usm::parseDisjointPoolConfig("1;host:2M,4,64K;device:4M,4,64K;"
+                                     "shared:0,0,2M;read_only_shared:4M,4,2M",
+                                     0);
+    ASSERT_EQ(compareConfigs(def, parsed1), true);
+
+    // Check partially set config
+    usm::DisjointPoolAllConfigs parsed2 =
+        usm::parseDisjointPoolConfig("1;device:4M;shared:0,0,2M", 0);
+    ASSERT_EQ(compareConfigs(def, parsed2), true);
+
+    // Check non-default config
+    usm::DisjointPoolAllConfigs test(def);
+    test.Configs[usm::DisjointPoolMemType::Shared].MaxPoolableSize = 128 * 1024;
+    test.Configs[usm::DisjointPoolMemType::Shared].Capacity = 4;
+    test.Configs[usm::DisjointPoolMemType::Shared].SlabMinSize = 64 * 1024;
+
+    usm::DisjointPoolAllConfigs parsed3 =
+        usm::parseDisjointPoolConfig("1;shared:128K,4,64K", 0);
+    ASSERT_EQ(compareConfigs(test, parsed3), true);
+}
+
 UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urUsmPoolManagerTest);