From 35d45c2b4f5a25ba12ab902f89d433ab451faf4e Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Wed, 27 Sep 2023 13:46:19 -0700
Subject: [PATCH] folly::sorted_vector_map: Add key-capture support

---
 test/integration/capture_keys.toml |  9 +++++++++
 types/map_seq_type.toml            | 21 +++++++++++++++------
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/test/integration/capture_keys.toml b/test/integration/capture_keys.toml
index 762dbab..c2f9fe2 100644
--- a/test/integration/capture_keys.toml
+++ b/test/integration/capture_keys.toml
@@ -329,3 +329,12 @@ class FixedAllocator {
       top_level = true
     '''
     expect_json_v2 = '[{"members": [{"data": 1}]}]'
+  [cases.folly_sorted_vector_map]
+    oid_skip = "Requires TreeBuilderV2"
+    param_types = ["const folly::sorted_vector_map<int, int>&"]
+    setup = "return {{{1,2}}};"
+    config_suffix = '''
+      [[codegen.capture_keys]]
+      top_level = true
+    '''
+    expect_json_v2 = '[{"members": [{"data": 1}]}]'
diff --git a/types/map_seq_type.toml b/types/map_seq_type.toml
index e19e14d..16b9615 100644
--- a/types/map_seq_type.toml
+++ b/types/map_seq_type.toml
@@ -74,7 +74,8 @@ traversal_func = '''
 
     for (const auto& kv : container) {
       tail = tail.delegate([&kv](auto ret) {
-        auto next = ret.delegate([&kv](typename TypeHandler<DB, T0>::type ret) {
+        auto start = maybeCaptureKey<captureKeys, DB, T0>(kv.first, ret);
+        auto next = start.delegate([&kv](typename TypeHandler<DB, T0>::type ret) {
           return OIInternal::getSizeType<DB>(kv.first, ret);
         });
         return OIInternal::getSizeType<DB>(kv.second, next);
@@ -95,10 +96,17 @@ el.container_stats.emplace(result::Element::ContainerStats{ .capacity = std::get
 '''
 
 [[codegen.processor]]
-type = '''types::st::List<DB, types::st::Pair<DB,
-  typename TypeHandler<DB, T0>::type,
-  typename TypeHandler<DB, T1>::type
->>'''
+type = '''
+std::conditional_t<captureKeys,
+  types::st::List<DB, types::st::Pair<DB,
+    typename CaptureKeyHandler<DB, T0>::type,
+    types::st::Pair<DB,
+      typename TypeHandler<DB, T0>::type,
+      typename TypeHandler<DB, T1>::type>>>,
+  types::st::List<DB, types::st::Pair<DB,
+    typename TypeHandler<DB, T0>::type,
+    typename TypeHandler<DB, T1>::type>>>
+'''
 func = '''
 using element_type = std::pair<T0, T1>;
 
@@ -106,13 +114,14 @@ static constexpr std::array<inst::Field, 2> entryFields{
   make_field<DB, T0>("key"),
   make_field<DB, T1>("value"),
 };
+static constexpr auto processors = maybeCaptureKeysProcessor<captureKeys, DB, T0>();
 static constexpr auto entry = inst::Field {
   sizeof(element_type),
   sizeof(element_type) - sizeof(T0) - sizeof(T1),
   "[]",
   std::array<std::string_view, 0>{},
   entryFields,
-  std::array<inst::ProcessorInst, 0>{},
+  processors,
 };
 
 auto list = std::get<ParsedData::List>(d.val);