From 2c43eaebeccb4cba48f3aa6c54b020e1dd477c19 Mon Sep 17 00:00:00 2001
From: Matt Bierner <matb@microsoft.com>
Date: Fri, 24 Aug 2018 11:01:09 +0200
Subject: [PATCH] Fix markdown document symbols not having full range of
 section

---
 .../src/features/foldingProvider.ts           | 22 +++-------------
 .../src/tableOfContentsProvider.ts            | 25 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/extensions/markdown-language-features/src/features/foldingProvider.ts b/extensions/markdown-language-features/src/features/foldingProvider.ts
index aa4d7a19eb5..42a4c083d76 100644
--- a/extensions/markdown-language-features/src/features/foldingProvider.ts
+++ b/extensions/markdown-language-features/src/features/foldingProvider.ts
@@ -3,11 +3,10 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import * as vscode from 'vscode';
-
-import { MarkdownEngine } from '../markdownEngine';
-import { TableOfContentsProvider, TocEntry } from '../tableOfContentsProvider';
 import { Token } from 'markdown-it';
+import * as vscode from 'vscode';
+import { MarkdownEngine } from '../markdownEngine';
+import { TableOfContentsProvider } from '../tableOfContentsProvider';
 
 const rangeLimit = 5000;
 
@@ -57,19 +56,6 @@ export default class MarkdownFoldingProvider implements vscode.FoldingRangeProvi
 	private async getHeaderFoldingRanges(document: vscode.TextDocument) {
 		const tocProvider = new TableOfContentsProvider(this.engine, document);
 		const toc = await tocProvider.getToc();
-		return toc.map((entry, startIndex) => {
-			const start = entry.line;
-			let end: number | undefined = undefined;
-			for (let i = startIndex + 1; i < toc.length; ++i) {
-				if (toc[i].level <= entry.level) {
-					end = toc[i].line - 1;
-					if (document.lineAt(end).isEmptyOrWhitespace && end >= start + 1) {
-						end = end - 1;
-					}
-					break;
-				}
-			}
-			return new vscode.FoldingRange(start, typeof end === 'number' ? end : document.lineCount - 1);
-		});
+		return toc.map((entry) => new vscode.FoldingRange(entry.line, entry.location.range.end.line));
 	}
 }
diff --git a/extensions/markdown-language-features/src/tableOfContentsProvider.ts b/extensions/markdown-language-features/src/tableOfContentsProvider.ts
index ff2cd01ad24..4991f06726d 100644
--- a/extensions/markdown-language-features/src/tableOfContentsProvider.ts
+++ b/extensions/markdown-language-features/src/tableOfContentsProvider.ts
@@ -17,6 +17,7 @@ export interface TocEntry {
 
 export interface SkinnyTextDocument {
 	readonly uri: vscode.Uri;
+	readonly lineCount: number;
 	getText(): string;
 	lineAt(line: number): vscode.TextLine;
 }
@@ -61,7 +62,29 @@ export class TableOfContentsProvider {
 				location: new vscode.Location(document.uri, line.range)
 			});
 		}
-		return toc;
+
+		// Get full range of section
+		return toc.map((entry, startIndex): TocEntry => {
+			const start = entry.line;
+			let end: number | undefined = undefined;
+			for (let i = startIndex + 1; i < toc.length; ++i) {
+				if (toc[i].level <= entry.level) {
+					end = toc[i].line - 1;
+					if (document.lineAt(end).isEmptyOrWhitespace && end >= start + 1) {
+						end = end - 1;
+					}
+					break;
+				}
+			}
+			const endLine = typeof end === 'number' ? end : document.lineCount - 1;
+			return {
+				...entry,
+				location: new vscode.Location(document.uri,
+					new vscode.Range(
+						entry.location.range.start,
+						new vscode.Position(endLine, document.lineAt(endLine).range.end.character)))
+			};
+		});
 	}
 
 	private static getHeaderLevel(markup: string): number {