Allow updating leading trivia (#1225)

* Refactor getLeadingTrivia to always return an editable array * Added test for verifying leading trivia can be edited: * Extra protection if annotations have to leading trivia
rokucommunity · Jun 14, 2024 · 136a5f6 · 136a5f6
1 parent 0bf5758
commit 136a5f6
Show file tree

Hide file tree

Showing 8 changed files with 155 additions and 58 deletions.
diff --git a/src/bscPlugin/SignatureHelpUtil.ts b/src/bscPlugin/SignatureHelpUtil.ts
@@ -78,7 +78,7 @@ export class SignatureHelpUtil {
         const funcStartPosition = func.range.start;
 
         // Get function comments in reverse order
-        const trivia = func.getLeadingTrivia().reverse();
+        const trivia = util.concatAnnotationLeadingTrivia(func).reverse();
         let functionComments = [] as string[];
 
         for (const currentToken of trivia) {

diff --git a/src/lexer/Lexer.ts b/src/lexer/Lexer.ts
@@ -111,7 +111,7 @@ export class Lexer {
                 ? util.createRange(this.lineBegin, this.columnBegin, this.lineEnd, this.columnEnd + 1)
                 : undefined,
             leadingWhitespace: this.leadingWhitespace,
-            leadingTrivia: this.leadingTrivia
+            leadingTrivia: this.leadingTrivia ?? []
         });
         this.leadingWhitespace = '';
         return this;

diff --git a/src/parser/Expression.ts b/src/parser/Expression.ts
@@ -256,11 +256,11 @@ export class FunctionExpression extends Expression implements TypedefProvider {
     public functionStatement?: FunctionStatement;
 
     public getLeadingTrivia(): Token[] {
-        return this.tokens.functionType?.leadingTrivia ?? [];
+        return this.tokens.functionType?.leadingTrivia;
     }
 
     public getEndTrivia(): Token[] {
-        return this.tokens.endFunctionType?.leadingTrivia ?? [];
+        return this.tokens.endFunctionType?.leadingTrivia;
     }
 
     /**
@@ -515,7 +515,7 @@ export class FunctionParameterExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.name.leadingTrivia ?? [];
+        return this.tokens.name.leadingTrivia;
     }
 }
 
@@ -779,7 +779,7 @@ export class GroupingExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.leftParen?.leadingTrivia ?? [];
+        return this.tokens.leftParen?.leadingTrivia;
     }
 }
 
@@ -833,7 +833,7 @@ export class LiteralExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.value.leadingTrivia ?? [];
+        return this.tokens.value.leadingTrivia;
     }
 }
 
@@ -938,11 +938,11 @@ export class ArrayLiteralExpression extends Expression {
         return new ArrayType(...innerTypes);
     }
     getLeadingTrivia(): Token[] {
-        return this.tokens.open?.leadingTrivia ?? [];
+        return this.tokens.open?.leadingTrivia;
     }
 
     getEndTrivia(): Token[] {
-        return this.tokens.close?.leadingTrivia ?? [];
+        return this.tokens.close?.leadingTrivia;
     }
 }
 
@@ -991,7 +991,7 @@ export class AAMemberExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.key.leadingTrivia ?? [];
+        return this.tokens.key.leadingTrivia;
     }
 
 }
@@ -1091,11 +1091,11 @@ export class AALiteralExpression extends Expression {
     }
 
     public getLeadingTrivia(): Token[] {
-        return this.tokens.open?.leadingTrivia ?? [];
+        return this.tokens.open?.leadingTrivia;
     }
 
     public getEndTrivia(): Token[] {
-        return this.tokens.close?.leadingTrivia ?? [];
+        return this.tokens.close?.leadingTrivia;
     }
 
 }
@@ -1150,7 +1150,7 @@ export class UnaryExpression extends Expression {
     }
 
     public getLeadingTrivia(): Token[] {
-        return this.tokens.operator.leadingTrivia ?? [];
+        return this.tokens.operator.leadingTrivia;
     }
 }
 
@@ -1221,7 +1221,7 @@ export class VariableExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.name.leadingTrivia ?? [];
+        return this.tokens.name.leadingTrivia;
     }
 }
 
@@ -1337,7 +1337,7 @@ export class SourceLiteralExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.value.leadingTrivia ?? [];
+        return this.tokens.value.leadingTrivia;
     }
 }
 
@@ -1407,7 +1407,7 @@ export class NewExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.new.leadingTrivia ?? [];
+        return this.tokens.new.leadingTrivia;
     }
 }
 
@@ -1823,7 +1823,7 @@ export class AnnotationExpression extends Expression {
     }
 
     public getLeadingTrivia(): Token[] {
-        return this.tokens.at?.leadingTrivia ?? [];
+        return this.tokens.at?.leadingTrivia;
     }
 
     transpile(state: BrsTranspileState) {
@@ -2104,7 +2104,7 @@ export class RegexLiteralExpression extends Expression {
     }
 
     getLeadingTrivia(): Token[] {
-        return this.tokens.regexLiteral?.leadingTrivia ?? [];
+        return this.tokens.regexLiteral?.leadingTrivia;
     }
 }
 

diff --git a/src/parser/Parser.spec.ts b/src/parser/Parser.spec.ts
@@ -1660,13 +1660,20 @@ describe('parser', () => {
                 ' hello comment 1
                 ' hello comment 2
                 @annotation
+                ' hello comment 3
+                @otherAnnotation
                 sub sayHello(name as string = "world")
                 end sub
             `, ParseMode.BrighterScript);
             const funcStatements = statements.filter(isFunctionStatement);
             const helloTrivia = funcStatements[0].getLeadingTrivia();
-            expect(helloTrivia.length).to.be.greaterThan(0);
-            expect(helloTrivia.filter(t => t.kind === TokenKind.Comment).length).to.eq(2);
+            expect(helloTrivia.filter(t => t.kind === TokenKind.Comment).length).to.eq(0);
+            const helloAnnotationTrivia = funcStatements[0].annotations[0].getLeadingTrivia();
+            expect(helloAnnotationTrivia.length).to.be.greaterThan(0);
+            expect(helloAnnotationTrivia.filter(t => t.kind === TokenKind.Comment).length).to.eq(2);
+            const otherAnnotationTrivia = funcStatements[0].annotations[1].getLeadingTrivia();
+            expect(otherAnnotationTrivia.length).to.be.greaterThan(0);
+            expect(otherAnnotationTrivia.filter(t => t.kind === TokenKind.Comment).length).to.eq(1);
         });
 
 

diff --git a/src/parser/Statement.spec.ts b/src/parser/Statement.spec.ts
@@ -1,13 +1,17 @@
 import { expect } from '../chai-config.spec';
-import type { NamespaceStatement, ClassStatement } from './Statement';
-import { Body, EmptyStatement } from './Statement';
+import { NamespaceStatement, ClassStatement } from './Statement';
+import { AssignmentStatement, Block, Body, CatchStatement, DottedSetStatement, EmptyStatement, EndStatement, ExitForStatement, ExitWhileStatement, ExpressionStatement, ForEachStatement, ForStatement, FunctionStatement, GotoStatement, IfStatement, ImportStatement, IncrementStatement, IndexedSetStatement, LabelStatement, LibraryStatement, PrintStatement, ReturnStatement, StopStatement, ThrowStatement, TryCatchStatement, WhileStatement } from './Statement';
 import { ParseMode, Parser } from './Parser';
 import { WalkMode } from '../astUtils/visitors';
 import { isClassStatement, isNamespaceStatement } from '../astUtils/reflection';
 import { Program } from '../Program';
 import { trim } from '../testHelpers.spec';
 import type { BrsFile } from '../files/BrsFile';
 import { tempDir } from '../testHelpers.spec';
+import { createStringLiteral, createToken, createVariableExpression } from '../astUtils/creators';
+import { TokenKind } from '../lexer/TokenKind';
+import { FunctionExpression } from './Expression';
+import type { AstNode } from './AstNode';
 
 describe('Statement', () => {
     let program: Program;
@@ -111,4 +115,94 @@ describe('Statement', () => {
         });
     });
 
+
+    describe('all Statements', () => {
+
+        let allStatements: AstNode[] = [];
+
+        beforeEach(() => {
+            const ident = createToken(TokenKind.Identifier, 'a');
+            const expr = createStringLiteral('');
+            const token = createToken(TokenKind.StringLiteral, '');
+            const body = new Body({ statements: [] });
+            const assignment = new AssignmentStatement({ equals: undefined, name: ident, value: expr });
+            const block = new Block({ statements: [] });
+            const expression = new ExpressionStatement({ expression: expr });
+            const exitFor = new ExitForStatement({ exitFor: token });
+            const exitWhile = new ExitWhileStatement({ exitWhile: token });
+            const funs = new FunctionStatement({
+                name: ident,
+                func: new FunctionExpression({
+                    parameters: [],
+                    body: block,
+                    functionType: token,
+                    leftParen: token,
+                    rightParen: token,
+                    endFunctionType: token
+                })
+            });
+            const ifs = new IfStatement({ if: token, condition: expr, thenBranch: block });
+            const increment = new IncrementStatement({ value: expr, operator: token });
+            const print = new PrintStatement({ print: token, expressions: [] });
+            const gotos = new GotoStatement({ goto: token, label: token });
+            const labels = new LabelStatement({ name: ident, colon: token });
+            const returns = new ReturnStatement({ return: token });
+            const ends = new EndStatement({ end: token });
+            const stop = new StopStatement({ stop: token });
+            const fors = new ForStatement({ for: token, counterDeclaration: assignment, to: token, finalValue: expr, body: block, endFor: token, step: token, increment: expr });
+            const foreach = new ForEachStatement({ forEach: token, in: token, endFor: token, item: token, target: expr, body: block });
+            const whiles = new WhileStatement({ while: token, endWhile: token, condition: expr, body: block });
+            const dottedSet = new DottedSetStatement({ obj: expr, name: ident, value: expr });
+            const indexedSet = new IndexedSetStatement({ obj: expr, indexes: [expr], value: expr, openingSquare: token, closingSquare: token });
+            const library = new LibraryStatement({ library: token, filePath: token });
+            const namespace = new NamespaceStatement({ namespace: token, nameExpression: createVariableExpression('a'), body: body, endNamespace: token });
+            const cls = new ClassStatement({ class: token, name: ident, body: [], endClass: token });
+            const imports = new ImportStatement({ import: token, path: token });
+            const catchStmt = new CatchStatement({ catch: token, exceptionVariable: ident, catchBranch: block });
+            const tryCatch = new TryCatchStatement({ try: token, tryBranch: block, catchStatement: catchStmt });
+            const throwSt = new ThrowStatement({ throw: createToken(TokenKind.Throw) });
+
+            allStatements = [
+                expression,
+                exitFor,
+                exitWhile,
+                funs,
+                ifs,
+                increment,
+                print,
+                gotos,
+                labels,
+                returns,
+                ends,
+                stop,
+                fors,
+                foreach,
+                whiles,
+                dottedSet,
+                indexedSet,
+                library,
+                namespace,
+                cls,
+                imports,
+                catchStmt,
+                tryCatch,
+                throwSt
+            ];
+
+        });
+
+        it('has editable leading trivia', () => {
+            for (const stmt of allStatements) {
+                const beforeTrivia = stmt.getLeadingTrivia();
+                expect(beforeTrivia.length, `${stmt.kind} already has leading trivia`).to.eq(0);
+                stmt.getLeadingTrivia().push(createToken(TokenKind.Comment, 'This is an added comment'));
+                const afterComments = stmt.getLeadingTrivia().filter(t => t.kind === TokenKind.Comment);
+                expect(afterComments.length, `${stmt.kind} leading trivia was not edited`).to.eq(1);
+                expect(afterComments[0].text, `${stmt.kind} leading trivia was not edited`).to.eq('This is an added comment');
+                stmt.getLeadingTrivia().pop();
+            }
+        });
+
+    });
+
 });