diff --git a/tools/api_proto_plugin/BUILD b/tools/api_proto_plugin/BUILD
new file mode 100644
index 000000000000..646490276955
--- /dev/null
+++ b/tools/api_proto_plugin/BUILD
@@ -0,0 +1,17 @@
+licenses(["notice"])  # Apache 2
+
+py_library(
+    name = "api_proto_plugin",
+    srcs = [
+        "annotations.py",
+        "plugin.py",
+        "traverse.py",
+        "type_context.py",
+        "visitor.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_protobuf//:protobuf_python",
+    ],
+)
diff --git a/tools/api_proto_plugin/__init__.py b/tools/api_proto_plugin/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tools/api_proto_plugin/annotations.py b/tools/api_proto_plugin/annotations.py
new file mode 100644
index 000000000000..eadd080b03aa
--- /dev/null
+++ b/tools/api_proto_plugin/annotations.py
@@ -0,0 +1,79 @@
+"""Envoy API annotations."""
+
+from collections import namedtuple
+
+import re
+
+# Key-value annotation regex.
+ANNOTATION_REGEX = re.compile('\[#([\w-]+?):(.*?)\]\s?', re.DOTALL)
+
+# Page/section titles with special prefixes in the proto comments
+DOC_TITLE_ANNOTATION = 'protodoc-title'
+
+# Not implemented yet annotation on leading comments, leading to insertion of
+# warning on field.
+NOT_IMPLEMENTED_WARN_ANNOTATION = 'not-implemented-warn'
+
+# Not implemented yet annotation on leading comments, leading to hiding of
+# field.
+NOT_IMPLEMENTED_HIDE_ANNOTATION = 'not-implemented-hide'
+
+# Comment that allows for easy searching for things that need cleaning up in the next major
+# API version.
+NEXT_MAJOR_VERSION_ANNOTATION = 'next-major-version'
+
+# Comment. Just used for adding text that will not go into the docs at all.
+COMMENT_ANNOTATION = 'comment'
+
+# proto compatibility status.
+PROTO_STATUS_ANNOTATION = 'proto-status'
+
+# Where v2 differs from v1..
+V2_API_DIFF_ANNOTATION = 'v2-api-diff'
+
+VALID_ANNOTATIONS = set([
+    DOC_TITLE_ANNOTATION,
+    NOT_IMPLEMENTED_WARN_ANNOTATION,
+    NOT_IMPLEMENTED_HIDE_ANNOTATION,
+    V2_API_DIFF_ANNOTATION,
+    NEXT_MAJOR_VERSION_ANNOTATION,
+    COMMENT_ANNOTATION,
+    PROTO_STATUS_ANNOTATION,
+])
+
+# These can propagate from file scope to message/enum scope (and be overridden).
+INHERITED_ANNOTATIONS = set([
+    PROTO_STATUS_ANNOTATION,
+])
+
+
+class AnnotationError(Exception):
+  """Base error class for the annotations module."""
+
+
+def ExtractAnnotations(s, inherited_annotations=None):
+  """Extract annotations map from a given comment string.
+
+  Args:
+    s: string that may contains annotations.
+    inherited_annotations: annotation map from file-level inherited annotations
+      (or None) if this is a file-level comment.
+
+  Returns:
+    Annotation map.
+  """
+  annotations = {
+      k: v for k, v in (inherited_annotations or {}).items() if k in INHERITED_ANNOTATIONS
+  }
+  # Extract annotations.
+  groups = re.findall(ANNOTATION_REGEX, s)
+  for group in groups:
+    annotation = group[0]
+    if annotation not in VALID_ANNOTATIONS:
+      raise AnnotationError('Unknown annotation: %s' % annotation)
+    annotations[group[0]] = group[1].lstrip()
+  return annotations
+
+
+def WithoutAnnotations(s):
+  return re.sub(ANNOTATION_REGEX, '', s)
diff --git a/tools/api_proto_plugin/plugin.py b/tools/api_proto_plugin/plugin.py
new file mode 100644
index 000000000000..0a56ea8e3f8b
--- /dev/null
+++ b/tools/api_proto_plugin/plugin.py
@@ -0,0 +1,57 @@
+"""Python protoc plugin for Envoy APIs."""
+
+import cProfile
+import io
+import os
+import pstats
+import sys
+
+from tools.api_proto_plugin import traverse
+
+from google.protobuf.compiler import plugin_pb2
+
+
+def Plugin(output_suffix, visitor):
+  """Protoc plugin entry point.
+
+  This defines protoc plugin and manages the stdin -> stdout flow. An
+  api_proto_plugin is defined by the provided visitor.
+
+  See
+  http://www.expobrain.net/2015/09/13/create-a-plugin-for-google-protocol-buffer/
+  for further details on protoc plugin basics.
+
+  Args:
+    output_suffix: output files are generated alongside their corresponding
+      input .proto, with this filename suffix.
+    visitor: visitor.Visitor defining the business logic of the plugin.
+  """
+  request = plugin_pb2.CodeGeneratorRequest()
+  request.ParseFromString(sys.stdin.buffer.read())
+  response = plugin_pb2.CodeGeneratorResponse()
+  cprofile_enabled = os.getenv('CPROFILE_ENABLED')
+
+  # We use request.file_to_generate rather than request.file_proto here since we
+  # are invoked inside a Bazel aspect, each node in the DAG will be visited once
+  # by the aspect and we only want to generate docs for the current node.
+  for file_to_generate in request.file_to_generate:
+    # Find the FileDescriptorProto for the file we actually are generating.
+    file_proto = [pf for pf in request.proto_file if pf.name == file_to_generate][0]
+    f = response.file.add()
+    f.name = file_proto.name + output_suffix
+    if cprofile_enabled:
+      pr = cProfile.Profile()
+      pr.enable()
+    # We don't actually generate any RST right now, we just string dump the
+    # input proto file descriptor into the output file.
+    f.content = traverse.TraverseFile(file_proto, visitor)
+    if cprofile_enabled:
+      pr.disable()
+      stats_stream = io.StringIO()
+      ps = pstats.Stats(pr,
+                        stream=stats_stream).sort_stats(os.getenv('CPROFILE_SORTBY', 'cumulative'))
+      stats_file = response.file.add()
+      stats_file.name = file_proto.name + output_suffix + '.profile'
+      ps.print_stats()
+      stats_file.content = stats_stream.getvalue()
+  sys.stdout.buffer.write(response.SerializeToString())
diff --git a/tools/api_proto_plugin/traverse.py b/tools/api_proto_plugin/traverse.py
new file mode 100644
index 000000000000..6ad97b8699aa
--- /dev/null
+++ b/tools/api_proto_plugin/traverse.py
@@ -0,0 +1,72 @@
+"""FileDescriptorProto traversal for api_proto_plugin framework."""
+
+from tools.api_proto_plugin import type_context
+
+
+def TraverseEnum(type_context, enum_proto, visitor):
+  """Traverse an enum definition.
+
+  Args:
+    type_context: type_context.TypeContext for enum type.
+    enum_proto: EnumDescriptorProto for enum.
+    visitor: visitor.Visitor defining the business logic of the plugin.
+
+  Returns:
+    Plugin specific output.
+  """
+  return visitor.VisitEnum(enum_proto, type_context)
+
+
+def TraverseMessage(type_context, msg_proto, visitor):
+  """Traverse a message definition.
+
+  Args:
+    type_context: type_context.TypeContext for message type.
+    msg_proto: DescriptorProto for message.
+    visitor: visitor.Visitor defining the business logic of the plugin.
+
+  Returns:
+    Plugin specific output.
+  """
+  # Skip messages synthesized to represent map types.
+  if msg_proto.options.map_entry:
+    return ''
+  # We need to do some extra work to recover the map type annotation from the
+  # synthesized messages.
+  type_context.map_typenames = {
+      '%s.%s' % (type_context.name, nested_msg.name): (nested_msg.field[0], nested_msg.field[1])
+      for nested_msg in msg_proto.nested_type
+      if nested_msg.options.map_entry
+  }
+  nested_msgs = [
+      TraverseMessage(type_context.ExtendNestedMessage(index, nested_msg.name), nested_msg, visitor)
+      for index, nested_msg in enumerate(msg_proto.nested_type)
+  ]
+  nested_enums = [
+      TraverseEnum(type_context.ExtendNestedEnum(index, nested_enum.name), nested_enum, visitor)
+      for index, nested_enum in enumerate(msg_proto.enum_type)
+  ]
+  return visitor.VisitMessage(msg_proto, type_context, nested_msgs, nested_enums)
+
+
+def TraverseFile(file_proto, visitor):
+  """Traverse a proto file definition.
+
+  Args:
+    file_proto: FileDescriptorProto for file.
+    visitor: visitor.Visitor defining the business logic of the plugin.
+
+  Returns:
+    Plugin specific output.
+  """
+  source_code_info = type_context.SourceCodeInfo(file_proto.name, file_proto.source_code_info)
+  package_type_context = type_context.TypeContext(source_code_info, file_proto.package)
+  msgs = [
+      TraverseMessage(package_type_context.ExtendMessage(index, msg.name), msg, visitor)
+      for index, msg in enumerate(file_proto.message_type)
+  ]
+  enums = [
+      TraverseEnum(package_type_context.ExtendEnum(index, enum.name), enum, visitor)
+      for index, enum in enumerate(file_proto.enum_type)
+  ]
+  return visitor.VisitFile(file_proto, package_type_context, msgs, enums)
diff --git a/tools/api_proto_plugin/type_context.py b/tools/api_proto_plugin/type_context.py
new file mode 100644
index 000000000000..d69c120a1a63
--- /dev/null
+++ b/tools/api_proto_plugin/type_context.py
@@ -0,0 +1,195 @@
+"""Type context for FileDescriptorProto traversal."""
+
+from collections import namedtuple
+
+from tools.api_proto_plugin import annotations
+
+# A comment is a (raw comment, annotation map) pair.
+Comment = namedtuple('Comment', ['raw', 'annotations'])
+
+
+class SourceCodeInfo(object):
+  """Wrapper for SourceCodeInfo proto."""
+
+  def __init__(self, name, source_code_info):
+    self.name = name
+    self.proto = source_code_info
+    # Map from path to SourceCodeInfo.Location
+    self._locations = {str(location.path): location for location in self.proto.location}
+    self._file_level_comments = None
+    self._file_level_annotations = None
+
+  @property
+  def file_level_comments(self):
+    """Obtain inferred file level comment."""
+    if self._file_level_comments:
+      return self._file_level_comments
+    comments = []
+    # We find the earliest detached comment by first finding the maximum start
+    # line for any location and then scanning for any earlier locations with
+    # detached comments.
+    earliest_detached_comment = max(location.span[0] for location in self.proto.location) + 1
+    for location in self.proto.location:
+      if location.leading_detached_comments and location.span[0] < earliest_detached_comment:
+        comments = location.leading_detached_comments
+        earliest_detached_comment = location.span[0]
+    self._file_level_comments = comments
+    return comments
+
+  @property
+  def file_level_annotations(self):
+    """Obtain inferred file level annotations."""
+    if self._file_level_annotations:
+      return self._file_level_annotations
+    self._file_level_annotations = dict(
+        sum([list(annotations.ExtractAnnotations(c).items()) for c in self.file_level_comments],
+            []))
+    return self._file_level_annotations
+
+  def LocationPathLookup(self, path):
+    """Lookup SourceCodeInfo.Location by path in SourceCodeInfo.
+
+    Args:
+      path: a list of path indexes as per
+        https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
+
+    Returns:
+      SourceCodeInfo.Location object if found, otherwise None.
+    """
+    return self._locations.get(str(path), None)
+
+  # TODO(htuch): consider integrating comment lookup with overall
+  # FileDescriptorProto, perhaps via two passes.
+  def LeadingCommentPathLookup(self, path):
+    """Lookup leading comment by path in SourceCodeInfo.
+
+    Args:
+      path: a list of path indexes as per
+        https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
+
+    Returns:
+      Comment object.
+    """
+    location = self.LocationPathLookup(path)
+    if location is not None:
+      return Comment(
+          location.leading_comments,
+          annotations.ExtractAnnotations(location.leading_comments, self.file_level_annotations))
+    return Comment('', {})
+
+
+class TypeContext(object):
+  """Contextual information for a message/field.
+
+  Provides information around namespaces and enclosing types for fields and
+  nested messages/enums.
+  """
+
+  def __init__(self, source_code_info, name):
+    # SourceCodeInfo as per
+    # https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto.
+    self.source_code_info = source_code_info
+    # path: a list of path indexes as per
+    #  https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
+    #  Extended as nested objects are traversed.
+    self.path = []
+    # Message/enum/field name. Extended as nested objects are traversed.
+    self.name = name
+    # Map from type name to the correct type annotation string, e.g. from
+    # ".envoy.api.v2.Foo.Bar" to "map<string, string>". This is lost during
+    # proto synthesis and is dynamically recovered in TraverseMessage.
+    self.map_typenames = {}
+    # Map from a message's oneof index to the fields sharing a oneof.
+    self.oneof_fields = {}
+    # Map from a message's oneof index to the name of oneof.
+    self.oneof_names = {}
+    # Map from a message's oneof index to the "required" bool property.
+    self.oneof_required = {}
+    self.type_name = 'file'
+
+  def _Extend(self, path, type_name, name):
+    if not self.name:
+      extended_name = name
+    else:
+      extended_name = '%s.%s' % (self.name, name)
+    extended = TypeContext(self.source_code_info, extended_name)
+    extended.path = self.path + path
+    extended.type_name = type_name
+    extended.map_typenames = self.map_typenames.copy()
+    extended.oneof_fields = self.oneof_fields.copy()
+    extended.oneof_names = self.oneof_names.copy()
+    extended.oneof_required = self.oneof_required.copy()
+    return extended
+
+  def ExtendMessage(self, index, name):
+    """Extend type context with a message.
+
+    Args:
+      index: message index in file.
+      name: message name.
+    """
+    return self._Extend([4, index], 'message', name)
+
+  def ExtendNestedMessage(self, index, name):
+    """Extend type context with a nested message.
+
+    Args:
+      index: nested message index in message.
+      name: message name.
+    """
+    return self._Extend([3, index], 'message', name)
+
+  def ExtendField(self, index, name):
+    """Extend type context with a field.
+
+    Args:
+      index: field index in message.
+      name: field name.
+    """
+    return self._Extend([2, index], 'field', name)
+
+  def ExtendEnum(self, index, name):
+    """Extend type context with an enum.
+
+    Args:
+      index: enum index in file.
+      name: enum name.
+    """
+    return self._Extend([5, index], 'enum', name)
+
+  def ExtendNestedEnum(self, index, name):
+    """Extend type context with a nested enum.
+
+    Args:
+      index: enum index in message.
+      name: enum name.
+    """
+    return self._Extend([4, index], 'enum', name)
+
+  def ExtendEnumValue(self, index, name):
+    """Extend type context with an enum enum.
+
+    Args:
+      index: enum value index in enum.
+      name: value name.
+    """
+    return self._Extend([2, index], 'enum_value', name)
+
+  def ExtendOneof(self, index, name):
+    """Extend type context with an oneof declaration.
+
+    Args:
+      index: oneof index in oneof_decl.
+      name: oneof name.
+    """
+    return self._Extend([8, index], 'oneof', name)
+
+  @property
+  def location(self):
+    """SourceCodeInfo.Location for type context."""
+    return self.source_code_info.LocationPathLookup(self.path)
+
+  @property
+  def leading_comment(self):
+    """Leading comment for type context."""
+    return self.source_code_info.LeadingCommentPathLookup(self.path)
diff --git a/tools/api_proto_plugin/visitor.py b/tools/api_proto_plugin/visitor.py
new file mode 100644
index 000000000000..0065537f0a6e
--- /dev/null
+++ b/tools/api_proto_plugin/visitor.py
@@ -0,0 +1,45 @@
+"""FileDescriptorProto visitor interface for api_proto_plugin implementations."""
+
+
+class Visitor(object):
+  """Abstract visitor interface for api_proto_plugin implementation."""
+
+  def VisitEnum(self, enum_proto, type_context):
+    """Visit an enum definition.
+
+    Args:
+      enum_proto: EnumDescriptorProto for enum.
+      type_context: type_context.TypeContext for enum type.
+
+    Returns:
+      Plugin specific output.
+    """
+    pass
+
+  def VisitMessage(self, msg_proto, type_context, nested_msgs, nested_enums):
+    """Visit a message definition.
+
+    Args:
+      msg_proto: DescriptorProto for message.
+      type_context: type_context.TypeContext for message type.
+      nested_msgs: a list of results from visiting nested messages.
+      nested_enums: a list of results from visiting nested enums.
+
+    Returns:
+      Plugin specific output.
+    """
+    pass
+
+  def VisitFile(self, file_proto, type_context, msgs, enums):
+    """Visit a proto file definition.
+
+    Args:
+      file_proto: FileDescriptorProto for file.
+      type_context: type_context.TypeContext for file.
+      msgs: a list of results from visiting messages.
+      enums: a list of results from visiting enums.
+
+    Returns:
+      Plugin specific output.
+    """
+    pass
diff --git a/tools/protodoc/BUILD b/tools/protodoc/BUILD
index b4b3c3f39acb..d2c9b12a6727 100644
--- a/tools/protodoc/BUILD
+++ b/tools/protodoc/BUILD
@@ -6,6 +6,7 @@ py_binary(
     python_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
+        "//tools/api_proto_plugin",
         "@com_envoyproxy_protoc_gen_validate//validate:validate_py",
         "@com_google_protobuf//:protobuf_python",
     ],
