Merge extension with extended type

lib/jazzy/assets/css/jazzy.css.scss

@@ -311,6 +311,11 @@ header {
     padding-left: 3px;
     margin-left: 15px;
   }
+  .declaration-note {
+    font-size: .85em;
+    color: rgba(128,128,128,1);
+    font-style: italic;
+  }
 }
 
 .pointer-container {

lib/jazzy/doc_builder.rb

@@ -234,22 +234,27 @@ def self.render_item(item, source_module)
       # Combine abstract and discussion into abstract
       abstract = (item.abstract || '') + (item.discussion || '')
       item_render = {
-        name: item.name,
-        abstract: Jazzy.markdown.render(abstract),
-        declaration: item.declaration,
-        usr: item.usr,
-        dash_type: item.type.dash_type,
+        name:                    item.name,
+        abstract:                render_markdown(abstract),
+        declaration:             item.declaration,
+        usr:                     item.usr,
+        dash_type:               item.type.dash_type,
+        github_token_url:        gh_token_url(item, source_module),
+        default_impl_abstract:   render_markdown(item.default_impl_abstract),
+        from_protocol_extension: item.from_protocol_extension,
+        return:                  render_markdown(item.return),
+        parameters:              (item.parameters if item.parameters.any?),
+        url:                     (item.url if item.children.any?),
+        start_line:              item.start_line,
+        end_line:                item.end_line,
       }
-      gh_token_url = gh_token_url(item, source_module)
-      item_render[:github_token_url] = gh_token_url
-      item_render[:return] = Jazzy.markdown.render(item.return) if item.return
-      item_render[:parameters] = item.parameters if item.parameters.any?
-      item_render[:url] = item.url if item.children.any?
-      item_render[:start_line] = item.start_line
-      item_render[:end_line] = item.end_line
       item_render.reject { |_, v| v.nil? }
     end
 
+    def self.render_markdown(markdown)
+      Jazzy.markdown.render(markdown) if markdown
+    end
+
     def self.make_task(mark, uid, items)
       {
         name: mark.name,

lib/jazzy/source_declaration.rb

@@ -14,6 +14,8 @@ class SourceDeclaration
     attr_accessor :name
     attr_accessor :declaration
     attr_accessor :abstract
+    attr_accessor :default_impl_abstract
+    attr_accessor :from_protocol_extension
     attr_accessor :discussion
     attr_accessor :return
     attr_accessor :children

lib/jazzy/source_declaration/type.rb

@@ -57,10 +57,18 @@ def declaration?
           kind.start_with?('sourcekitten.source.lang.objc.decl')
       end
 
-      def extension?
+      def swift_extension?
         kind =~ /^source\.lang\.swift\.decl\.extension.*/
       end
 
+      def swift_extensible?
+        kind =~ /^source\.lang\.swift\.decl\.(class|struct|protocol|enum)$/
+      end
+
+      def swift_protocol?
+        kind == 'source.lang.swift.decl.protocol'
+      end
+
       def param?
         # SourceKit strangely categorizes initializer parameters as local
         # variables, so both kinds represent a parameter in jazzy.

lib/jazzy/sourcekitten.rb

@@ -161,7 +161,7 @@ def self.should_document?(doc)
       # Document extensions & enum elements, since we can't tell their ACL.
       type = SourceDeclaration::Type.new(doc['key.kind'])
       return true if type.swift_enum_element?
-      if type.extension?
+      if type.swift_extension?
         return Array(doc['key.substructure']).any? do |subdoc|
           should_document?(subdoc)
         end
@@ -302,11 +302,80 @@ def self.doc_coverage
         (@undocumented_tokens.count + @documented_count)
     end
 
+    # Merges multiple extensions of the same entity into a single document.
+    #
+    # Merges extensions into the protocol/class/struct/enum they extend, if it
+    # occurs in the same project.
+    #
+    # Merges redundant declarations when documenting podspecs.
     def self.deduplicate_declarations(declarations)
-      duplicates = declarations.group_by { |d| [d.usr, d.type.kind] }.values
-      duplicates.map do |decls|
-        decls.first.tap do |d|
-          d.children = deduplicate_declarations(decls.flat_map(&:children).uniq)
+      duplicate_groups = declarations
+                         .group_by { |d| deduplication_key(d) }
+                         .values
+
+      duplicate_groups.map do |group|
+        # Put extended type (if present) before extensions
+        merge_declarations(group)
+      end
+    end
+
+    # Two declarations get merged if they have the same deduplication key.
+    def self.deduplication_key(decl)
+      if decl.type.swift_extensible? || decl.type.swift_extension?
+        [decl.usr]
+      else
+        [decl.usr, decl.type.kind]
+      end
+    end
+
+    # Merges all of the given types and extensions into a single document.
+    def self.merge_declarations(decls)
+      extensions, typedecls = decls.partition { |d| d.type.swift_extension? }
+
+      if typedecls.size > 1
+        warn 'Found conflicting type declarations with the same name, which ' \
+          'may indicate a build issue or a bug in Jazzy: ' +
+          typedecls.map { |t| "#{t.type.name.downcase} #{t.name}" }.join(', ')
+      end
+      typedecl = typedecls.first
+
+      if typedecl && typedecl.type.swift_protocol?
+        merge_default_implementations_into_protocol(typedecl, extensions)
+        mark_members_from_protocol_extension(extensions)
+        extensions.reject! { |ext| ext.children.empty? }
+      end
+
+      decls = typedecls + extensions
+      decls.first.tap do |d|
+        d.children = deduplicate_declarations(decls.flat_map(&:children).uniq)
+      end
+    end
+
+    # If any of the extensions provide default implementations for methods in
+    # the given protocol, merge those members into the protocol doc instead of
+    # keeping them on the extension. These get a “Default implementation”
+    # annotation in the generated docs.
+    def self.merge_default_implementations_into_protocol(protocol, extensions)
+      protocol.children.each do |proto_method|
+        extensions.each do |ext|
+          defaults, ext.children = ext.children.partition do |ext_member|
+            ext_member.name == proto_method.name
+          end
+          unless defaults.empty?
+            proto_method.default_impl_abstract =
+              defaults.flat_map { |d| [d.abstract, d.discussion] }.join("\n\n")
+          end
+        end
+      end
+    end
+
+    # Protocol methods provided only in an extension and not in the protocol
+    # itself are a special beast: they do not use dynamic dispatch. These get an
+    # “Extension method” annotation in the generated docs.
+    def self.mark_members_from_protocol_extension(extensions)
+      extensions.each do |ext|
+        ext.children.each do |ext_member|
+          ext_member.from_protocol_extension = true
         end
       end
     end

lib/jazzy/templates/task.mustache

@@ -17,6 +17,16 @@
         <a name="//apple_ref/{{language_stub}}/{{dash_type}}/{{name}}" class="dashAnchor"></a>
         <a class="token" href="#/{{usr}}">{{name}}</a>
         </code>
+        {{#default_impl_abstract}}
+          <span class="declaration-note">
+            Default implementation
+          </span>
+        {{/default_impl_abstract}}
+        {{#from_protocol_extension}}
+          <span class="declaration-note">
+            Extension method
+          </span>
+        {{/from_protocol_extension}}
       </div>
       <div class="height-container">
         <div class="pointer-container"></div>
@@ -30,6 +40,12 @@
             {{/url}}
           </div>
           {{/abstract}}
+          {{#default_impl_abstract}}
+          <h4>Default Implementation</h4>
+          <div class="default_impl abstract">
+            {{{default_impl_abstract}}}
+          </div>
+          {{/default_impl_abstract}}
           {{#declaration}}
           <div class="declaration">
             <h4>Declaration</h4>