container regex matching

facebookexperimental · Dec 21, 2022 · 2dc5479 · 2dc5479
1 parent abec0dc
commit 2dc5479
Show file tree

Hide file tree

Showing 7 changed files with 17 additions and 12 deletions.
diff --git a/src/ContainerInfo.cpp b/src/ContainerInfo.cpp
@@ -72,6 +72,14 @@ std::unique_ptr<ContainerInfo> ContainerInfo::loadFromFile(
     return nullptr;
   }
 
+  std::regex matcher;
+  if (std::optional<std::string> str =
+          (*info)["matcher"].value<std::string>()) {
+    matcher = std::regex(*str, std::regex_constants::grep);
+  } else {
+    matcher = std::regex("^" + typeName, std::regex_constants::grep);
+  }
+
   std::optional<size_t> numTemplateParams =
       (*info)["numTemplateParams"].value<size_t>();
 
@@ -122,6 +130,7 @@ std::unique_ptr<ContainerInfo> ContainerInfo::loadFromFile(
 
   return std::unique_ptr<ContainerInfo>(new ContainerInfo{
       std::move(typeName),
+      std::move(matcher),
       numTemplateParams,
       ctype,
       std::move(header),

diff --git a/src/ContainerInfo.h b/src/ContainerInfo.h
@@ -16,6 +16,7 @@
 #pragma once
 #include <filesystem>
 #include <optional>
+#include <regex>
 #include <string>
 #include <vector>
 
@@ -73,6 +74,7 @@ const char *containerTypeEnumToStr(ContainerTypeEnum ty);
 
 struct ContainerInfo {
   std::string typeName;
+  std::regex matcher;
   std::optional<size_t> numTemplateParams;
   ContainerTypeEnum ctype = UNKNOWN_TYPE;
   std::string header;

diff --git a/src/OICodeGen.cpp b/src/OICodeGen.cpp
@@ -166,20 +166,11 @@ std::optional<ContainerInfo> OICodeGen::getContainerInfo(
     return std::nullopt;
   }
 
+  std::string nameStr = std::string(*name);
   for (auto it = containerInfoList.rbegin(); it != containerInfoList.rend();
        it++) {
     const auto &info = *it;
-    if (name->starts_with(info->typeName)) {
-      // Blob must match exactly. Otherwise it also matches BlobsMap
-      if (info->ctype == CAFFE2_BLOB_TYPE && *name != "caffe2::Blob") {
-        return std::nullopt;
-      }
-
-      // IOBuf must also match exactly, or it also matches IOBufQueue
-      if (info->ctype == FOLLY_IOBUF_TYPE && *name != "folly::IOBuf") {
-        continue;
-      }
-
+    if (std::regex_search(nameStr, info->matcher)) {
       return *info;
     }
   }

diff --git a/src/Serialize.h b/src/Serialize.h
@@ -39,7 +39,7 @@
   BOOST_CLASS_VERSION(Type, version)
 
 DEFINE_TYPE_VERSION(PaddingInfo, 120, 3)
-DEFINE_TYPE_VERSION(ContainerInfo, 168, 4)
+DEFINE_TYPE_VERSION(ContainerInfo, 200, 5)
 DEFINE_TYPE_VERSION(struct drgn_location_description, 32, 2)
 DEFINE_TYPE_VERSION(struct drgn_object_locator, 72, 2)
 DEFINE_TYPE_VERSION(FuncDesc::Arg, 128, 2)

diff --git a/types/caffe2_blob_type.toml b/types/caffe2_blob_type.toml
@@ -1,5 +1,6 @@
 [info]
 typeName = "caffe2::Blob"
+matcher = "^caffe2::Blob$"
 numTemplateParams = 0
 ctype = "CAFFE2_BLOB_TYPE"
 header = "caffe2/core/blob.h"

diff --git a/types/folly_iobuf_queue_type.toml b/types/folly_iobuf_queue_type.toml
@@ -1,5 +1,6 @@
 [info]
 typeName = "folly::IOBufQueue"
+matcher = "^folly::IOBufQueue$"
 numTemplateParams = 0
 ctype = "FOLLY_IOBUFQUEUE_TYPE"
 header = "folly/io/IOBufQueue.h"

diff --git a/types/folly_iobuf_type.toml b/types/folly_iobuf_type.toml
@@ -1,5 +1,6 @@
 [info]
 typeName = "folly::IOBuf"
+matcher = "^folly::IOBuf$"
 numTemplateParams = 0
 ctype = "FOLLY_IOBUF_TYPE"
 header = "folly/io/IOBuf.h"