diff --git a/tools/protodoc/protodoc.py b/tools/protodoc/protodoc.py
index 83bf0946a961..e386757a5228 100755
--- a/tools/protodoc/protodoc.py
+++ b/tools/protodoc/protodoc.py
@@ -4,15 +4,14 @@
 # https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html for Sphinx RST syntax.
 
 from collections import defaultdict
-import cProfile
 import functools
-import io
 import os
-import pstats
 import re
-import sys
 
-from google.protobuf.compiler import plugin_pb2
+from tools.api_proto_plugin import annotations
+from tools.api_proto_plugin import plugin
+from tools.api_proto_plugin import visitor
+
 from validate import validate_pb2
 
 # Namespace prefix for Envoy core APIs.
@@ -30,65 +29,55 @@
 # http://www.fileformat.info/info/unicode/char/2063/index.htm
 UNICODE_INVISIBLE_SEPARATOR = u'\u2063'
 
-# Key-value annotation regex.
-ANNOTATION_REGEX = re.compile('\[#([\w-]+?):(.*?)\]\s?', re.DOTALL)
-
-# Page/section titles with special prefixes in the proto comments
-DOC_TITLE_ANNOTATION = 'protodoc-title'
-
-# Not implemented yet annotation on leading comments, leading to insertion of
-# warning on field.
-NOT_IMPLEMENTED_WARN_ANNOTATION = 'not-implemented-warn'
+# Template for data plane API URLs.
+DATA_PLANE_API_URL_FMT = 'https://github.com/envoyproxy/envoy/blob/{}/api/%s#L%d'.format(
+    os.environ['ENVOY_BLOB_SHA'])
 
