Fix shortcode handling in ToC

Fixes #1010

2 files changed, 29 insertions, 13 deletions

diff --git a/hugolib/handler_page.go b/hugolib/handler_page.go
index ed254dcdd..31415f78b 100644
--- a/hugolib/handler_page.go
+++ b/hugolib/handler_page.go
@@ -60,13 +60,14 @@ func (h markdownHandler) PageConvert(p *Page, t tpl.Template) HandledResult {
 	tmpContent, tmpTableOfContents := helpers.ExtractTOC(p.renderContent(helpers.RemoveSummaryDivider(p.rawContent)))
 
 	if len(p.contentShortCodes) > 0 {
-		tmpContentWithTokensReplaced, err := replaceShortcodeTokens(tmpContent, shortcodePlaceholderPrefix, true, p.contentShortCodes)
-
+		replaced, err := replaceShortcodeTokensInsources(shortcodePlaceholderPrefix, true, p.contentShortCodes,
+			tmpContent, tmpTableOfContents)
 		if err != nil {
-			jww.FATAL.Printf("Fail to replace short code tokens in %s:\n%s", p.BaseFileName(), err.Error())
+			jww.FATAL.Printf("Fail to replace shortcode tokens in %s:\n%s", p.BaseFileName(), err.Error())
 			return HandledResult{err: err}
 		}
-		tmpContent = tmpContentWithTokensReplaced
+		tmpContent = replaced[0]
+		tmpTableOfContents = replaced[1]
 	}
 
 	p.Content = helpers.BytesToHTML(tmpContent)
@@ -99,14 +100,14 @@ func (h asciidocHandler) PageConvert(p *Page, t tpl.Template) HandledResult {
 	tmpContent, tmpTableOfContents := helpers.ExtractTOC(p.renderContent(helpers.RemoveSummaryDivider(p.rawContent)))
 
 	if len(p.contentShortCodes) > 0 {
-		tmpContentWithTokensReplaced, err := replaceShortcodeTokens(tmpContent, shortcodePlaceholderPrefix, true, p.contentShortCodes)
-
+		replaced, err := replaceShortcodeTokensInsources(shortcodePlaceholderPrefix, true, p.contentShortCodes,
+			tmpContent, tmpTableOfContents)
 		if err != nil {
-			jww.FATAL.Printf("Fail to replace short code tokens in %s:\n%s", p.BaseFileName(), err.Error())
+			jww.FATAL.Printf("Fail to replace shortcode tokens in %s:\n%s", p.BaseFileName(), err.Error())
 			return HandledResult{err: err}
-		} else {
-			tmpContent = tmpContentWithTokensReplaced
 		}
+		tmpContent = replaced[0]
+		tmpTableOfContents = replaced[1]
 	}
 
 	p.Content = helpers.BytesToHTML(tmpContent)
@@ -127,13 +128,14 @@ func (h rstHandler) PageConvert(p *Page, t tpl.Template) HandledResult {
 	tmpContent, tmpTableOfContents := helpers.ExtractTOC(p.renderContent(helpers.RemoveSummaryDivider(p.rawContent)))
 
 	if len(p.contentShortCodes) > 0 {
-		tmpContentWithTokensReplaced, err := replaceShortcodeTokens(tmpContent, shortcodePlaceholderPrefix, true, p.contentShortCodes)
-
+		replaced, err := replaceShortcodeTokensInsources(shortcodePlaceholderPrefix, true, p.contentShortCodes,
+			tmpContent, tmpTableOfContents)
 		if err != nil {
-			jww.FATAL.Printf("Fail to replace short code tokens in %s:\n%s", p.BaseFileName(), err.Error())
+			jww.FATAL.Printf("Fail to replace shortcode tokens in %s:\n%s", p.BaseFileName(), err.Error())
 			return HandledResult{err: err}
 		}
-		tmpContent = tmpContentWithTokensReplaced
+		tmpContent = replaced[0]
+		tmpTableOfContents = replaced[1]
 	}
 
 	p.Content = helpers.BytesToHTML(tmpContent)
diff --git a/hugolib/shortcode.go b/hugolib/shortcode.go
index f1db1dc19..c4eec974d 100644
--- a/hugolib/shortcode.go
+++ b/hugolib/shortcode.go
@@ -434,6 +434,20 @@ Loop:
 
 }
 
+// replaceShortcodeTokensInsources calls replaceShortcodeTokens for every source given.
+func replaceShortcodeTokensInsources(prefix string, wrapped bool, replacements map[string]string, sources ...[]byte) (b [][]byte, err error) {
+	result := make([][]byte, len(sources))
+	for i, s := range sources {
+		b, err := replaceShortcodeTokens(s, prefix, wrapped, replacements)
+
+		if err != nil {
+			return nil, err
+		}
+		result[i] = b
+	}
+	return result, nil
+}
+
 // Replace prefixed shortcode tokens (HUGOSHORTCODE-1, HUGOSHORTCODE-2) with the real content.
 // wrapped = true means that the token has been wrapped in {@{@/@}@}
 func replaceShortcodeTokens(source []byte, prefix string, wrapped bool, replacements map[string]string) (b []byte, err error) {