Implement access specifier folding range collection and add unit tests

Extension/src/LanguageServer/Providers/foldingRangeProvider.ts

@@ -8,11 +8,16 @@ import { ManualPromise } from '../../Utility/Async/manualPromise';
 import { CppFoldingRange, DefaultClient, FoldingRangeKind, GetFoldingRangesParams, GetFoldingRangesRequest, GetFoldingRangesResult } from '../client';
 import { RequestCancelled, ServerCancelled } from '../protocolFilter';
 import { CppSettings } from '../settings';
+import { collectAccessSpecifierFoldingRanges, mergeFoldingRangesWithLimit } from './foldingRangeUtils';
 
 interface FoldingRangeRequestInfo {
     promise: ManualPromise<vscode.FoldingRange[] | undefined> | undefined;
 }
 
+interface FoldingContextWithRangeLimit {
+    rangeLimit?: number;
+}
+
 export class FoldingRangeProvider implements vscode.FoldingRangeProvider {
     private client: DefaultClient;
     public onDidChangeFoldingRangesEvent = new vscode.EventEmitter<void>();
@@ -45,8 +50,12 @@ export class FoldingRangeProvider implements vscode.FoldingRangeProvider {
             promise: undefined
         };
         this.pendingRequests.set(document.uri.toString(), foldingRangeRequestInfo);
+        const rangeLimit: number | undefined = (context as FoldingContextWithRangeLimit).rangeLimit;
 
-        const promise: Promise<vscode.FoldingRange[] | undefined> = this.requestRanges(document.uri.toString(), token);
+        const promise: Promise<vscode.FoldingRange[] | undefined> = this.requestRanges(document.uri.toString(), token).then((ranges: vscode.FoldingRange[] | undefined) => {
+            const accessSpecifierRanges = collectAccessSpecifierFoldingRanges(document.getText());
+            return mergeFoldingRangesWithLimit(ranges, accessSpecifierRanges, rangeLimit) as vscode.FoldingRange[];
+        });
         await promise;
         this.pendingRequests.delete(document.uri.toString());
         if (foldingRangeRequestInfo.promise !== undefined) {

Extension/src/LanguageServer/Providers/foldingRangeUtils.ts

@@ -0,0 +1,196 @@
+export interface FoldingRangeLike {
+    start: number;
+    end: number;
+}
+
+export function mergeFoldingRangesWithLimit(primary: FoldingRangeLike[] | undefined, secondary: FoldingRangeLike[], rangeLimit: number | undefined): FoldingRangeLike[] {
+    const mergedRanges: FoldingRangeLike[] = (primary ?? []).concat(secondary);
+
+    if (rangeLimit === undefined) {
+        return mergedRanges;
+    }
+
+    if (rangeLimit <= 0) {
+        return [];
+    }
+
+    // Keep existing server ranges first, then append access-specifier ranges until the limit.
+    return mergedRanges.slice(0, rangeLimit);
+}
+
+const accessSpecifierPattern: RegExp = /^\s*(public|protected|private)\s*:\s*$/;
+const classDeclarationStartPattern: RegExp = /^\s*(?:template\s*<.*>\s*)?(class|struct|union)\b/;
+
+interface FoldingScanState {
+    inBlockComment: boolean;
+    rawStringDelimiter?: string;
+}
+
+function tryConsumeRawStringStart(line: string, index: number): { consumed: number; rawStringDelimiter: string; } | undefined {
+    const remaining = line.slice(index);
+    const match = /^(?:u8|u|U|L)?R"([^ ()\\\t\r\n]{0,16})\(/.exec(remaining);
+    if (match === null) {
+        return undefined;
+    }
+
+    return {
+        consumed: match[0].length,
+        rawStringDelimiter: match[1]
+    };
+}
+
+function stripLineForFolding(line: string, state: FoldingScanState): { text: string; state: FoldingScanState; } {
+    let result = '';
+    let index = 0;
+    let inString: '"' | '\'' | undefined;
+
+    while (index < line.length) {
+        const character = line[index];
+        const nextCharacter = line[index + 1];
+
+        if (state.rawStringDelimiter !== undefined) {
+            const rawStringTerminator = `)${state.rawStringDelimiter}"`;
+            const rawStringEndIndex = line.indexOf(rawStringTerminator, index);
+            if (rawStringEndIndex < 0) {
+                index = line.length;
+                continue;
+            }
+
+            state.rawStringDelimiter = undefined;
+            index = rawStringEndIndex + rawStringTerminator.length;
+            continue;
+        }
+
+        if (state.inBlockComment) {
+            if (character === '*' && nextCharacter === '/') {
+                state.inBlockComment = false;
+                index += 2;
+                continue;
+            }
+
+            index++;
+            continue;
+        }
+
+        if (inString !== undefined) {
+            if (character === '\\') {
+                index += 2;
+                continue;
+            }
+
+            if (character === inString) {
+                inString = undefined;
+            }
+
+            index++;
+            continue;
+        }
+
+        if (character === '/' && nextCharacter === '/') {
+            break;
+        }
+
+        if (character === '/' && nextCharacter === '*') {
+            state.inBlockComment = true;
+            index += 2;
+            continue;
+        }
+
+        const rawStringStart = tryConsumeRawStringStart(line, index);
+        if (rawStringStart !== undefined) {
+            state.rawStringDelimiter = rawStringStart.rawStringDelimiter;
+            index += rawStringStart.consumed;
+            continue;
+        }
+
+        if (character === '"' || character === '\'') {
+            inString = character;
+            index++;
+            continue;
+        }
+
+        result += character;
+        index++;
+    }
+
+    return { text: result, state };
+}
+
+function countCharacter(line: string, character: string): number {
+    return (line.match(new RegExp(`\\${character}`, 'g')) ?? []).length;
+}
+
+function addFoldingRange(ranges: FoldingRangeLike[], startLine: number, endLine: number): void {
+    if (endLine > startLine) {
+        ranges.push({ start: startLine, end: endLine });
+    }
+}
+
+export function collectAccessSpecifierFoldingRanges(text: string): FoldingRangeLike[] {
+    const ranges: FoldingRangeLike[] = [];
+    const activeSectionsByDepth: Map<number, number> = new Map<number, number>();
+    const classBodyDepths: number[] = [];
+    const lines: string[] = text.split(/\r?\n/);
+
+    const scanState: FoldingScanState = {
+        inBlockComment: false
+    };
+    let braceDepth = 0;
+    let pendingClassDeclaration = false;
+
+    for (let lineIndex = 0; lineIndex < lines.length; lineIndex++) {
+        const strippedLine = stripLineForFolding(lines[lineIndex], scanState);
+
+        const lineText = strippedLine.text;
+        const currentDepth = braceDepth;
+        const currentClassDepth = classBodyDepths[classBodyDepths.length - 1];
+
+        if (currentClassDepth === currentDepth && accessSpecifierPattern.test(lineText)) {
+            const activeSectionStart = activeSectionsByDepth.get(currentDepth);
+            if (activeSectionStart !== undefined) {
+                addFoldingRange(ranges, activeSectionStart, lineIndex - 1);
+            }
+
+            activeSectionsByDepth.set(currentDepth, lineIndex);
+        }
+
+        if (classDeclarationStartPattern.test(lineText)) {
+            pendingClassDeclaration = true;
+        }
+
+        const openingBraces = countCharacter(lineText, '{');
+        const closingBraces = countCharacter(lineText, '}');
+
+        if (pendingClassDeclaration) {
+            if (openingBraces > 0) {
+                const classBodyDepth = currentDepth + openingBraces - closingBraces;
+                if (classBodyDepth > currentDepth) {
+                    classBodyDepths.push(classBodyDepth);
+                }
+                pendingClassDeclaration = false;
+            } else if (lineText.includes(';')) {
+                pendingClassDeclaration = false;
+            }
+        }
+
+        braceDepth = currentDepth + openingBraces - closingBraces;
+
+        while (classBodyDepths.length > 0 && classBodyDepths[classBodyDepths.length - 1] > braceDepth) {
+            const endedClassDepth = classBodyDepths.pop();
+            if (endedClassDepth === undefined) {
+                break;
+            }
+
+            const activeSectionStart = activeSectionsByDepth.get(endedClassDepth);
+            if (activeSectionStart !== undefined) {
+                addFoldingRange(ranges, activeSectionStart, lineIndex - 1);
+                activeSectionsByDepth.delete(endedClassDepth);
+            }
+        }
+    }
+
+    const lastLine = lines.length - 1;
+    activeSectionsByDepth.forEach((startLine: number) => addFoldingRange(ranges, startLine, lastLine));
+
+    return ranges;
+}

Extension/test/unit/foldingRangeProvider.test.ts

@@ -0,0 +1,99 @@
+/* --------------------------------------------------------------------------------------------
+ * Copyright (c) Microsoft Corporation. All Rights Reserved.
+ * See 'LICENSE' in the project root for license information.
+ * -------------------------------------------------------------------------------------------- */
+
+import { deepStrictEqual } from 'assert';
+import { describe, it } from 'mocha';
+import { collectAccessSpecifierFoldingRanges, mergeFoldingRangesWithLimit } from '../../src/LanguageServer/Providers/foldingRangeUtils';
+
+function toRangeTuples(text: string): [number, number][] {
+    return collectAccessSpecifierFoldingRanges(text).map(range => [range.start, range.end]);
+}
+
+describe('Access specifier folding', () => {
+    it('creates fold ranges for public/protected/private sections', () => {
+        const source = [
+            'class A',
+            '{',
+            'public:',
+            '    void foo();',
+            'private:',
+            '    int value;',
+            'protected:',
+            '    void bar();',
+            '};'
+        ].join('\n');
+
+        deepStrictEqual(toRangeTuples(source), [
+            [2, 3],
+            [4, 5],
+            [6, 7]
+        ]);
+    });
+
+    it('respects rangeLimit when merging ranges', () => {
+        const primary = [
+            { start: 0, end: 1 },
+            { start: 2, end: 3 }
+        ];
+        const secondary = [
+            { start: 4, end: 5 },
+            { start: 6, end: 7 }
+        ];
+
+        deepStrictEqual(mergeFoldingRangesWithLimit(primary, secondary, 3), [
+            { start: 0, end: 1 },
+            { start: 2, end: 3 },
+            { start: 4, end: 5 }
+        ]);
+    });
+
+    it('returns all merged ranges when rangeLimit is undefined', () => {
+        const primary = [{ start: 10, end: 20 }];
+        const secondary = [{ start: 30, end: 40 }];
+
+        deepStrictEqual(mergeFoldingRangesWithLimit(primary, secondary, undefined), [
+            { start: 10, end: 20 },
+            { start: 30, end: 40 }
+        ]);
+    });
+
+    it('ignores access-specifier-like lines and braces inside multiline raw strings', () => {
+        const source = [
+            'class A',
+            '{',
+            'public:',
+            '    const char* text = R"raw(',
+            'private:',
+            '}',
+            ')raw";',
+            '    void foo();',
+            'private:',
+            '    int value;',
+            '};'
+        ].join('\n');
+
+        deepStrictEqual(toRangeTuples(source), [
+            [2, 7],
+            [8, 9]
+        ]);
+    });
+
+    it('detects class declarations with same-line template prefix', () => {
+        const source = [
+            'template<typename T> class A',
+            '{',
+            'public:',
+            '    void foo();',
+            'private:',
+            '    int value;',
+            '};'
+        ].join('\n');
+
+        deepStrictEqual(toRangeTuples(source), [
+            [2, 3],
+            [4, 5]
+        ]);
+    });
+});