-# Not implemented yet annotation on leading comments, leading to hiding of
-# field.
-NOT_IMPLEMENTED_HIDE_ANNOTATION = 'not-implemented-hide'
 
-# Comment that allows for easy searching for things that need cleaning up in the next major
-# API version.
-NEXT_MAJOR_VERSION_ANNOTATION = 'next-major-version'
+class ProtodocError(Exception):
+  """Base error class for the protodoc module."""
 
-# Comment. Just used for adding text that will not go into the docs at all.
-COMMENT_ANNOTATION = 'comment'
 
-# proto compatibility status.
-PROTO_STATUS_ANNOTATION = 'proto-status'
+def HideNotImplemented(comment):
+  """Should a given type_context.Comment be hidden because it is tagged as [#not-implemented-hide:]?"""
+  return annotations.NOT_IMPLEMENTED_HIDE_ANNOTATION in comment.annotations
 
-# Where v2 differs from v1..
-V2_API_DIFF_ANNOTATION = 'v2-api-diff'
 
-VALID_ANNOTATIONS = set([
-    DOC_TITLE_ANNOTATION,
-    NOT_IMPLEMENTED_WARN_ANNOTATION,
-    NOT_IMPLEMENTED_HIDE_ANNOTATION,
-    V2_API_DIFF_ANNOTATION,
-    NEXT_MAJOR_VERSION_ANNOTATION,
-    COMMENT_ANNOTATION,
-    PROTO_STATUS_ANNOTATION,
-])
+def GithubUrl(type_context):
+  """Obtain data plane API Github URL by path from a TypeContext.
 
-# These can propagate from file scope to message/enum scope (and be overridden).
-INHERITED_ANNOTATIONS = set([
-    PROTO_STATUS_ANNOTATION,
-])
+  Args:
+    type_context: type_context.TypeContext for node.
 
-# Template for data plane API URLs.
-DATA_PLANE_API_URL_FMT = 'https://github.com/envoyproxy/envoy/blob/{}/api/%s#L%d'.format(
-    os.environ['ENVOY_BLOB_SHA'])
+  Returns:
+    A string with a corresponding data plane API GitHub Url.
+  """
+  if type_context.location is not None:
+    return DATA_PLANE_API_URL_FMT % (type_context.source_code_info.name,
+                                     type_context.location.span[0])
+  return ''
 
 
-class ProtodocError(Exception):
-  """Base error class for the protodoc module."""
+def FormatCommentWithAnnotations(comment, type_name=''):
+  """Format a comment string with additional RST for annotations.
 
+  Args:
+    comment: comment string.
+    type_name: optional, 'message' or 'enum' may be specified for additional
+      message/enum specific annotations.
 
-def FormatCommentWithAnnotations(s, annotations, type_name):
-  if NOT_IMPLEMENTED_WARN_ANNOTATION in annotations:
+  Returns:
+    A string with additional RST from annotations.
+  """
+  s = annotations.WithoutAnnotations(StripLeadingSpace(comment.raw) + '\n')
+  if annotations.NOT_IMPLEMENTED_WARN_ANNOTATION in comment.annotations:
     s += '\n.. WARNING::\n  Not implemented yet\n'
