Improve implementation of doc comment collection and expose in SchemaParser interface.

c++/src/capnp/compiler/capnp.c++

@@ -433,11 +433,7 @@ public:
       nodes.setWithCaveats(i, schemas[i].getProto());
     }
 
-    auto docs = compiler->getLoader().getAllDocs();
-    auto docnodes = request.initNodeDocs(docs.size());
-    for (size_t i = 0; i < docs.size(); i++) {
-      docnodes.setWithCaveats(i, docs[i]);
-    }
+    request.adoptSourceInfo(compiler->getAllSourceInfo(message.getOrphanage()));
 
     auto requestedFiles = request.initRequestedFiles(sourceFiles.size());
     for (size_t i = 0; i < sourceFiles.size(); i++) {

c++/src/capnp/compiler/compiler.h

@@ -91,10 +91,16 @@ public:
   // exception if the parent ID is not recognized; returns null if the parent has no child of the
   // given name.  Neither the parent nor the child schema node is actually compiled.
 
+  kj::Maybe<schema::Node::SourceInfo::Reader> getSourceInfo(uint64_t id) const;
+  // Get the SourceInfo for the given type ID, if available.
+
   Orphan<List<schema::CodeGeneratorRequest::RequestedFile::Import>>
       getFileImportTable(Module& module, Orphanage orphanage) const;
   // Build the import table for the CodeGeneratorRequest for the given module.
 
+  Orphan<List<schema::Node::SourceInfo>> getAllSourceInfo(Orphanage orphanage) const;
+  // Gets the SourceInfo structs for all nodes parsed by the compiler.
+
   enum Eagerness: uint32_t {
     // Flags specifying how eager to be about compilation.  These are intended to be bitwise OR'd.
     // Used with the method `eagerlyCompile()`.

c++/src/capnp/compiler/node-translator.h

@@ -137,6 +137,9 @@ public:
     // Auxiliary nodes that were produced when translating this node and should be loaded along
     // with it.  In particular, structs that contain groups (or named unions) spawn extra nodes
     // representing those, and interfaces spawn struct nodes representing method params/results.
+
+    kj::Array<schema::Node::SourceInfo::Reader> sourceInfo;
+    // The SourceInfo for the node and all aux nodes.
   };
 
   NodeSet getBootstrapNode();
@@ -161,9 +164,6 @@ public:
   // `brandBuilder` may be used to construct a message which will fill in ResolvedDecl::brand in
   // the result.
 
-  kj::Maybe<Orphan<schema::NodeDoc>>& getDoc() { return wipNodeDoc; };
-  void addFieldDoc(uint codeOrder, ::capnp::Text::Reader docComment);
-
 private:
   class DuplicateNameDetector;
   class DuplicateOrdinalDetector;
@@ -179,17 +179,21 @@ private:
   kj::Own<BrandScope> localBrand;
 
   Orphan<schema::Node> wipNode;
-  kj::Maybe<Orphan<schema::NodeDoc>> wipNodeDoc;
-  // The work-in-progress schema node and its docstring
+  // The work-in-progress schema node.
 
-  kj::Vector<std::pair<uint, ::capnp::Text::Reader>> fieldDocs;
-  // TODO(now): Don't use std::pair.
+  Orphan<schema::Node::SourceInfo> sourceInfo;
+  // Doc comments and other source info for this node.
+
+  struct AuxNode {
+    Orphan<schema::Node> node;
+    Orphan<schema::Node::SourceInfo> sourceInfo;
+  };
 
-  kj::Vector<Orphan<schema::Node>> groups;
+  kj::Vector<AuxNode> groups;
   // If this is a struct node and it contains groups, these are the nodes for those groups,  which
   // must be loaded together with the top-level node.
 
-  kj::Vector<Orphan<schema::Node>> paramStructs;
+  kj::Vector<AuxNode> paramStructs;
   // If this is an interface, these are the auto-generated structs representing params and results.
 
   struct UnfinishedValue {

c++/src/capnp/schema-loader.c++

@@ -139,8 +139,6 @@ public:
   const _::RawBrandedSchema* getUnbound(const _::RawSchema* schema);
 
   kj::Array<Schema> getAllLoaded() const;
-  kj::Array<schema::NodeDoc::Reader> getAllDocs() const;
-  void loadDoc(schema::NodeDoc::Reader docReader);
 
   void requireStructSize(uint64_t id, uint dataWordCount, uint pointerCount);
   // Require any struct nodes loaded with this ID -- in the past and in the future -- to have at
@@ -159,7 +157,6 @@ private:
   std::unordered_map<uint64_t, _::RawSchema*> schemas;
   std::unordered_map<SchemaBindingsPair, _::RawBrandedSchema*, SchemaBindingsPairHash> brands;
   std::unordered_map<const _::RawSchema*, _::RawBrandedSchema*> unboundBrands;
-  std::unordered_map<uint64_t, schema::NodeDoc::Reader> docs;
 
   struct RequiredSize {
     uint16_t dataWordCount;
@@ -1849,24 +1846,6 @@ kj::Array<Schema> SchemaLoader::Impl::getAllLoaded() const {
   return result;
 }
 
-kj::Array<schema::NodeDoc::Reader> SchemaLoader::Impl::getAllDocs() const {
-  size_t count = 0;
-  for (auto& doc: docs) {
-    ++count;
-  }
-
-  kj::Array<schema::NodeDoc::Reader> result = kj::heapArray<schema::NodeDoc::Reader>(count);
-  size_t i = 0;
-  for (auto& doc: docs) {
-    result[i++] = doc.second;
-  }
-  return result;
-}
-
-void SchemaLoader::Impl::loadDoc(schema::NodeDoc::Reader docReader) {
-  docs[docReader.getId()] = docReader;
-}
-
 void SchemaLoader::Impl::requireStructSize(uint64_t id, uint dataWordCount, uint pointerCount) {
   auto& slot = structSizeRequirements[id];
   slot.dataWordCount = kj::max(slot.dataWordCount, dataWordCount);
@@ -2126,14 +2105,6 @@ kj::Array<Schema> SchemaLoader::getAllLoaded() const {
   return impl.lockShared()->get()->getAllLoaded();
 }
 
-kj::Array<schema::NodeDoc::Reader> SchemaLoader::getAllDocs() const {
-  return impl.lockShared()->get()->getAllDocs();
-}
-
-void SchemaLoader::loadDoc(schema::NodeDoc::Reader docReader) const {
-  return impl.lockExclusive()->get()->loadDoc(docReader);
-}
-
 void SchemaLoader::loadNative(const _::RawSchema* nativeSchema) {
   impl.lockExclusive()->get()->loadNative(nativeSchema);
 }

c++/src/capnp/schema-loader.h

@@ -152,9 +152,6 @@ public:
   // loadCompiledTypeAndDependencies<T>() in order to get a flat list of all of T's transitive
   // dependencies.
 
-  kj::Array<schema::NodeDoc::Reader> getAllDocs() const;
-  void loadDoc(schema::NodeDoc::Reader docReader) const;
-
 private:
   class Validator;
   class CompatibilityChecker;

c++/src/capnp/schema-parser-test.c++

@@ -181,5 +181,81 @@ TEST(SchemaParser, Constants) {
   EXPECT_EQ("text", genericConst.get("value").as<Text>());
 }
 
+void expectSourceInfo(schema::Node::SourceInfo::Reader sourceInfo,
+                   uint64_t expectedId, kj::StringPtr expectedComment,
+                   std::initializer_list<const kj::StringPtr> expectedMembers) {
+  KJ_EXPECT(sourceInfo.getId() == expectedId, sourceInfo, expectedId);
+  KJ_EXPECT(sourceInfo.getDocComment() == expectedComment, sourceInfo, expectedComment);
+
+  auto members = sourceInfo.getMembers();
+  KJ_ASSERT(members.size() == expectedMembers.size());
+  for (auto i: kj::indices(expectedMembers)) {
+    KJ_EXPECT(members[i].getDocComment() == expectedMembers.begin()[i],
+              members[i], expectedMembers.begin()[i]);
+  }
+}
+
+TEST(SchemaParser, SourceInfo) {
+  SchemaParser parser;
+  FakeFileReader reader;
+
+  reader.add("foo.capnp",
+      "@0x84a2c6051e1061ed;\n"
+      "# file doc comment\n"
+      "\n"
+      "struct Foo @0xc6527d0a670dc4c3 {\n"
+      "  # struct doc comment\n"
+      "  # second line\n"
+      "\n"
+      "  bar @0 :UInt32;\n"
+      "  # field doc comment\n"
+      "  baz :group {\n"
+      "    # group doc comment\n"
+      "    qux @1 :Text;\n"
+      "    # group field doc comment\n"
+      "  }\n"
+      "}\n"
+      "\n"
+      "enum Corge @0xae08878f1a016f14 {\n"
+      "  # enum doc comment\n"
+      "  grault @0;\n"
+      "  # enumerant doc comment\n"
+      "  garply @1;\n"
+      "}\n"
+      "\n"
+      "interface Waldo @0xc0f1b0aff62b761e {\n"
+      "  # interface doc comment\n"
+      "  fred @0 (plugh :Int32) -> (xyzzy :Text);\n"
+      "  # method doc comment\n"
+      "}\n"
+      "\n"
+      "struct Thud @0xcca9972702b730b4 {}\n"
+      "# post-comment\n");
+
+  ParsedSchema file = parser.parseFile(SchemaFile::newDiskFile(
+      "foo.capnp", "foo.capnp", nullptr, reader));
+  ParsedSchema foo = file.getNested("Foo");
+
+  expectSourceInfo(file.getSourceInfo(), 0x84a2c6051e1061edull, "file doc comment\n", {});
+
+  expectSourceInfo(foo.getSourceInfo(), 0xc6527d0a670dc4c3ull, "struct doc comment\nsecond line\n",
+      { "field doc comment\n", "group doc comment\n" });
+
+  auto group = foo.asStruct().getFieldByName("baz").getType().asStruct();
+  expectSourceInfo(KJ_ASSERT_NONNULL(parser.getSourceInfo(group)),
+      group.getProto().getId(), "group doc comment\n", { "group field doc comment\n" });
+
+  ParsedSchema corge = file.getNested("Corge");
+  expectSourceInfo(corge.getSourceInfo(), 0xae08878f1a016f14, "enum doc comment\n",
+      { "enumerant doc comment\n", "" });
+
+  ParsedSchema waldo = file.getNested("Waldo");
+  expectSourceInfo(waldo.getSourceInfo(), 0xc0f1b0aff62b761e, "interface doc comment\n",
+      { "method doc comment\n" });
+
+  ParsedSchema thud = file.getNested("Thud");
+  expectSourceInfo(thud.getSourceInfo(), 0xcca9972702b730b4, "post-comment\n", {});
+}
+
 }  // namespace
 }  // namespace capnp

c++/src/capnp/schema-parser.c++

@@ -196,6 +196,10 @@ ParsedSchema SchemaParser::parseFile(kj::Own<SchemaFile>&& file) const {
   return ParsedSchema(impl->compiler.getLoader().get(id), *this);
 }
 
+kj::Maybe<schema::Node::SourceInfo::Reader> SchemaParser::getSourceInfo(Schema schema) const {
+  return impl->compiler.getSourceInfo(schema.getProto().getId());
+}
+
 SchemaParser::ModuleImpl& SchemaParser::getModuleImpl(kj::Own<SchemaFile>&& file) const {
   auto lock = impl->fileMap.lockExclusive();
 
@@ -226,6 +230,10 @@ ParsedSchema ParsedSchema::getNested(kj::StringPtr nestedName) const {
   }
 }
 
+schema::Node::SourceInfo::Reader ParsedSchema::getSourceInfo() const {
+  return KJ_ASSERT_NONNULL(parser->getSourceInfo(*this));
+}
+
 // =======================================================================================
 
 namespace {

c++/src/capnp/schema-parser.h

@@ -77,6 +77,11 @@ public:
   // normally.  In this case, the result is a best-effort attempt to compile the schema, but it
   // may be invalid or corrupt, and using it for anything may cause exceptions to be thrown.
 
+  kj::Maybe<schema::Node::SourceInfo::Reader> getSourceInfo(Schema schema) const;
+  // Look up source info (e.g. doc comments) for the given schema, which must have come from this
+  // SchemaParser. Note that this will also work for implicit group and param types that don't have
+  // a type name hence don't have a `ParsedSchema`.
+
   template <typename T>
   inline void loadCompiledTypeAndDependencies() {
     // See SchemaLoader::loadCompiledTypeAndDependencies().
@@ -110,6 +115,9 @@ public:
   // Gets the nested node with the given name, or throws an exception if there is no such nested
   // declaration.
 
+  schema::Node::SourceInfo::Reader getSourceInfo() const;
+  // Get the source info for this schema.
+
 private:
   inline ParsedSchema(Schema inner, const SchemaParser& parser): Schema(inner), parser(&parser) {}
 

c++/src/capnp/schema.capnp

@@ -169,19 +169,33 @@ struct Node {
       targetsAnnotation @30 :Bool;
     }
   }
-}
 
-struct NodeDoc {
-  # separate carrier for documentation comments on Nodes,
-  # to keep them out of the binary descriptors
+  struct SourceInfo {
+    # Additional information about a node which is not needed at runtime, but may be useful for
+    # documentation or debugging purposes. This is kept in a separate struct to make sure it
+    # doesn't accidentally get included in contexts where it is not needed. The
+    # `CodeGeneratorRequest` includes this information in a separate array.
 
-  id @0 :Id;
-  # ID should exist as Node in the same request
+    id @0 :Id;
+    # ID of the Node which this info describes.
 
-  docComment @1 :Text;
+    docComment @1 :Text;
+    # The top-level doc comment for the Node.
 
-  fieldDocs @2 :List(FieldDoc);
-  # valid only if Node is a "struct"
+    members @2 :List(Member);
+    # Information about each member -- i.e. fields (for structs), enumerants (for enums), or
+    # methods (for interfaces).
+    #
+    # This list is the same length and order as the corresponding list in the Node, i.e.
+    # Node.struct.fields, Node.enum.enumerants, or Node.interface.methods.
+
+    struct Member {
+      docComment @0 :Text;
+      # Doc comment on the member.
+    }
+
+    # TODO(someday): Record location of the declaration in the original source code.
+  }
 }
 
 struct Field {
@@ -242,13 +256,6 @@ struct Field {
   }
 }
 
-struct FieldDoc {
-  # separate container to carry field docstrings
-
-  codeOrder @0 :UInt16;
-  docComment @1 :Text;
-}
-
 struct Enumerant {
   # Schema for member of an enum.
 
@@ -488,8 +495,9 @@ struct CodeGeneratorRequest {
   # All nodes parsed by the compiler, including for the files on the command line and their
   # imports.
 
-  nodeDocs @3 :List(NodeDoc);
-  # documentation comments for nodes, where present
+  sourceInfo @3 :List(Node.SourceInfo);
+  # Information about the original source code for each node, where available. This array may be
+  # omitted or may be missing some nodes if no info is available for them.
 
   requestedFiles @1 :List(RequestedFile);
   # Files which were listed on the command line.