diff --git a/lib/puppet-strings/markdown.rb b/lib/puppet-strings/markdown.rb
index b9e988954..da7577f41 100644
--- a/lib/puppet-strings/markdown.rb
+++ b/lib/puppet-strings/markdown.rb
@@ -11,23 +11,52 @@ module PuppetStrings::Markdown
   require_relative 'markdown/resource_types'
   require_relative 'markdown/puppet_tasks'
   require_relative 'markdown/puppet_plans'
-  require_relative 'markdown/table_of_contents'
+
+  # Get modules that handle collecting and rendering each section.
+  #
+  # @return [Array[module]] The modules
+  def self.groups
+    [
+      PuppetStrings::Markdown::PuppetClasses,
+      PuppetStrings::Markdown::DefinedTypes,
+      PuppetStrings::Markdown::ResourceTypes,
+      PuppetStrings::Markdown::Functions,
+      PuppetStrings::Markdown::DataTypes,
+      PuppetStrings::Markdown::PuppetTasks,
+      PuppetStrings::Markdown::PuppetPlans,
+    ]
+  end
 
   # generates markdown documentation
   # @return [String] markdown doc
   def self.generate
-    final = "# Reference\n\n"
-    final += "<!-- DO NOT EDIT: This document was generated by Puppet Strings -->\n\n"
-    final += PuppetStrings::Markdown::TableOfContents.render
-    final += PuppetStrings::Markdown::PuppetClasses.render
-    final += PuppetStrings::Markdown::DefinedTypes.render
-    final += PuppetStrings::Markdown::ResourceTypes.render
-    final += PuppetStrings::Markdown::Functions.render
-    final += PuppetStrings::Markdown::DataTypes.render
-    final += PuppetStrings::Markdown::PuppetTasks.render
-    final += PuppetStrings::Markdown::PuppetPlans.render
-
-    final
+    output = [
+      "# Reference\n\n",
+      "<!-- DO NOT EDIT: This document was generated by Puppet Strings -->\n\n",
+      "## Table of Contents\n\n",
+    ]
+
+    # Create table of contents
+    template = erb(File.join(__dir__, 'markdown', 'templates', 'table_of_contents.erb'))
+    groups.each do |group|
+      group_name = group.name
+      items = group.items.map { |item| item.toc_info }
+      has_private = items.any? { |item| item[:private] }
+      has_public = items.any? { |item| !item[:private] }
+
+      output << template.result(binding)
+    end
+
+    # Create actual contents
+    groups.each do |group|
+      items = group.items.reject { |item| item.private? }
+      unless items.empty?
+        output << "## #{group.name}\n\n"
+        output.append(items.map { |item| item.render })
+      end
+    end
+
+    output.join('')
   end
 
   # mimicks the behavior of the json render, although path will never be nil
@@ -41,4 +70,17 @@ def self.render(path = nil)
       YARD::Logger.instance.debug "Wrote markdown to #{path}"
     end
   end
+
+  # Helper function to load an ERB template.
+  #
+  # @param [String] path The full path to the template file.
+  # @return [ERB] Template
+  def self.erb(path)
+    if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.6.0')
+      ERB.new(File.read(path), trim_mode: '-')
+    else
+      # This outputs warnings in Ruby 2.6+.
+      ERB.new(File.read(path), nil, '-')
+    end
+  end
 end
diff --git a/lib/puppet-strings/markdown/base.rb b/lib/puppet-strings/markdown/base.rb
index 2a40cf515..c73e8a7c0 100644
--- a/lib/puppet-strings/markdown/base.rb
+++ b/lib/puppet-strings/markdown/base.rb
@@ -199,17 +199,4 @@ def clean_link(input)
       input.tr('^a-zA-Z0-9_-', '-')
     end
   end
-
-  # Helper function to load an ERB template.
-  #
-  # @param [String] path The full path to the template file.
-  # @return [ERB] Template
-  def self.erb(path)
-    if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.6.0')
-      ERB.new(File.read(path), trim_mode: '-')
-    else
-      # This outputs warnings in Ruby 2.6+.
-      ERB.new(File.read(path), nil, '-')
-    end
-  end
 end