-  if V2_API_DIFF_ANNOTATION in annotations:
-    s += '\n.. NOTE::\n  **v2 API difference**: ' + annotations[V2_API_DIFF_ANNOTATION] + '\n'
+  if annotations.V2_API_DIFF_ANNOTATION in comment.annotations:
+    s += '\n.. NOTE::\n  **v2 API difference**: ' + comment.annotations[
+        annotations.V2_API_DIFF_ANNOTATION] + '\n'
   if type_name == 'message' or type_name == 'enum':
-    if PROTO_STATUS_ANNOTATION in annotations:
-      status = annotations[PROTO_STATUS_ANNOTATION]
+    if annotations.PROTO_STATUS_ANNOTATION in comment.annotations:
+      status = comment.annotations[annotations.PROTO_STATUS_ANNOTATION]
       if status not in ['frozen', 'draft', 'experimental']:
         raise ProtodocError('Unknown proto status: %s' % status)
       if status == 'draft' or status == 'experimental':
@@ -97,209 +86,13 @@ def FormatCommentWithAnnotations(s, annotations, type_name):
   return s
 
 
-def ExtractAnnotations(s, inherited_annotations=None, type_name='file'):
-  """Extract annotations from a given comment string.
-
-  Args:
-    s: string that may contains annotations.
-    inherited_annotations: annotation map from file-level inherited annotations
-      (or None) if this is a file-level comment.
-  Returns:
-    Pair of string with with annotations stripped and annotation map.
-  """
-  annotations = {
-      k: v for k, v in (inherited_annotations or {}).items() if k in INHERITED_ANNOTATIONS
-  }
-  # Extract annotations.
-  groups = re.findall(ANNOTATION_REGEX, s)
-  # Remove annotations.
-  without_annotations = re.sub(ANNOTATION_REGEX, '', s)
-  for group in groups:
-    annotation = group[0]
-    if annotation not in VALID_ANNOTATIONS:
-      raise ProtodocError('Unknown annotation: %s' % annotation)
-    annotations[group[0]] = group[1].lstrip()
-  return FormatCommentWithAnnotations(without_annotations, annotations, type_name), annotations
-
-
-class SourceCodeInfo(object):
-  """Wrapper for SourceCodeInfo proto."""
-
-  def __init__(self, name, source_code_info):
-    self._name = name
-    self._proto = source_code_info
-    self._leading_comments = {
-        str(location.path): location.leading_comments for location in self._proto.location
-    }
-    self._file_level_comment = None
-
-  @property
-  def file_level_comment(self):
-    """Obtain inferred file level comment."""
-    if self._file_level_comment:
-      return self._file_level_comment
-    comment = ''
-    earliest_detached_comment = max(max(location.span) for location in self._proto.location)
-    for location in self._proto.location:
-      if location.leading_detached_comments and location.span[0] < earliest_detached_comment:
-        comment = StripLeadingSpace(''.join(location.leading_detached_comments)) + '\n'
-        earliest_detached_comment = location.span[0]
-    self._file_level_comment = comment
-    return comment
-
-  def LeadingCommentPathLookup(self, path, type_name):
-    """Lookup leading comment by path in SourceCodeInfo.
-
-    Args:
-      path: a list of path indexes as per
-        https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
-      type_name: name of type the comment belongs to.
-    Returns:
-      Pair of attached leading comment and Annotation objects, where there is a
-      leading comment
-      otherwise ('', []).
-    """
-    leading_comment = self._leading_comments.get(str(path), None)
-    if leading_comment is not None:
-      _, file_annotations = ExtractAnnotations(self.file_level_comment)
-      return ExtractAnnotations(
-          StripLeadingSpace(leading_comment) + '\n', file_annotations, type_name)
-    return '', []
-
-  def GithubUrl(self, path):
-    """Obtain data plane API Github URL by path from SourceCodeInfo.
-
-    Args:
-      path: a list of path indexes as per
-        https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
-    Returns:
-      A string with a corresponding data plane API GitHub Url.
-    """
-    for location in self._proto.location:
-      if location.path == path:
-        return DATA_PLANE_API_URL_FMT % (self._name, location.span[0])
-    return ''
-
-
-class TypeContext(object):
-  """Contextual information for a message/field.
-
-  Provides information around namespaces and enclosing types for fields and
-  nested messages/enums.
-  """
-
-  def __init__(self, source_code_info, name):
-    # SourceCodeInfo as per
-    # https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto.
-    self.source_code_info = source_code_info
-    # path: a list of path indexes as per
-    #  https://github.com/google/protobuf/blob/a08b03d4c00a5793b88b494f672513f6ad46a681/src/google/protobuf/descriptor.proto#L717.
-    #  Extended as nested objects are traversed.
-    self.path = []
-    # Message/enum/field name. Extended as nested objects are traversed.
-    self.name = name
-    # Map from type name to the correct type annotation string, e.g. from
-    # ".envoy.api.v2.Foo.Bar" to "map<string, string>". This is lost during
-    # proto synthesis and is dynamically recovered in FormatMessage.
-    self.map_typenames = {}
-    # Map from a message's oneof index to the fields sharing a oneof.
-    self.oneof_fields = {}
-    # Map from a message's oneof index to the name of oneof.
-    self.oneof_names = {}
-    # Map from a message's oneof index to the "required" bool property.
-    self.oneof_required = {}
-    self.type_name = 'file'
-
-  def _Extend(self, path, type_name, name):
-    if not self.name:
-      extended_name = name
-    else:
-      extended_name = '%s.%s' % (self.name, name)
-    extended = TypeContext(self.source_code_info, extended_name)
-    extended.path = self.path + path
-    extended.type_name = type_name
-    extended.map_typenames = self.map_typenames.copy()
-    extended.oneof_fields = self.oneof_fields.copy()
-    extended.oneof_names = self.oneof_names.copy()
-    extended.oneof_required = self.oneof_required.copy()
-    return extended
-
-  def ExtendMessage(self, index, name):
-    """Extend type context with a message.
-
-    Args:
-      index: message index in file.
-      name: message name.
-    """
-    return self._Extend([4, index], 'message', name)
-
-  def ExtendNestedMessage(self, index, name):
-    """Extend type context with a nested message.
-
-    Args:
-      index: nested message index in message.
-      name: message name.
-    """
-    return self._Extend([3, index], 'message', name)
-
-  def ExtendField(self, index, name):
-    """Extend type context with a field.
-
-    Args:
-      index: field index in message.
-      name: field name.
-    """
-    return self._Extend([2, index], 'field', name)
-
-  def ExtendEnum(self, index, name):
-    """Extend type context with an enum.
-
-    Args:
-      index: enum index in file.
-      name: enum name.
-    """
-    return self._Extend([5, index], 'enum', name)
-
-  def ExtendNestedEnum(self, index, name):
-    """Extend type context with a nested enum.
-
-    Args:
-      index: enum index in message.
-      name: enum name.
-    """
-    return self._Extend([4, index], 'enum', name)
-
-  def ExtendEnumValue(self, index, name):
-    """Extend type context with an enum enum.
-
-    Args:
-      index: enum value index in enum.
-      name: value name.
-    """
-    return self._Extend([2, index], 'enum_value', name)
-
-  def ExtendOneof(self, index, name):
-    """Extend type context with an oneof declaration.
-
-    Args:
-      index: oneof index in oneof_decl.
-      name: oneof name.
-    """
-    return self._Extend([8, index], "oneof", name)
-
-  def LeadingCommentPathLookup(self):
-    return self.source_code_info.LeadingCommentPathLookup(self.path, self.type_name)
-
-  def GithubUrl(self):
-    return self.source_code_info.GithubUrl(self.path)
-
-
 def MapLines(f, s):
   """Apply a function across each line in a flat string.
 
   Args:
     f: A string transform function for a line.
     s: A string consisting of potentially multiple lines.
+
   Returns:
     A flat string with f applied to each line.
   """
