From b53e9cb21a765b20b967b0f9f7f9d468535cc8e5 Mon Sep 17 00:00:00 2001
From: odknt <4185305+odknt@users.noreply.github.com>
Date: Wed, 6 Feb 2019 16:24:31 +0900
Subject: [PATCH] Finds and replaces TOC in output of asciidoctor with
 blackfriday compatible TOC

---
 helpers/content.go      | 40 ++++++++++++++++++++++++++++++++++++++--
 helpers/content_test.go | 23 +++++++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/helpers/content.go b/helpers/content.go
index f8479cd1b9a..75cb26d9c9c 100644
--- a/helpers/content.go
+++ b/helpers/content.go
@@ -640,13 +640,49 @@ func getAsciidocContent(ctx *RenderingContext) []byte {
 	}
 
 	jww.INFO.Println("Rendering", ctx.DocumentName, "with", path, "...")
-	args := []string{"--no-header-footer", "--safe"}
+	args := []string{"--no-header-footer", "--safe", "--attribute=toc"}
 	if isAsciidoctor {
 		// asciidoctor-specific arg to show stack traces on errors
 		args = append(args, "--trace")
 	}
 	args = append(args, "-")
-	return externallyRenderContent(ctx, path, args)
+	return replaceAsciidocTOC(externallyRenderContent(ctx, path, args))
+}
+
+// replaceAsciidocTOC replaces html tags for TOC for compatibility with blackfriday rendered TOC.
+func replaceAsciidocTOC(content []byte) []byte {
+	first := []byte(`<div id="toc" class="toc">
+<div id="toctitle">Table of Contents</div>
+<ul class="sectlevel1">`)
+	last := []byte(`</ul>
+</div>`)
+	newFirst := []byte(`<nav>
+<ul>`)
+	newLast := []byte(`</ul>
+</nav>`)
+
+	startOfTOC := bytes.Index(content, first)
+	if startOfTOC < 0 {
+		return content
+	}
+
+	peekEnd := len(content)
+	if peekEnd > 70+startOfTOC+len(first) {
+		peekEnd = 70 + startOfTOC + len(first)
+	}
+
+	correctNav := bytes.Index(content[startOfTOC:peekEnd], []byte(`<li><a href="#`))
+	if correctNav < 0 { // no match found
+		return content
+	}
+	lengthOfTOC := bytes.Index(content[startOfTOC:], last)
+	endOfTOC := startOfTOC + lengthOfTOC + len(last)
+
+	newcontent := make([]byte, 0, len(content))
+	newcontent = append(content[:startOfTOC], newFirst...)
+	newcontent = append(newcontent, content[startOfTOC+len(first):lengthOfTOC]...)
+	newcontent = append(newcontent, newLast...)
+	return append(newcontent, content[endOfTOC:]...)
 }
 
 // HasRst returns whether rst2html is installed on this computer.
diff --git a/helpers/content_test.go b/helpers/content_test.go
index 5297df2de2a..3d123121357 100644
--- a/helpers/content_test.go
+++ b/helpers/content_test.go
@@ -453,6 +453,29 @@ func TestExtractNoTOC(t *testing.T) {
 	}
 }
 
+func TestReplaceAsciidocTOC(t *testing.T) {
+	// rendered by asciidoctor.
+	content := []byte(`<div id="toc" class="toc">
+<div id="toctitle">Table of Contents</div>
+<ul class="sectlevel1">
+<li><a href="#sect1">sect1</a></li>
+</ul>
+</div>`)
+
+	// compatible TOC with blackfriday.
+	expected := []byte(`<nav>
+<ul>
+<li><a href="#sect1">sect1</a></li>
+</ul>
+</nav>`)
+
+	// replaces TOC for extraction by ExtractTOC()
+	actual := replaceAsciidocTOC(content)
+	if !bytes.Equal(actual, expected) {
+		t.Errorf("Actual (%s) did not equal expected (%s) content\n", actual, expected)
+	}
+}
+
 var totalWordsBenchmarkString = strings.Repeat("Hugo Rocks ", 200)
 
 func TestTotalWords(t *testing.T) {