Introduce SyntaxGraph::StructuredNode::at

Summary: re-use type_system::FastFieldHandle

Reviewed By: praihan

Differential Revision: D79918154

fbshipit-source-id: 857b6f7756ac56ea57c8f4d13b861ce99538b931

Author: thedavekwon

third-party/thrift/src/thrift/lib/cpp2/schema/SyntaxGraph.h

@@ -630,6 +630,23 @@ class StructuredNode : detail::WithDefinition, detail::WithUri {
   using detail::WithUri::uri;
   folly::span<const FieldNode> fields() const { return fields_; }
 
+  /**
+   * Looks up a field by ID.
+   *
+   * Throws:
+   *   - std::out_of_range if the field ID is not present.
+   */
+  const FieldNode& at(FieldId id) const { return at(fieldHandleFor(id)); }
+  /**
+   * Looks up a field by name.
+   *
+   * Throws:
+   *   - std::out_of_range if the field ID is not present.
+   */
+  const FieldNode& at(std::string_view name) const {
+    return at(fieldHandleFor(name));
+  }
+
  protected:
   StructuredNode(
       const detail::Resolver& resolver,
@@ -638,10 +655,65 @@ class StructuredNode : detail::WithDefinition, detail::WithUri {
       std::vector<FieldNode> fields)
       : detail::WithDefinition(resolver, definitionKey),
         detail::WithUri(uri),
-        fields_(std::move(fields)) {}
+        fields_(std::move(fields)) {
+    std::uint16_t ordinal = 1;
+    for (const FieldNode& field : fields_) {
+      fieldHandleById_.emplace(
+          field.id(), type_system::FastFieldHandle{ordinal});
+      fieldHandleByName_.emplace(
+          field.name(), type_system::FastFieldHandle{ordinal});
+      ++ordinal;
+    }
+  }
 
  private:
+  /**
+   * Looks up a field by a fast field handle, previously obtained from
+   * fieldHandleFor(...);
+   *
+   * Preconditions:
+   *   - The provided handle was obtained by calling fieldHandleFor(...) on this
+   *     instance.
+   *
+   * Throws:
+   *   - std::out_of_range if the field handle is invalid or out of range.
+   */
+  const FieldNode& at(type_system::FastFieldHandle handle) const {
+    if (!handle.valid() || handle.ordinal > fields_.size()) {
+      folly::throw_exception<std::out_of_range>(
+          fmt::format("invalid field handle: {}", handle.ordinal));
+    }
+    return fields_.at(handle.ordinal - 1);
+  }
+
+  /**
+   * Returns a field handle for the given field ID, if it exists, returning
+   * `FastFieldHandle::invalid()` otherwise.
+   */
+  type_system::FastFieldHandle fieldHandleFor(FieldId id) const noexcept {
+    if (const type_system::FastFieldHandle* handle =
+            folly::get_ptr(fieldHandleById_, id)) {
+      return *handle;
+    }
+    return type_system::FastFieldHandle::invalid();
+  }
+  /**
+   * Returns a field handle for the given field name, if it exists, returning
+   * `FastFieldHandle::invalid()` otherwise.
+   */
+  type_system::FastFieldHandle fieldHandleFor(
+      std::string_view name) const noexcept {
+    if (const type_system::FastFieldHandle* handle =
+            folly::get_ptr(fieldHandleByName_, name)) {
+      return *handle;
+    }
+    return type_system::FastFieldHandle::invalid();
+  }
+
   std::vector<FieldNode> fields_;
+  folly::F14FastMap<FieldId, type_system::FastFieldHandle> fieldHandleById_;
+  folly::F14FastMap<std::string_view, type_system::FastFieldHandle>
+      fieldHandleByName_;
 };
 
 class StructNode final : folly::MoveOnly,

third-party/thrift/src/thrift/lib/cpp2/schema/test/SyntaxGraphTest.cpp

@@ -185,22 +185,34 @@ TEST_F(ServiceSchemaTest, Struct) {
 
   EXPECT_EQ(s.fields().size(), 2);
 
-  EXPECT_EQ(s.fields()[0].id(), FieldId{1});
-  EXPECT_EQ(
-      s.fields()[0].presence(), FieldNode::PresenceQualifier::UNQUALIFIED);
-  EXPECT_EQ(s.fields()[0].type().asPrimitive(), Primitive::I32);
-  EXPECT_EQ(s.fields()[0].name(), "field1");
-  EXPECT_STREQ(s.fields()[0].name().data(), "field1");
-  EXPECT_EQ(s.fields()[0].customDefault()->as_i32(), 10);
-
-  EXPECT_EQ(s.fields()[1].id(), FieldId{2});
-  EXPECT_EQ(s.fields()[1].presence(), FieldNode::PresenceQualifier::OPTIONAL_);
-  EXPECT_EQ(
-      &s.fields()[1].type().asEnum(),
-      &program.definitionsByName().at("TestEnum")->asEnum());
-  EXPECT_EQ(s.fields()[1].name(), "field2");
-  EXPECT_STREQ(s.fields()[1].name().data(), "field2");
-  EXPECT_EQ(s.fields()[1].customDefault(), nullptr);
+  auto checkField1 = [](const auto& field) {
+    EXPECT_EQ(field.id(), FieldId{1});
+    EXPECT_EQ(field.presence(), FieldNode::PresenceQualifier::UNQUALIFIED);
+    EXPECT_EQ(field.type().asPrimitive(), Primitive::I32);
+    EXPECT_EQ(field.name(), "field1");
+    EXPECT_STREQ(field.name().data(), "field1");
+    EXPECT_EQ(field.customDefault()->as_i32(), 10);
+  };
+
+  checkField1(s.at(FieldId{1}));
+  checkField1(s.at("field1"));
+
+  auto checkField2 = [&](const auto& field) {
+    EXPECT_EQ(field.id(), FieldId{2});
+    EXPECT_EQ(field.presence(), FieldNode::PresenceQualifier::OPTIONAL_);
+    EXPECT_EQ(
+        &field.type().asEnum(),
+        &program.definitionsByName().at("TestEnum")->asEnum());
+    EXPECT_EQ(field.name(), "field2");
+    EXPECT_STREQ(field.name().data(), "field2");
+    EXPECT_EQ(field.customDefault(), nullptr);
+  };
+
+  checkField2(s.at(FieldId{2}));
+  checkField2(s.at("field2"));
+
+  EXPECT_THROW(s.at(FieldId{3}), std::out_of_range);
+  EXPECT_THROW(s.at("field3"), std::out_of_range);
 
   EXPECT_EQ(
       s.toDebugString(),
@@ -229,17 +241,33 @@ TEST_F(ServiceSchemaTest, Union) {
 
   EXPECT_EQ(u.fields().size(), 2);
 
-  EXPECT_EQ(u.fields()[0].id(), FieldId{1});
-  EXPECT_EQ(u.fields()[0].name(), "s");
-  EXPECT_EQ(
-      &u.fields()[0].type().asStruct(),
-      &program.definitionsByName().at("TestStruct")->asStruct());
-
-  EXPECT_EQ(u.fields()[1].id(), FieldId{2});
-  EXPECT_EQ(u.fields()[1].name(), "e");
-  EXPECT_EQ(
-      &u.fields()[1].type().asEnum(),
-      &program.definitionsByName().at("TestEnum")->asEnum());
+  auto checkField1 = [&](const auto& field) {
+    EXPECT_EQ(field.id(), FieldId{1});
+    EXPECT_EQ(field.presence(), FieldNode::PresenceQualifier::UNQUALIFIED);
+    EXPECT_EQ(
+        &field.type().asStruct(),
+        &program.definitionsByName().at("TestStruct")->asStruct());
+    EXPECT_EQ(field.name(), "s");
+    EXPECT_STREQ(field.name().data(), "s");
+  };
+
+  checkField1(u.at(FieldId{1}));
+  checkField1(u.at("s"));
+
+  auto checkField2 = [&](const auto& field) {
+    EXPECT_EQ(field.id(), FieldId{2});
+    EXPECT_EQ(field.presence(), FieldNode::PresenceQualifier::UNQUALIFIED);
+    EXPECT_EQ(
+        &field.type().asEnum(),
+        &program.definitionsByName().at("TestEnum")->asEnum());
+    EXPECT_EQ(field.name(), "e");
+  };
+
+  checkField2(u.at(FieldId{2}));
+  checkField2(u.at("e"));
+
+  EXPECT_THROW(u.at(FieldId{3}), std::out_of_range);
+  EXPECT_THROW(u.at("field3"), std::out_of_range);
 
   EXPECT_EQ(
       u.toDebugString(),