@@ -330,28 +123,33 @@ def FormatHeader(style, text):
   Args:
     style: underline style, e.g. '=', '-'.
     text: header text
+
   Returns:
     RST formatted header.
   """
   return '%s\n%s\n\n' % (text, style * len(text))
 
 
-def FormatHeaderFromFile(style, file_level_comment, alt):
+def FormatHeaderFromFile(style, source_code_info, proto_name):
   """Format RST header based on special file level title
 
   Args:
     style: underline style, e.g. '=', '-'.
-    file_level_comment: detached comment at top of file.
-    alt: If the file_level_comment does not contain a user
-         specified title, use the alt text as page title.
+    source_code_info: SourceCodeInfo object.
+    proto_name: If the file_level_comment does not contain a user specified
+      title, use this as page title.
+
   Returns:
     RST formatted header, and file level comment without page title strings.
   """
-  anchor = FormatAnchor(FileCrossRefLabel(alt))
-  stripped_comment, annotations = ExtractAnnotations(file_level_comment)
-  if DOC_TITLE_ANNOTATION in annotations:
-    return anchor + FormatHeader(style, annotations[DOC_TITLE_ANNOTATION]), stripped_comment
-  return anchor + FormatHeader(style, alt), stripped_comment
+  anchor = FormatAnchor(FileCrossRefLabel(proto_name))
+  stripped_comment = annotations.WithoutAnnotations(
+      StripLeadingSpace('\n'.join(c + '\n' for c in source_code_info.file_level_comments)))
+  if annotations.DOC_TITLE_ANNOTATION in source_code_info.file_level_annotations:
+    return anchor + FormatHeader(
+        style,
+        source_code_info.file_level_annotations[annotations.DOC_TITLE_ANNOTATION]), stripped_comment
+  return anchor + FormatHeader(style, proto_name), stripped_comment
 
 
 def FormatFieldTypeAsJson(type_context, field):
@@ -360,10 +158,9 @@ def FormatFieldTypeAsJson(type_context, field):
   Args:
     type_context: contextual information for message/enum/field.
     field: FieldDescriptor proto.
-  Return:
-    RST formatted pseudo-JSON string representation of field type.
+  Return: RST formatted pseudo-JSON string representation of field type.
   """
