From 4c6f23276684cef34418de134b759e1d720288d3 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Mon, 9 Oct 2023 13:39:33 -0700
Subject: [PATCH] containers: add required features (#374)

Summary:
Adds the option for required features to container definitions. These cause the container not to be passed to `DrgnParser` if that feature is not enabled during code generation. The thrift isset type does not currently work with `tree-builder-v2` and only provides benefit with `capture-thrift-isset`. This change makes sure the container is ignored if it won't be useful, allowing code generation under `tree-builder-v2`.


Test Plan: - CI

Differential Revision: D49960512

Pulled By: JakeHillion
---
 oi/CMakeLists.txt                  |  5 ++++-
 oi/CodeGen.cpp                     |  4 ++++
 oi/ContainerInfo.cpp               | 13 +++++++++++++
 oi/ContainerInfo.h                 |  2 ++
 oi/EnumBitset.h                    | 15 +++++++++++++++
 test/integration/thrift_isset.toml |  1 -
 types/README.md                    |  7 +++++++
 types/thrift_isset_type.toml       |  1 +
 8 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/oi/CMakeLists.txt b/oi/CMakeLists.txt
index 1751d23..3754f41 100644
--- a/oi/CMakeLists.txt
+++ b/oi/CMakeLists.txt
@@ -26,17 +26,20 @@ target_link_libraries(symbol_service
   dw
 )
 
+add_library(features Features.cpp)
+target_link_libraries(features glog::glog)
+
 add_library(container_info
   ContainerInfo.cpp
 )
 target_link_libraries(container_info
+  features
   glog::glog
   toml
 )
 
 add_library(codegen
   CodeGen.cpp
-  Features.cpp
   FuncGen.cpp
   OICodeGen.cpp
 )
diff --git a/oi/CodeGen.cpp b/oi/CodeGen.cpp
index b7b113e..19f8de5 100644
--- a/oi/CodeGen.cpp
+++ b/oi/CodeGen.cpp
@@ -1119,6 +1119,10 @@ bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
 
 void CodeGen::registerContainer(const fs::path& path) {
   auto info = std::make_unique<ContainerInfo>(path);
+  if (info->requiredFeatures != (config_.features & info->requiredFeatures)) {
+    VLOG(1) << "Skipping container (feature conflict): " << info->typeName;
+    return;
+  }
   VLOG(1) << "Registered container: " << info->typeName;
   containerInfos_.emplace_back(std::move(info));
 }
diff --git a/oi/ContainerInfo.cpp b/oi/ContainerInfo.cpp
index 7d3d9ff..a5703be 100644
--- a/oi/ContainerInfo.cpp
+++ b/oi/ContainerInfo.cpp
@@ -242,6 +242,19 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
 
   underlyingContainerIndex = info["underlying_container_index"].value<size_t>();
 
+  if (toml::array* arr = info["required_features"].as_array()) {
+    arr->for_each([&](auto&& el) {
+      if constexpr (toml::is_string<decltype(el)>) {
+        oi::detail::Feature f = oi::detail::featureFromStr(*el);
+        if (f == oi::detail::Feature::UnknownFeature) {
+          LOG(WARNING) << "unknown feature in container config: " << el;
+          return;
+        }
+        requiredFeatures[f] = true;
+      }
+    });
+  }
+
   if (!container["codegen"].is_table()) {
     throw ContainerInfoError(
         path, "a container info file requires a `codegen` table");
diff --git a/oi/ContainerInfo.h b/oi/ContainerInfo.h
index 776ab57..0ad9da9 100644
--- a/oi/ContainerInfo.h
+++ b/oi/ContainerInfo.h
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "oi/ContainerTypeEnum.h"
+#include "oi/Features.h"
 
 ContainerTypeEnum containerTypeEnumFromStr(std::string& str);
 const char* containerTypeEnumToStr(ContainerTypeEnum ty);
@@ -94,6 +95,7 @@ struct ContainerInfo {
   std::optional<size_t> underlyingContainerIndex{};
   std::vector<size_t> stubTemplateParams{};
   bool captureKeys = false;
+  oi::detail::FeatureSet requiredFeatures;
 
   Codegen codegen;
 
diff --git a/oi/EnumBitset.h b/oi/EnumBitset.h
index 96b798b..d9828ee 100644
--- a/oi/EnumBitset.h
+++ b/oi/EnumBitset.h
@@ -47,11 +47,26 @@ class EnumBitset {
     return bitset.none();
   }
 
+  bool operator==(const EnumBitset<T, N>& that) const {
+    return bitset == that.bitset;
+  }
   EnumBitset<T, N>& operator|=(const EnumBitset<T, N>& that) {
     bitset |= that.bitset;
     return *this;
   }
+  EnumBitset<T, N>& operator&=(const EnumBitset<T, N>& that) {
+    bitset &= that.bitset;
+    return *this;
+  }
 
  private:
   BitsetType bitset;
 };
+
+template <typename T, size_t N>
+EnumBitset<T, N> operator&(const EnumBitset<T, N>& lhs,
+                           const EnumBitset<T, N>& rhs) {
+  auto out = lhs;
+  out &= rhs;
+  return out;
+}
diff --git a/test/integration/thrift_isset.toml b/test/integration/thrift_isset.toml
index e32f416..064cb13 100644
--- a/test/integration/thrift_isset.toml
+++ b/test/integration/thrift_isset.toml
@@ -213,7 +213,6 @@ namespace cpp2 {
       ]}]'''
 
   [cases.no_capture]
-    oil_skip = 'oil does not support thrift isset yet' # https://github.com/facebookexperimental/object-introspection/issues/296
     param_types = ["const cpp2::MyThriftStructBoxed&"]
     setup = '''
       cpp2::MyThriftStructBoxed ret;
diff --git a/types/README.md b/types/README.md
index 1b70b4b..ebde86c 100644
--- a/types/README.md
+++ b/types/README.md
@@ -22,6 +22,13 @@ This document describes the format of the container definition files contained i
   Only used for container adapters. Points OI to the template parameter
   representing the underlying container to be measured.
 
+- `required_features`
+
+  A set of feature names such as `tree-builder-v2` which must be enabled for
+  this container description to be included. Currently only supported with
+  CodeGen v2 as that's their only use case and an implementation for CodeGen v1
+  would be untested.
+
 ### codegen
 - `decl`
 
diff --git a/types/thrift_isset_type.toml b/types/thrift_isset_type.toml
index edbc118..213cb58 100644
--- a/types/thrift_isset_type.toml
+++ b/types/thrift_isset_type.toml
@@ -2,6 +2,7 @@
 type_name = "apache::thrift::detail::isset_bitset"
 ctype = "THRIFT_ISSET_TYPE"
 header = "thrift/lib/cpp2/gen/module_types_h.h"
+required_features = ["capture-thrift-isset"]
 
 # Old:
 typeName = "apache::thrift::detail::isset_bitset<"