From 2dc5479c32d44b305a857b135c06aace5f23e3c4 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Tue, 20 Dec 2022 04:54:40 -0800
Subject: [PATCH] container regex matching

---
 src/ContainerInfo.cpp             |  9 +++++++++
 src/ContainerInfo.h               |  2 ++
 src/OICodeGen.cpp                 | 13 ++-----------
 src/Serialize.h                   |  2 +-
 types/caffe2_blob_type.toml       |  1 +
 types/folly_iobuf_queue_type.toml |  1 +
 types/folly_iobuf_type.toml       |  1 +
 7 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/ContainerInfo.cpp b/src/ContainerInfo.cpp
index 75b9924..31b0985 100644
--- 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
index 28ddd4b..7e470b1 100644
--- 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
index 4af46fd..68ec556 100644
--- 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
index bae2c1b..c7181ec 100644
--- 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
index 283c18b..09a6ea7 100644
--- 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
index 799dd05..7f795bb 100644
--- 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
index f463add..5326f36 100644
--- 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"