-  if NormalizeFQN(field.type_name) in type_context.map_typenames:
+  if TypeNameFromFQN(field.type_name) in type_context.map_typenames:
     return '"{...}"'
   if field.label == field.LABEL_REPEATED:
     return '[]'
@@ -378,14 +175,13 @@ def FormatMessageAsJson(type_context, msg):
   Args:
     type_context: contextual information for message/enum/field.
     msg: message definition DescriptorProto.
-  Return:
-    RST formatted pseudo-JSON string representation of message definition.
+  Return: RST formatted pseudo-JSON string representation of message definition.
   """
   lines = []
   for index, field in enumerate(msg.field):
     field_type_context = type_context.ExtendField(index, field.name)
-    leading_comment, comment_annotations = field_type_context.LeadingCommentPathLookup()
-    if NOT_IMPLEMENTED_HIDE_ANNOTATION in comment_annotations:
+    leading_comment = field_type_context.leading_comment
+    if HideNotImplemented(leading_comment):
       continue
     lines.append('"%s": %s' % (field.name, FormatFieldTypeAsJson(type_context, field)))
 
@@ -395,21 +191,44 @@ def FormatMessageAsJson(type_context, msg):
     return '.. code-block:: json\n\n  {}\n\n'
 
 
-def NormalizeFQN(fqn):
-  """Normalize a fully qualified field type name.
+def NormalizeFieldTypeName(field_fqn):
+  """Normalize a fully qualified field type name, e.g.
+
+  .envoy.foo.bar.
+
+  Strips leading ENVOY_API_NAMESPACE_PREFIX and ENVOY_PREFIX.
 
-  Strips leading ENVOY_API_NAMESPACE_PREFIX and ENVOY_PREFIX and makes pretty wrapped type names.
+  Args:
+    field_fqn: a fully qualified type name from FieldDescriptorProto.type_name.
+  Return: Normalized type name.
+  """
+  if field_fqn.startswith(ENVOY_API_NAMESPACE_PREFIX):
+    return field_fqn[len(ENVOY_API_NAMESPACE_PREFIX):]
+  if field_fqn.startswith(ENVOY_PREFIX):
+    return field_fqn[len(ENVOY_PREFIX):]
+  return field_fqn
+
+
+def NormalizeTypeContextName(type_name):
+  """Normalize a type name, e.g.
+
+  envoy.foo.bar.
+
+  Strips leading ENVOY_API_NAMESPACE_PREFIX and ENVOY_PREFIX.
 
   Args:
-    fqn: a fully qualified type name from FieldDescriptorProto.type_name.
-  Return:
-    Normalized type name.
+    type_name: a name from a TypeContext.
+  Return: Normalized type name.
   """
-  if fqn.startswith(ENVOY_API_NAMESPACE_PREFIX):
-    return fqn[len(ENVOY_API_NAMESPACE_PREFIX):]
-  if fqn.startswith(ENVOY_PREFIX):
-    return fqn[len(ENVOY_PREFIX):]
-  return fqn
+  return NormalizeFieldTypeName(QualifyTypeName(type_name))
+
+
+def QualifyTypeName(type_name):
+  return '.' + type_name
+
+
+def TypeNameFromFQN(fqn):
+  return fqn[1:]
 
 
 def FormatEmph(s):
@@ -426,15 +245,17 @@ def FormatFieldType(type_context, field):
   Args:
     type_context: contextual information for message/enum/field.
     field: FieldDescriptor proto.
-  Return:
-    RST formatted field type.
+  Return: RST formatted field type.
   """
   if field.type_name.startswith(ENVOY_API_NAMESPACE_PREFIX) or field.type_name.startswith(
       ENVOY_PREFIX):
-    type_name = NormalizeFQN(field.type_name)
+    type_name = NormalizeFieldTypeName(field.type_name)
     if field.type == field.TYPE_MESSAGE:
-      if type_context.map_typenames and type_name in type_context.map_typenames:
-        return type_context.map_typenames[type_name]
+      if type_context.map_typenames and TypeNameFromFQN(
+          field.type_name) in type_context.map_typenames:
+        return 'map<%s, %s>' % tuple(
+            map(functools.partial(FormatFieldType, type_context),
+                type_context.map_typenames[TypeNameFromFQN(field.type_name)]))
       return FormatInternalLink(type_name, MessageCrossRefLabel(type_name))
     if field.type == field.TYPE_ENUM:
       return FormatInternalLink(type_name, EnumCrossRefLabel(type_name))
@@ -516,49 +337,55 @@ def FormatFieldAsDefinitionListItem(outer_type_context, type_context, field):
     outer_type_context: contextual information for enclosing message.
     type_context: contextual information for message/enum/field.
     field: FieldDescriptorProto.
+
   Returns:
     RST formatted definition list item.
   """
-  annotations = []
+  field_annotations = []
 
-  anchor = FormatAnchor(FieldCrossRefLabel(type_context.name))
+  anchor = FormatAnchor(FieldCrossRefLabel(NormalizeTypeContextName(type_context.name)))
   if field.options.HasExtension(validate_pb2.rules):
     rule = field.options.Extensions[validate_pb2.rules]
     if ((rule.HasField('message') and rule.message.required) or
         (rule.HasField('string') and rule.string.min_bytes > 0) or
         (rule.HasField('repeated') and rule.repeated.min_items > 0)):
-      annotations = ['*REQUIRED*']
-  leading_comment, comment_annotations = type_context.LeadingCommentPathLookup()
-  if NOT_IMPLEMENTED_HIDE_ANNOTATION in comment_annotations:
+      field_annotations = ['*REQUIRED*']
+  leading_comment = type_context.leading_comment
+  formatted_leading_comment = FormatCommentWithAnnotations(leading_comment)
+  if HideNotImplemented(leading_comment):
     return ''
 
   if field.HasField('oneof_index'):
     oneof_context = outer_type_context.ExtendOneof(field.oneof_index,
                                                    type_context.oneof_names[field.oneof_index])