diff --git a/lib/puppet-strings/markdown/data_types.rb b/lib/puppet-strings/markdown/data_types.rb
index 0fb1bd02c..d7c4dde0a 100644
--- a/lib/puppet-strings/markdown/data_types.rb
+++ b/lib/puppet-strings/markdown/data_types.rb
@@ -3,41 +3,19 @@
 require_relative 'data_type'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about data types
   module DataTypes
+    def self.name
+      'Data types'
+    end
 
     # @return [Array] list of data types
-    def self.in_dtypes
+    def self.items
       arr = YARD::Registry.all(:puppet_data_type).map!(&:to_hash)
       arr.concat(YARD::Registry.all(:puppet_data_type_alias).map!(&:to_hash))
 
       arr.sort! { |a,b| a[:name] <=> b[:name] }
       arr.map! { |a| PuppetStrings::Markdown::DataType.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_dtypes.nil?
-        in_dtypes.find { |type| type.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_dtypes.length > 0 ? "## Data types\n\n" : ""
-      in_dtypes.each do |type|
-        final += type.render unless type.private?
-      end
-      final
-    end
-
-
-    def self.toc_info
-      final = ["Data types"]
-
-      in_dtypes.each do |type|
-        final.push(type.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/defined_types.rb b/lib/puppet-strings/markdown/defined_types.rb
index a4c281a1c..daf2a9dd5 100644
--- a/lib/puppet-strings/markdown/defined_types.rb
+++ b/lib/puppet-strings/markdown/defined_types.rb
@@ -3,37 +3,16 @@
 require_relative 'defined_type'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about defined types
   module DefinedTypes
+    def self.name
+      'Defined types'
+    end
 
     # @return [Array] list of defined types
-    def self.in_dtypes
+    def self.items
       arr = YARD::Registry.all(:puppet_defined_type).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::DefinedType.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_dtypes.nil?
-        in_dtypes.find { |type| type.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_dtypes.length > 0 ? "## Defined types\n\n" : ""
-      in_dtypes.each do |type|
-        final += type.render unless type.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Defined types"]
-
-      in_dtypes.each do |type|
-        final.push(type.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/functions.rb b/lib/puppet-strings/markdown/functions.rb
index c111478c5..5345261c4 100644
--- a/lib/puppet-strings/markdown/functions.rb
+++ b/lib/puppet-strings/markdown/functions.rb
@@ -3,38 +3,17 @@
 require_relative 'function'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about functions
   module Functions
+    def self.name
+      'Functions'
+    end
 
     # @return [Array] list of functions
-    def self.in_functions
+    def self.items
       arr = YARD::Registry.all(:puppet_function).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::Function.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_functions.nil?
-        in_functions.find { |func| func.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_functions.length > 0 ? "## Functions\n\n" : ""
-      in_functions.each do |func|
-        final += func.render unless func.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Functions"]
-
-      in_functions.each do |func|
-        final.push(func.toc_info)
-      end
-
-      final
-    end
   end
 end
 
diff --git a/lib/puppet-strings/markdown/puppet_classes.rb b/lib/puppet-strings/markdown/puppet_classes.rb
index f22dd5789..6fcf1e7fe 100644
--- a/lib/puppet-strings/markdown/puppet_classes.rb
+++ b/lib/puppet-strings/markdown/puppet_classes.rb
@@ -3,37 +3,16 @@
 require_relative 'puppet_class'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about classes
   module PuppetClasses
+    def self.name
+      'Classes'
+    end
 
     # @return [Array] list of classes
-    def self.in_classes
+    def self.items
       arr = YARD::Registry.all(:puppet_class).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::PuppetClass.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_classes.nil?
-        in_classes.find { |klass| klass.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_classes.length > 0 ? "## Classes\n\n" : ""
-      in_classes.each do |klass|
-        final += klass.render unless klass.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Classes"]
-
-      in_classes.each do |klass|
-        final.push(klass.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/puppet_plans.rb b/lib/puppet-strings/markdown/puppet_plans.rb
index 7586eab01..e15286b49 100644
--- a/lib/puppet-strings/markdown/puppet_plans.rb
+++ b/lib/puppet-strings/markdown/puppet_plans.rb
@@ -3,37 +3,16 @@
 require_relative 'puppet_plan'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about plans
   module PuppetPlans
+    def self.name
+      'Plans'
+    end
 
     # @return [Array] list of classes
-    def self.in_plans
+    def self.items
       arr = YARD::Registry.all(:puppet_plan).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::PuppetPlan.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_plans.nil?
-        in_plans.find { |plan| plan.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_plans.length > 0 ? "## Plans\n\n" : ""
-      in_plans.each do |plan|
-        final += plan.render unless plan.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Plans"]
-
-      in_plans.each do |plan|
-        final.push(plan.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/puppet_tasks.rb b/lib/puppet-strings/markdown/puppet_tasks.rb
index 9d394f3fc..1c5437bd1 100644
--- a/lib/puppet-strings/markdown/puppet_tasks.rb
+++ b/lib/puppet-strings/markdown/puppet_tasks.rb
@@ -3,34 +3,16 @@
 require_relative 'puppet_task'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about tasks
   module PuppetTasks
+    def self.name
+      'Tasks'
+    end
 
     # @return [Array] list of classes
-    def self.in_tasks
+    def self.items
       arr = YARD::Registry.all(:puppet_task).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::PuppetTask.new(a) }
     end
-
-    def self.contains_private?
-      false
-    end
-
-    def self.render
-      final = in_tasks.length > 0 ? "## Tasks\n\n" : ""
-      in_tasks.each do |task|
-        final += task.render unless task.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Tasks"]
-
-      in_tasks.each do |task|
-        final.push(task.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/resource_types.rb b/lib/puppet-strings/markdown/resource_types.rb
index 01abe1c65..2af9d7162 100644
--- a/lib/puppet-strings/markdown/resource_types.rb
+++ b/lib/puppet-strings/markdown/resource_types.rb
@@ -3,37 +3,16 @@
 require_relative 'resource_type'
 
 module PuppetStrings::Markdown
+  # Adapter to get information about resource types
   module ResourceTypes
+    def self.name
+      'Resource types'
+    end
 
     # @return [Array] list of resource types
-    def self.in_rtypes
+    def self.items
       arr = YARD::Registry.all(:puppet_type).sort_by!(&:name).map!(&:to_hash)
       arr.map! { |a| PuppetStrings::Markdown::ResourceType.new(a) }
     end
-
-    def self.contains_private?
-      result = false
-      unless in_rtypes.nil?
-        in_rtypes.find { |type| type.private? }.nil? ? false : true
-      end
-    end
-
-    def self.render
-      final = in_rtypes.length > 0 ? "## Resource types\n\n" : ""
-      in_rtypes.each do |type|
-        final += type.render unless type.private?
-      end
-      final
-    end
-
-    def self.toc_info
-      final = ["Resource types"]
-
-      in_rtypes.each do |type|
-        final.push(type.toc_info)
-      end
-
-      final
-    end
   end
 end
diff --git a/lib/puppet-strings/markdown/table_of_contents.rb b/lib/puppet-strings/markdown/table_of_contents.rb
deleted file mode 100644
index 5f7e15382..000000000
--- a/lib/puppet-strings/markdown/table_of_contents.rb
+++ /dev/null
@@ -1,26 +0,0 @@
-# frozen_string_literal: true
-
-module PuppetStrings::Markdown
-  module TableOfContents
-    def self.render
-      final = "## Table of Contents\n\n"
-
-      [PuppetStrings::Markdown::PuppetClasses,
-      PuppetStrings::Markdown::DefinedTypes,
-      PuppetStrings::Markdown::ResourceTypes,
-      PuppetStrings::Markdown::Functions,
-      PuppetStrings::Markdown::DataTypes,
-      PuppetStrings::Markdown::PuppetTasks,
-      PuppetStrings::Markdown::PuppetPlans].each do |r|
-        toc = r.toc_info
-        group_name = toc.shift
-        group = toc
-        priv = r.contains_private?
-
-        template = File.join(File.dirname(__FILE__), 'templates/table_of_contents.erb')
-        final += PuppetStrings::Markdown.erb(template).result(binding)
-      end
-      final
-    end
-  end
-end
diff --git a/lib/puppet-strings/markdown/templates/table_of_contents.erb b/lib/puppet-strings/markdown/templates/table_of_contents.erb
index f20582a3a..72d91e207 100644
--- a/lib/puppet-strings/markdown/templates/table_of_contents.erb
+++ b/lib/puppet-strings/markdown/templates/table_of_contents.erb
@@ -1,26 +1,26 @@
-<% if group.length > 0 -%>
+<% unless items.empty? -%>
 ### <%= group_name %>
+<% if has_public -%>
+<% if has_private # only display public heading if we have both -%>
 
-<% if priv -%>
 #### Public <%= group_name %>
+<% end -%>
 
-<% group.each do |item| -%>
+<% items.each do |item| -%>
 <% unless item[:private] -%>
 * [`<%= item[:name] %>`](#<%= item[:link] %>)<% unless item[:desc].nil? || item[:desc].empty? %>: <%= item[:desc] %><% end %>
 <% end -%>
 <% end -%>
+<% end -%>
+<% if has_private -%>
 
 #### Private <%= group_name %>
 
-<% group.each do |item| -%>
+<% items.each do |item| -%>
 <% if item[:private] -%>
 * `<%= item[:name] %>`<% unless item[:desc].nil? || item[:desc].empty? %>: <%= item[:desc] %><% end %>
 <% end -%>
 <% end -%>
-<% else -%>
-<% group.each do |item| -%>
-* [`<%= item[:name] %>`](#<%= item[:link] %>)<% unless item[:desc].nil? || item[:desc].empty? %>: <%= item[:desc] %><% end %>
-<% end -%>
 <% end -%>
 
 <% end -%>
diff --git a/spec/unit/puppet-strings/markdown_spec.rb b/spec/unit/puppet-strings/markdown_spec.rb
index 96246e58b..e9fe6ca34 100644
--- a/spec/unit/puppet-strings/markdown_spec.rb
+++ b/spec/unit/puppet-strings/markdown_spec.rb
@@ -2,56 +2,55 @@
 
 require 'spec_helper'
 require 'puppet-strings/markdown'
-require 'puppet-strings/markdown/table_of_contents'
 require 'tempfile'
 
 describe PuppetStrings::Markdown do
-  let(:fixture_path) do
-    File.expand_path("../../fixtures", __dir__)
-  end
+  describe 'rendering fixtures' do
+    let(:fixture_path) do
+      File.expand_path("../../fixtures", __dir__)
+    end
 
-  def fixture_content(fixture)
-    @fixtures ||= {}
-    @fixtures[fixture] ||= File.read(File.join(fixture_path, fixture))
-  end
+    def fixture_content(fixture)
+      @fixtures ||= {}
+      @fixtures[fixture] ||= File.read(File.join(fixture_path, fixture))
+    end
 
-  def parse_shared_content
-    # Populate the YARD registry with both Puppet and Ruby source
-    YARD::Parser::SourceParser.parse_string(fixture_content("puppet/class.pp"), :puppet)
-    YARD::Parser::SourceParser.parse_string(fixture_content("puppet/function.pp"), :puppet)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func4x.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func4x_1.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func3x.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func3x.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/provider.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/resource_type.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/resource_api.rb"), :ruby)
-
-    # task metadata derives the task name from the filename, so we have to parse
-    # directly from the filesystem to correctly pick up the name
-    YARD::Parser::SourceParser.parse(File.join(fixture_path, "json/backup.json"))
-  end
+    def parse_shared_content
+      # Populate the YARD registry with both Puppet and Ruby source
+      YARD::Parser::SourceParser.parse_string(fixture_content("puppet/class.pp"), :puppet)
+      YARD::Parser::SourceParser.parse_string(fixture_content("puppet/function.pp"), :puppet)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func4x.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func4x_1.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func3x.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/func3x.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/provider.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/resource_type.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/resource_api.rb"), :ruby)
+
+      # task metadata derives the task name from the filename, so we have to parse
+      # directly from the filesystem to correctly pick up the name
+      YARD::Parser::SourceParser.parse(File.join(fixture_path, "json/backup.json"))
+    end
 
-  def parse_plan_content
-    # the parser behaves differently when parsing puppet files in the the plans directory,
-    # so we have to parse directly from the filesystem to correctly pick up the name
-    YARD::Parser::SourceParser.parse(File.join(fixture_path, "plans/plan.pp"))
-  end
+    def parse_plan_content
+      # the parser behaves differently when parsing puppet files in the the plans directory,
+      # so we have to parse directly from the filesystem to correctly pick up the name
+      YARD::Parser::SourceParser.parse(File.join(fixture_path, "plans/plan.pp"))
+    end
 
-  def parse_data_type_content
-    YARD::Parser::SourceParser.parse_string(fixture_content("ruby/data_type.rb"), :ruby)
-    YARD::Parser::SourceParser.parse_string(fixture_content("puppet/type_alias.pp"), :puppet)
-  end
+    def parse_data_type_content
+      YARD::Parser::SourceParser.parse_string(fixture_content("ruby/data_type.rb"), :ruby)
+      YARD::Parser::SourceParser.parse_string(fixture_content("puppet/type_alias.pp"), :puppet)
+    end
 
-  let(:output) { PuppetStrings::Markdown.generate }
+    let(:output) { PuppetStrings::Markdown.generate }
 
-  RSpec.shared_examples 'markdown lint checker' do |parameter|
-    it 'should not generate markdown lint errors from the rendered markdown' do
-      expect(output).to have_no_markdown_lint_errors
+    RSpec.shared_examples 'markdown lint checker' do |parameter|
+      it 'should not generate markdown lint errors from the rendered markdown' do
+        expect(output).to have_no_markdown_lint_errors
+      end
     end
-  end
 
-  describe 'markdown rendering' do
     before(:each) do
       parse_shared_content
     end
@@ -160,4 +159,105 @@ def parse_data_type_content
       end
     end
   end
+
+  it 'renders only private functions correctly' do
+    expect(YARD::Parser::SourceParser.parse_string(<<~'PUPPET', :puppet).enumerator.length).to eq(1)
+      # @return void
+      # @api private
+      function func_private() {}
+    PUPPET
+
+    expect(described_class.generate).to eq(<<~'MARKDOWN')
+      # Reference
+
+      <!-- DO NOT EDIT: This document was generated by Puppet Strings -->
+
+      ## Table of Contents
+
+      ### Functions
+
+      #### Private Functions
+
+      * `func_private`
+
+    MARKDOWN
+  end
+
+  it 'renders only public functions correctly' do
+    expect(YARD::Parser::SourceParser.parse_string(<<~'PUPPET', :puppet).enumerator.length).to eq(1)
+      # @return void
+      function func_public() {}
+    PUPPET
+
+    expect(described_class.generate).to eq(<<~'MARKDOWN')
+      # Reference
+
+      <!-- DO NOT EDIT: This document was generated by Puppet Strings -->
+
+      ## Table of Contents
+
+      ### Functions
+
+      * [`func_public`](#func_public)
+
+      ## Functions
+
+      ### <a name="func_public"></a>`func_public`
+
+      Type: Puppet Language
+
+      The func_public function.
+
+      #### `func_public()`
+
+      The func_public function.
+
+      Returns: `Any` void
+
+    MARKDOWN
+  end
+
+  it 'renders both public and private functions correctly' do
+    expect(YARD::Parser::SourceParser.parse_string(<<~'PUPPET', :puppet).enumerator.length).to eq(2)
+      # @return void
+      function func_public() {}
+
+      # @return void
+      # @api private
+      function func_private() {}
+    PUPPET
+
+    expect(described_class.generate).to eq(<<~'MARKDOWN')
+      # Reference
+
+      <!-- DO NOT EDIT: This document was generated by Puppet Strings -->
+
+      ## Table of Contents
+
+      ### Functions
+
+      #### Public Functions
+
+      * [`func_public`](#func_public)
+
+      #### Private Functions
+
+      * `func_private`
+
+      ## Functions
+
+      ### <a name="func_public"></a>`func_public`
+
+      Type: Puppet Language
+
+      The func_public function.
+
+      #### `func_public()`
+
+      The func_public function.
+
+      Returns: `Any` void
+
+    MARKDOWN
+  end
 end