From e431c00c9f61ec577b0b98ad1c620a1a9ef64536 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Wed, 20 Dec 2023 15:46:51 +0000
Subject: [PATCH] tbv2: implement folly::IOBuf

folly::IOBuf does not have TreeBuilder v2 container support. Add it.

The implementation is a direct clone of v1. It still lacks tests.

Test Plan:
- It codegens on a prod type.
- No runtime testing... Bad form, I know.
- Issue created to add integration tests: https://github.com/facebookexperimental/object-introspection/issues/436
---
 types/folly_iobuf_type.toml | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/types/folly_iobuf_type.toml b/types/folly_iobuf_type.toml
index c7e1043..6b88d5c 100644
--- a/types/folly_iobuf_type.toml
+++ b/types/folly_iobuf_type.toml
@@ -39,3 +39,34 @@ void getSizeType(const %1% &container, size_t& returnArg)
     SAVE_SIZE(fullCapacity);
 }
 """
+
+traversal_func = """
+// We calculate the length of all IOBufs in the chain manually.
+// IOBuf has built-in computeChainCapacity()/computeChainLength()
+// functions which do this for us. But dead code optimization in TAO
+// caused these functions to be removed which causes relocation
+// errors.
+
+std::size_t fullLength = container.length();
+std::size_t fullCapacity = container.capacity();
+for (const folly::IOBuf* current = container.next(); current != &container;
+    current = current->next()) {
+  fullLength += current->length();
+  fullCapacity += current->capacity();
+}
+
+return returnArg.write(fullCapacity).write(fullLength);
+"""
+
+[[codegen.processor]]
+type = "types::st::VarInt<DB>"
+func = """
+el.container_stats.emplace(result::Element::ContainerStats{ .capacity = std::get<ParsedData::VarInt>(d.val).value });
+el.exclusive_size += el.container_stats->capacity;
+"""
+
+[[codegen.processor]]
+type = "types::st::VarInt<DB>"
+func = """
+el.container_stats->length = std::get<ParsedData::VarInt>(d.val).value;
+"""