[BinaryAnnotator] Add more options that control the generation of .afb files (#8323)

* [BinaryAnnotator] Add more options that control the generation of `.afb` files.

* [BinaryAnnotator] Update the include paths.

.gitignore

@@ -153,3 +153,5 @@ cmake-build-debug/
 _deps/
 **/.gradle/**
 kotlin/**/generated
+MODULE.bazel
+MODULE.bazel.lock

src/annotated_binary_text_gen.cpp

@@ -402,7 +402,8 @@ static void GenerateSection(std::ostream &os, const BinarySection &section,
 }  // namespace
 
 bool AnnotatedBinaryTextGenerator::Generate(
-    const std::string &filename, const std::string &schema_filename) {
+    const std::string &filename, const std::string &schema_filename,
+    const std::string &output_filename) {
   OutputConfig output_config;
   output_config.max_bytes_per_line = options_.max_bytes_per_line;
   output_config.include_vector_contents = options_.include_vector_contents;
@@ -435,18 +436,23 @@ bool AnnotatedBinaryTextGenerator::Generate(
     }
   }
 
-  // Modify the output filename.
+  std::string out = output_filename;
-  std::string output_filename = StripExtension(filename);
+  if (out.empty()) {
-  output_filename += options_.output_postfix;
+    // Modify the output filename.
-  output_filename +=
+    out = StripExtension(filename);
-      "." + (options_.output_extension.empty() ? GetExtension(filename)
+    out += options_.output_postfix;
-                                               : options_.output_extension);
+    out +=
+        "." + (options_.output_extension.empty() ? GetExtension(filename)
+                                                 : options_.output_extension);
+  }
 
-  std::ofstream ofs(output_filename.c_str());
+  std::ofstream ofs(out.c_str());
 
   ofs << "// Annotated Flatbuffer Binary" << std::endl;
   ofs << "//" << std::endl;
-  ofs << "// Schema file: " << schema_filename << std::endl;
+  if (!schema_filename.empty()) {
+    ofs << "// Schema file: " << schema_filename << std::endl;
+  }
   ofs << "// Binary file: " << filename << std::endl;
 
   // Generate each of the binary sections

src/annotated_binary_text_gen.h

@@ -56,8 +56,8 @@ class AnnotatedBinaryTextGenerator {
 
   // Generate the annotated binary for the given `filename`. Returns true if the
   // annotated binary was successfully saved.
-  bool Generate(const std::string &filename,
+  bool Generate(const std::string &filename, const std::string &schema_filename,
-                const std::string &schema_filename);
+                const std::string &output_filename = "");
 
  private:
   const std::map<uint64_t, BinarySection> annotations_;

src/binary_annotator.cpp

@@ -120,12 +120,13 @@ static BinarySection GenerateMissingSection(const uint64_t offset,
 }  // namespace
 
 std::map<uint64_t, BinarySection> BinaryAnnotator::Annotate() {
-  flatbuffers::Verifier verifier(bfbs_, static_cast<size_t>(bfbs_length_));
+  if (bfbs_ != nullptr && bfbs_length_ != 0) {
-
+    flatbuffers::Verifier verifier(bfbs_, static_cast<size_t>(bfbs_length_));
-  if ((is_size_prefixed_ &&
+    if ((is_size_prefixed_ &&
-       !reflection::VerifySizePrefixedSchemaBuffer(verifier)) ||
+         !reflection::VerifySizePrefixedSchemaBuffer(verifier)) ||
-      !reflection::VerifySchemaBuffer(verifier)) {
+        !reflection::VerifySchemaBuffer(verifier)) {
-    return {};
+      return {};
+    }
   }
 
   // The binary is too short to read as a flatbuffers.
@@ -141,8 +142,7 @@ std::map<uint64_t, BinarySection> BinaryAnnotator::Annotate() {
 
   if (IsValidOffset(root_table_offset)) {
     // Build the root table, and all else will be referenced from it.
-    BuildTable(root_table_offset, BinarySectionType::RootTable,
+    BuildTable(root_table_offset, BinarySectionType::RootTable, RootTable());
-               schema_->root_table());
   }
 
   // Now that all the sections are built, make sure the binary sections are
@@ -203,7 +203,7 @@ uint64_t BinaryAnnotator::BuildHeader(const uint64_t header_offset) {
 
   BinaryRegionComment root_offset_comment;
   root_offset_comment.type = BinaryRegionCommentType::RootTableOffset;
-  root_offset_comment.name = schema_->root_table()->name()->str();
+  root_offset_comment.name = RootTable()->name()->str();
 
   if (!IsValidOffset(root_table_loc)) {
     SetError(root_offset_comment,
@@ -1516,4 +1516,11 @@ bool BinaryAnnotator::ContainsSection(const uint64_t offset) {
                                              it->second.regions.back().length;
 }
 
+const reflection::Object *BinaryAnnotator::RootTable() const {
+  if (!root_table_.empty()) {
+    return schema_->objects()->LookupByKey(root_table_);
+  }
+  return schema_->root_table();
+}
+
 }  // namespace flatbuffers

src/binary_annotator.h

@@ -17,15 +17,21 @@
 #ifndef FLATBUFFERS_BINARY_ANNOTATOR_H_
 #define FLATBUFFERS_BINARY_ANNOTATOR_H_
 
+#include <cstddef>
+#include <cstdint>
+#include <iomanip>
+#include <ios>
 #include <list>
 #include <map>
+#include <sstream>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "flatbuffers/base.h"
 #include "flatbuffers/reflection.h"
+#include "flatbuffers/reflection_generated.h"
 #include "flatbuffers/stl_emulation.h"
-#include "flatbuffers/util.h"
 
 namespace flatbuffers {
 
@@ -227,7 +233,7 @@ inline static std::string ToString(const BinaryRegionType type) {
     case BinaryRegionType::Uint8: return "uint8_t";
     case BinaryRegionType::Uint16: return "uint16_t";
     case BinaryRegionType::Uint32: return "uint32_t";
-    case BinaryRegionType::Uint64: return "uint64_t"; ;
+    case BinaryRegionType::Uint64: return "uint64_t";
     case BinaryRegionType::Int8: return "int8_t";
     case BinaryRegionType::Int16: return "int16_t";
     case BinaryRegionType::Int32: return "int32_t";
@@ -250,6 +256,18 @@ class BinaryAnnotator {
       : bfbs_(bfbs),
         bfbs_length_(bfbs_length),
         schema_(reflection::GetSchema(bfbs)),
+        root_table_(""),
+        binary_(binary),
+        binary_length_(binary_length),
+        is_size_prefixed_(is_size_prefixed) {}
+
+  BinaryAnnotator(const reflection::Schema *schema,
+                  const std::string &root_table, const uint8_t *binary,
+                  uint64_t binary_length, bool is_size_prefixed)
+      : bfbs_(nullptr),
+        bfbs_length_(0),
+        schema_(schema),
+        root_table_(root_table),
         binary_(binary),
         binary_length_(binary_length),
         is_size_prefixed_(is_size_prefixed) {}
@@ -329,7 +347,7 @@ class BinaryAnnotator {
   }
 
   // Adds the provided `section` keyed by the `offset` it occurs at. If a
-  // section is already added at that offset, it doesn't replace the exisiting
+  // section is already added at that offset, it doesn't replace the existing
   // one.
   void AddSection(const uint64_t offset, const BinarySection &section) {
     sections_.insert(std::make_pair(offset, section));
@@ -384,10 +402,13 @@ class BinaryAnnotator {
 
   bool ContainsSection(const uint64_t offset);
 
+  const reflection::Object *RootTable() const;
+
   // The schema for the binary file
   const uint8_t *bfbs_;
   const uint64_t bfbs_length_;
   const reflection::Schema *schema_;
+  const std::string root_table_;
 
   // The binary data itself.
   const uint8_t *binary_;