Merge pull request #71218 from mikeash/prebuilt-string-map-nul-termination-fix

[Runtime] Fix PrebuiltStringMap find with non-terminated keys.

Author: mikeash

include/swift/Runtime/PrebuiltStringMap.h

@@ -160,7 +160,10 @@ struct PrebuiltStringMap {
 
     size_t numSearched = 0;
     while (const char *key = array()[index].key) {
-      if (strcmp(key, toFind) == 0)
+      // key is NUL terminated but toFind may not be. Check that they have equal
+      // contents up to len, and check that key has a terminating NUL at the
+      // right point.
+      if (strncmp(key, toFind, len) == 0 && key[len] == 0)
         return &array()[index];
 
       index = index + 1;

unittests/runtime/PrebuiltStringMap.cpp

@@ -16,40 +16,57 @@
 static bool stringIsNull(const char *str) { return str == nullptr; }
 
 TEST(PrebuiltStringMapTest, basic) {
-  unsigned testEntryCount = 1000;
-
-  std::vector<std::pair<std::string, unsigned>> testVector;
-  testVector.reserve(testEntryCount);
-  for (unsigned i = 0; i < testEntryCount; i++) {
-    std::string key;
-    for (unsigned j = 0; j < i; j++)
-      key += std::to_string(j);
-    testVector.push_back({key, i});
-  }
-
-  using Map = swift::PrebuiltStringMap<const char *, unsigned, stringIsNull>;
-
-  unsigned mapSize = testEntryCount * 4 / 3;
-  void *mapAllocation = calloc(1, Map::byteSize(mapSize));
-  Map *map = new (mapAllocation) Map(mapSize);
-
-  for (auto &[key, value] : testVector) {
-    const char *keyCStr = key.c_str();
-    auto *element = map->insert(keyCStr);
-    EXPECT_NE(element, nullptr);
-
-    element->key = keyCStr;
-    element->value = value;
-  }
-
-  for (auto &[key, value] : testVector) {
-    const char *keyCStr = key.c_str();
-    auto *element = map->find(keyCStr);
-    EXPECT_NE(element, nullptr);
-
-    EXPECT_EQ(element->key, key);
-    EXPECT_EQ(element->value, value);
-  }
-
-  free(mapAllocation);
+  auto testOnce = [&](unsigned testEntryCount) {
+    std::vector<std::pair<std::string, unsigned>> testVector;
+    testVector.reserve(testEntryCount);
+    for (unsigned i = 0; i < testEntryCount; i++) {
+      std::string key;
+      for (unsigned j = 0; j < i; j++)
+        key += std::to_string(j);
+      testVector.push_back({key, i});
+    }
+
+    using Map = swift::PrebuiltStringMap<const char *, unsigned, stringIsNull>;
+
+    unsigned mapSize = testEntryCount * 4 / 3;
+    void *mapAllocation = calloc(1, Map::byteSize(mapSize));
+    Map *map = new (mapAllocation) Map(mapSize);
+
+    for (auto &[key, value] : testVector) {
+      const char *keyCStr = key.c_str();
+      auto *element = map->insert(keyCStr);
+      EXPECT_NE(element, nullptr);
+
+      element->key = keyCStr;
+      element->value = value;
+    }
+
+    for (auto &[key, value] : testVector) {
+      const char *keyCStr = key.c_str();
+      auto *element = map->find(keyCStr);
+      EXPECT_NE(element, nullptr);
+
+      EXPECT_EQ(element->key, key);
+      EXPECT_EQ(element->value, value);
+    }
+
+    // Test the interface that takes a string and length without NUL
+    // termination.
+    for (auto &[key, value] : testVector) {
+      auto keyCopy = key;
+      keyCopy += "xyz";
+      const char *keyCStr = keyCopy.c_str();
+      auto *element = map->find(keyCStr, key.size());
+      EXPECT_NE(element, nullptr);
+
+      EXPECT_EQ(element->key, key);
+      EXPECT_EQ(element->value, value);
+    }
+
+    free(mapAllocation);
+  };
+
+  testOnce(10);
+  testOnce(100);
+  testOnce(1000);
 }