-    oneof_comment, oneof_comment_annotations = oneof_context.LeadingCommentPathLookup()
-    if NOT_IMPLEMENTED_HIDE_ANNOTATION in oneof_comment_annotations:
+    oneof_comment = oneof_context.leading_comment
+    formatted_oneof_comment = FormatCommentWithAnnotations(oneof_comment)
+    if HideNotImplemented(oneof_comment):
       return ''
 
     # If the oneof only has one field and marked required, mark the field as required.
     if len(type_context.oneof_fields[field.oneof_index]) == 1 and type_context.oneof_required[
         field.oneof_index]:
-      annotations = ['*REQUIRED*']
+      field_annotations = ['*REQUIRED*']
 
     if len(type_context.oneof_fields[field.oneof_index]) > 1:
       # Fields in oneof shouldn't be marked as required when we have oneof comment below it.
-      annotations = []
+      field_annotations = []
       oneof_template = '\nPrecisely one of %s must be set.\n' if type_context.oneof_required[
           field.oneof_index] else '\nOnly one of %s may be set.\n'
-      oneof_comment += oneof_template % ', '.join(
-          FormatInternalLink(f, FieldCrossRefLabel(outer_type_context.ExtendField(i, f).name))
+      formatted_oneof_comment += oneof_template % ', '.join(
+          FormatInternalLink(
+              f,
+              FieldCrossRefLabel(NormalizeTypeContextName(
+                  outer_type_context.ExtendField(i, f).name)))
           for i, f in type_context.oneof_fields[field.oneof_index])
   else:
-    oneof_comment = ''
+    formatted_oneof_comment = ''
 
   comment = '(%s) ' % ', '.join([FormatFieldType(type_context, field)] +
-                                annotations) + leading_comment
+                                field_annotations) + formatted_leading_comment
   return anchor + field.name + '\n' + MapLines(functools.partial(Indent, 2),
-                                               comment + oneof_comment)
+                                               comment + formatted_oneof_comment)
 
 
 def FormatMessageAsDefinitionList(type_context, msg):
@@ -567,6 +394,7 @@ def FormatMessageAsDefinitionList(type_context, msg):
   Args:
     type_context: contextual information for message/enum/field.
     msg: DescriptorProto.
+
   Returns:
     RST formatted definition list item.
   """
@@ -575,9 +403,8 @@ def FormatMessageAsDefinitionList(type_context, msg):
   type_context.oneof_names = defaultdict(list)
   for index, field in enumerate(msg.field):
     if field.HasField('oneof_index'):
-      _, comment_annotations = type_context.ExtendField(index,
-                                                        field.name).LeadingCommentPathLookup()
-      if NOT_IMPLEMENTED_HIDE_ANNOTATION in comment_annotations:
+      leading_comment = type_context.ExtendField(index, field.name).leading_comment
+      if HideNotImplemented(leading_comment):
         continue
       type_context.oneof_fields[field.oneof_index].append((index, field.name))
   for index, oneof_decl in enumerate(msg.oneof_decl):
@@ -589,59 +416,23 @@ def FormatMessageAsDefinitionList(type_context, msg):
                                       field) for index, field in enumerate(msg.field)) + '\n'
 
 
-def FormatMessage(type_context, msg):
-  """Format a DescriptorProto as RST section.
-
-  Args:
-    type_context: contextual information for message/enum/field.
-    msg: DescriptorProto.
-  Returns:
-    RST formatted section.
-  """
-  # Skip messages synthesized to represent map types.
-  if msg.options.map_entry:
-    return ''
-  # We need to do some extra work to recover the map type annotation from the
-  # synthesized messages.
-  type_context.map_typenames = {
-      '%s.%s' % (type_context.name, nested_msg.name):
-      'map<%s, %s>' % tuple(map(functools.partial(FormatFieldType, type_context), nested_msg.field))
-      for nested_msg in msg.nested_type
-      if nested_msg.options.map_entry
-  }
-  nested_msgs = '\n'.join(
-      FormatMessage(type_context.ExtendNestedMessage(index, nested_msg.name), nested_msg)
-      for index, nested_msg in enumerate(msg.nested_type))
-  nested_enums = '\n'.join(
-      FormatEnum(type_context.ExtendNestedEnum(index, nested_enum.name), nested_enum)
-      for index, nested_enum in enumerate(msg.enum_type))
-  anchor = FormatAnchor(MessageCrossRefLabel(type_context.name))
-  header = FormatHeader('-', type_context.name)
-  proto_link = FormatExternalLink('[%s proto]' % type_context.name,
-                                  type_context.GithubUrl()) + '\n\n'
-  leading_comment, annotations = type_context.LeadingCommentPathLookup()
-  if NOT_IMPLEMENTED_HIDE_ANNOTATION in annotations:
-    return ''
-  return anchor + header + proto_link + leading_comment + FormatMessageAsJson(
-      type_context, msg) + FormatMessageAsDefinitionList(type_context,
-                                                         msg) + nested_msgs + '\n' + nested_enums
-
-
 def FormatEnumValueAsDefinitionListItem(type_context, enum_value):
   """Format a EnumValueDescriptorProto as RST definition list item.
 
   Args:
     type_context: contextual information for message/enum/field.
     enum_value: EnumValueDescriptorProto.
+
   Returns:
     RST formatted definition list item.
   """
-  anchor = FormatAnchor(EnumValueCrossRefLabel(type_context.name))
+  anchor = FormatAnchor(EnumValueCrossRefLabel(NormalizeTypeContextName(type_context.name)))
   default_comment = '*(DEFAULT)* ' if enum_value.number == 0 else ''
-  leading_comment, annotations = type_context.LeadingCommentPathLookup()
-  if NOT_IMPLEMENTED_HIDE_ANNOTATION in annotations:
+  leading_comment = type_context.leading_comment
+  formatted_leading_comment = FormatCommentWithAnnotations(leading_comment)
+  if HideNotImplemented(leading_comment):
     return ''
-  comment = default_comment + UNICODE_INVISIBLE_SEPARATOR + leading_comment
+  comment = default_comment + UNICODE_INVISIBLE_SEPARATOR + formatted_leading_comment
   return anchor + enum_value.name + '\n' + MapLines(functools.partial(Indent, 2), comment)
 
 
@@ -651,6 +442,7 @@ def FormatEnumAsDefinitionList(type_context, enum):
   Args:
     type_context: contextual information for message/enum/field.
     enum: DescriptorProto.
+
   Returns:
     RST formatted definition list item.
   """
@@ -660,88 +452,57 @@ def FormatEnumAsDefinitionList(type_context, enum):
       for index, enum_value in enumerate(enum.value)) + '\n'
 
 
