[llvm] b15aa7f - [ORC] Add unit test for MemoryFlags APIs, don't dereference end() iterator.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2024-07-22T17:16:32+10:00
New Revision: b15aa7f88c21d126bac4cf77eb05c522e0f1870d

URL: https://github.com/llvm/llvm-project/commit/b15aa7f88c21d126bac4cf77eb05c522e0f1870d
DIFF: https://github.com/llvm/llvm-project/commit/b15aa7f88c21d126bac4cf77eb05c522e0f1870d.diff

LOG: [ORC] Add unit test for MemoryFlags APIs, don't dereference end() iterator.

In AllocGroupSmallMap::find(AllocGroup) we were calling lower_bound(...) and
then unconditionally dereferencing the resulting iterator, however
lower_bound(...) may return end() if the value being searched for is higher
than any value present in the map. This patch adds a check for end() before
the dereference to guard against dereference of end().

This commit also adds some basic unit tests for MemProt and AllocGroupSmallMap.

rdar://129662981

Added: 
    llvm/unittests/ExecutionEngine/Orc/MemoryFlagsTest.cpp

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h
    llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h
index b8b5f90b6b0fb..fa522ff88e1db 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h
@@ -184,7 +184,7 @@ template <typename T> class AllocGroupSmallMap {
   iterator end() { return Elems.end(); }
   iterator find(AllocGroup G) {
     auto I = lower_bound(Elems, G, compareKey);
-    return (I->first == G) ? I : end();
+    return (I == end() || I->first == G) ? I : end();
   }
 
   bool empty() const { return Elems.empty(); }

diff  --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index af431658c9b7d..cc868e7587dc6 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -28,6 +28,7 @@ add_llvm_unittest(OrcJITTests
   LookupAndRecordAddrsTest.cpp
   MachOPlatformTest.cpp
   MapperJITLinkMemoryManagerTest.cpp
+  MemoryFlagsTest.cpp
   MemoryMapperTest.cpp
   ObjectFormatsTest.cpp
   ObjectLinkingLayerTest.cpp

diff  --git a/llvm/unittests/ExecutionEngine/Orc/MemoryFlagsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MemoryFlagsTest.cpp
new file mode 100644
index 0000000000000..270ea2f442b9e
--- /dev/null
+++ b/llvm/unittests/ExecutionEngine/Orc/MemoryFlagsTest.cpp
@@ -0,0 +1,55 @@
+//===------- MemoryFlagsTest.cpp - Test MemoryFlags and related APIs ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h"
+#include "gtest/gtest.h"
+
+#include <future>
+
+using namespace llvm;
+using namespace llvm::orc;
+
+TEST(MemProtTest, Basics) {
+  MemProt MPNone = MemProt::None;
+
+  EXPECT_EQ(MPNone & MemProt::Read, MemProt::None);
+  EXPECT_EQ(MPNone & MemProt::Write, MemProt::None);
+  EXPECT_EQ(MPNone & MemProt::Exec, MemProt::None);
+
+  EXPECT_EQ(MPNone | MemProt::Read, MemProt::Read);
+  EXPECT_EQ(MPNone | MemProt::Write, MemProt::Write);
+  EXPECT_EQ(MPNone | MemProt::Exec, MemProt::Exec);
+
+  MemProt MPAll = MemProt::Read | MemProt::Write | MemProt::Exec;
+  EXPECT_EQ(MPAll & MemProt::Read, MemProt::Read);
+  EXPECT_EQ(MPAll & MemProt::Write, MemProt::Write);
+  EXPECT_EQ(MPAll & MemProt::Exec, MemProt::Exec);
+}
+
+TEST(AllocGroupSmallMap, EmptyMap) {
+  AllocGroupSmallMap<bool> EM;
+  EXPECT_TRUE(EM.empty());
+  EXPECT_EQ(EM.size(), 0u);
+}
+
+TEST(AllocGroupSmallMap, NonEmptyMap) {
+  AllocGroupSmallMap<unsigned> NEM;
+  NEM[MemProt::Read] = 42;
+
+  EXPECT_FALSE(NEM.empty());
+  EXPECT_EQ(NEM.size(), 1U);
+  EXPECT_EQ(NEM[MemProt::Read], 42U);
+  EXPECT_EQ(NEM.find(MemProt::Read), NEM.begin());
+  EXPECT_EQ(NEM.find(MemProt::Read | MemProt::Write), NEM.end());
+
+  NEM[MemProt::Read | MemProt::Write] = 7;
+  EXPECT_EQ(NEM.size(), 2U);
+  EXPECT_EQ(NEM.begin()->second, 42U);
+  EXPECT_EQ((NEM.begin() + 1)->second, 7U);
+}
+