-def FormatEnum(type_context, enum):
-  """Format an EnumDescriptorProto as RST section.
+def FormatProtoAsBlockComment(proto):
+  """Format a proto as a RST block comment.
 
-  Args:
-    type_context: contextual information for message/enum/field.
-    enum: EnumDescriptorProto.
-  Returns:
-    RST formatted section.
+  Useful in debugging, not usually referenced.
   """
-  anchor = FormatAnchor(EnumCrossRefLabel(type_context.name))
-  header = FormatHeader('-', 'Enum %s' % type_context.name)
-  proto_link = FormatExternalLink('[%s proto]' % type_context.name,
-                                  type_context.GithubUrl()) + '\n\n'
-  leading_comment, annotations = type_context.LeadingCommentPathLookup()
-  if NOT_IMPLEMENTED_HIDE_ANNOTATION in annotations:
-    return ''
-  return anchor + header + proto_link + leading_comment + FormatEnumAsDefinitionList(
-      type_context, enum)
+  return '\n\nproto::\n\n' + MapLines(functools.partial(Indent, 2), str(proto)) + '\n'
 
 
-def FormatProtoAsBlockComment(proto):
-  """Format as RST a proto as a block comment.
+class RstFormatVisitor(visitor.Visitor):
+  """Visitor to generate a RST representation from a FileDescriptor proto.
 
-  Useful in debugging, not usually referenced.
+  See visitor.Visitor for visitor method docs comments.
   """
-  return '\n\nproto::\n\n' + MapLines(functools.partial(Indent, 2), str(proto)) + '\n'
 
+  def VisitEnum(self, enum_proto, type_context):
+    normal_enum_type = NormalizeTypeContextName(type_context.name)
+    anchor = FormatAnchor(EnumCrossRefLabel(normal_enum_type))
+    header = FormatHeader('-', 'Enum %s' % normal_enum_type)
+    github_url = GithubUrl(type_context)
+    proto_link = FormatExternalLink('[%s proto]' % normal_enum_type, github_url) + '\n\n'
+    leading_comment = type_context.leading_comment
+    formatted_leading_comment = FormatCommentWithAnnotations(leading_comment, 'enum')
+    if HideNotImplemented(leading_comment):
+      return ''
+    return anchor + header + proto_link + formatted_leading_comment + FormatEnumAsDefinitionList(
+        type_context, enum_proto)
+
+  def VisitMessage(self, msg_proto, type_context, nested_msgs, nested_enums):
+    normal_msg_type = NormalizeTypeContextName(type_context.name)
+    anchor = FormatAnchor(MessageCrossRefLabel(normal_msg_type))
+    header = FormatHeader('-', normal_msg_type)
+    github_url = GithubUrl(type_context)
+    proto_link = FormatExternalLink('[%s proto]' % normal_msg_type, github_url) + '\n\n'
+    leading_comment = type_context.leading_comment
+    formatted_leading_comment = FormatCommentWithAnnotations(leading_comment, 'message')
+    if HideNotImplemented(leading_comment):
+      return ''
+    return anchor + header + proto_link + formatted_leading_comment + FormatMessageAsJson(
+        type_context, msg_proto) + FormatMessageAsDefinitionList(
+            type_context, msg_proto) + '\n'.join(nested_msgs) + '\n' + '\n'.join(nested_enums)
 
-def GenerateRst(proto_file):
-  """Generate a RST representation from a FileDescriptor proto."""
-  source_code_info = SourceCodeInfo(proto_file.name, proto_file.source_code_info)
-  # Find the earliest detached comment, attribute it to file level.
-  # Also extract file level titles if any.
-  header, comment = FormatHeaderFromFile('=', source_code_info.file_level_comment, proto_file.name)
-  package_prefix = NormalizeFQN('.' + proto_file.package + '.')[:-1]
-  package_type_context = TypeContext(source_code_info, package_prefix)
-  msgs = '\n'.join(
-      FormatMessage(package_type_context.ExtendMessage(index, msg.name), msg)
-      for index, msg in enumerate(proto_file.message_type))
-  enums = '\n'.join(
-      FormatEnum(package_type_context.ExtendEnum(index, enum.name), enum)
-      for index, enum in enumerate(proto_file.enum_type))
-  debug_proto = FormatProtoAsBlockComment(proto_file)
-  return header + comment + msgs + enums  # + debug_proto
+  def VisitFile(self, file_proto, type_context, msgs, enums):
+    # Find the earliest detached comment, attribute it to file level.
+    # Also extract file level titles if any.
+    header, comment = FormatHeaderFromFile('=', type_context.source_code_info, file_proto.name)
+    debug_proto = FormatProtoAsBlockComment(file_proto)
+    return header + comment + '\n'.join(msgs) + '\n'.join(enums)  # + debug_proto
 
 
 def Main():
-  # http://www.expobrain.net/2015/09/13/create-a-plugin-for-google-protocol-buffer/
-  request = plugin_pb2.CodeGeneratorRequest()
-  request.ParseFromString(sys.stdin.buffer.read())
-  response = plugin_pb2.CodeGeneratorResponse()
-  cprofile_enabled = os.getenv('CPROFILE_ENABLED')
-
-  # We use file_to_generate rather than proto_file here since we are invoked
-  # inside a Bazel aspect, each node in the DAG will be visited once by the
-  # aspect and we only want to generate docs for the current node.
-  for file_to_generate in request.file_to_generate:
-    # Find the FileDescriptorProto for the file we actually are generating.
-    proto_file = None
-    for pf in request.proto_file:
-      if pf.name == file_to_generate:
-        proto_file = pf
-        break
-    assert (proto_file is not None)
-    f = response.file.add()
-    f.name = proto_file.name + '.rst'
-    if cprofile_enabled:
-      pr = cProfile.Profile()
-      pr.enable()
-    # We don't actually generate any RST right now, we just string dump the
-    # input proto file descriptor into the output file.
-    f.content = GenerateRst(proto_file)
-    if cprofile_enabled:
-      pr.disable()
-      stats_stream = io.StringIO()
-      ps = pstats.Stats(pr,
-                        stream=stats_stream).sort_stats(os.getenv('CPROFILE_SORTBY', 'cumulative'))
-      stats_file = response.file.add()
-      stats_file.name = proto_file.name + '.rst.profile'
-      ps.print_stats()
-      stats_file.content = stats_stream.getvalue()
-  sys.stdout.buffer.write(response.SerializeToString())
+  plugin.Plugin('.rst', RstFormatVisitor())
 
 
 if __name__ == '__main__':