AST editing

docs/plugins.md

@@ -228,32 +228,44 @@ export default function () {
 };
 ```
 
-### Example AST modifier plugin
+## Modifying code
+Sometimes plugins will want to modify code before the project is transpiled. While you can technically edit the AST directly at any point in the file's lifecycle, this is not recommended as those changes will remain changed as long as that file exists in memory and could cause issues with file validation if the plugin is used in a language-server context (i.e. inside vscode).
 
-AST modification can be done after parsing (`afterFileParsed`), but it is recommended to modify the AST only before transpilation (`beforeFileTranspile`), otherwise it could cause problems if the plugin is used in a language-server context.
+Instead, we provide an instace of an `AstEditor` class in the `beforeFileTranspile` event that allows you to modify AST before the file is transpiled, and then those modifications are undone `afterFileTranspile`.
+
+For example, consider the following brightscript code:
+```brightscript
+sub main()
+    print "hello <FIRST_NAME>"
+end sub
+```
+
+Here's the plugin:
 
 ```typescript
-// removePrint.ts
-import { CompilerPlugin, Program, TranspileObj } from 'brighterscript';
-import { EmptyStatement } from 'brighterscript/dist/parser';
-import { isBrsFile, createStatementEditor, editStatements } from 'brighterscript/dist/parser/ASTUtils';
+import { CompilerPlugin, BeforeFileTranspileEvent, isBrsFile, WalkMode, createVisitor, TokenKind } from './';
 
 // plugin factory
 export default function () {
     return {
         name: 'removePrint',
         // transform AST before transpilation
-        beforeFileTranspile: (entry: TranspileObj) => {
-            if (isBrsFile(entry.file)) {
-                // visit functions bodies and replace `PrintStatement` nodes with `EmptyStatement`
-                entry.file.parser.functionExpressions.forEach((fun) => {
-                    const visitor = createStatementEditor({
-                        PrintStatement: (statement) => new EmptyStatement()
-                    });
-                    editStatements(fun.body, visitor);
+        beforeFileTranspile: (event: BeforeFileTranspileEvent) => {
+            if (isBrsFile(event.file)) {
+                event.file.ast.walk(createVisitor({
+                    LiteralExpression: (literal) => {
+                        //replace every occurance of <FIRST_NAME> in strings with "world"
+                        if (literal.token.kind === TokenKind.StringLiteral && literal.token.text.includes('<FIRST_NAME>')) {
+                            event.astEditor.setProperty(literal.token, 'text', literal.token.text.replace('<FIRST_NAME>', 'world'));
+                        }
+                    }
+                }), {
+                    walkMode: WalkMode.visitExpressionsRecursive
                 });
             }
         }
     } as CompilerPlugin;
 };
 ```
+
+This plugin will search through every LiteralExpression in the entire project, and every time we find a string literal, we will replace `<FIRST_NAME>` with `world`. This is done with the `event.astEditor` object. `astEditor` allows you to apply edits to the AST, and then the brighterscript compiler will `undo` those edits once the file has been transpiled.

src/Program.spec.ts

@@ -16,8 +16,9 @@ import { EmptyStatement } from './parser/Statement';
 import { expectZeroDiagnostics, trim, trimMap } from './testHelpers.spec';
 import { doesNotThrow } from 'assert';
 import { Logger } from './Logger';
-import { createToken } from './astUtils';
+import { createToken, createVisitor, isBrsFile, WalkMode } from './astUtils';
 import { TokenKind } from './lexer';
+import type { LiteralExpression } from './parser/Expression';
 
 let sinon = sinonImport.createSandbox();
 let tmpPath = s`${process.cwd()}/.tmp`;
@@ -1640,6 +1641,45 @@ describe('Program', () => {
     });
 
     describe('transpile', () => {
+
+        it('handles AstEditor flow properly', async () => {
+            program.addOrReplaceFile('source/main.bs', `
+                sub main()
+                    print "hello world"
+                end sub
+            `);
+            let literalExpression: LiteralExpression;
+            //replace all strings with "goodbye world"
+            program.plugins.add({
+                name: 'TestPlugin',
+                beforeFileTranspile: (event) => {
+                    if (isBrsFile(event.file)) {
+                        event.file.ast.walk(createVisitor({
+                            LiteralExpression: (literal) => {
+                                literalExpression = literal;
+                                event.astEditor.setProperty(literal.token, 'text', '"goodbye world"');
+                            }
+                        }), {
+                            walkMode: WalkMode.visitExpressionsRecursive
+                        });
+                    }
+                }
+            });
+            //transpile the file
+            await program.transpile([], stagingFolderPath);
+            //our changes should be there
+            expect(
+                fsExtra.readFileSync(`${stagingFolderPath}/source/main.brs`).toString()
+            ).to.eql(trim`
+                sub main()
+                    print "goodbye world"
+                end sub`
+            );
+
+            //our literalExpression should have been restored to its original value
+            expect(literalExpression.token.text).to.eql('"hello world"');
+        });
+
         it('copies bslib.brs when no ropm version was found', async () => {
             await program.transpile([], stagingFolderPath);
             expect(fsExtra.pathExistsSync(`${stagingFolderPath}/source/bslib.brs`)).to.be.true;

src/Program.ts

@@ -25,6 +25,7 @@ import type { FunctionStatement, Statement } from './parser/Statement';
 import { ParseMode } from './parser';
 import { TokenKind } from './lexer';
 import { BscPlugin } from './bscPlugin/BscPlugin';
+import { AstEditor } from './astUtils/AstEditor';
 const startOfSourcePkgPath = `source${path.sep}`;
 const bslibNonAliasedRokuModulesPkgPath = s`source/roku_modules/rokucommunity_bslib/bslib.brs`;
 const bslibAliasedRokuModulesPkgPath = s`source/roku_modules/bslib/bslib.brs`;
@@ -1188,7 +1189,8 @@ export class Program {
             outputPath = s`${stagingFolderPath}/${outputPath}`;
             return {
                 file: file,
-                outputPath: outputPath
+                outputPath: outputPath,
+                astEditor: new AstEditor()
             };
         });
 
@@ -1199,6 +1201,7 @@ export class Program {
             if (isBrsFile(entry.file) && entry.file.isTypedef) {
                 return;
             }
+
             this.plugins.emit('beforeFileTranspile', entry);
             const { file, outputPath } = entry;
             const result = file.transpile();
@@ -1222,6 +1225,9 @@ export class Program {
             }
 
             this.plugins.emit('afterFileTranspile', entry);
+
+            //undo all astEditor edits that may have been applied to this file.
+            entry.astEditor.undoAll();
         });
 
         //if there's no bslib file already loaded into the program, copy it to the staging directory

src/astUtils/AstEditor.spec.ts

@@ -0,0 +1,151 @@
+import { expect } from 'chai';
+import { AstEditor } from './AstEditor';
+
+describe('AstEditor', () => {
+    let changer: AstEditor;
+    let obj: ReturnType<typeof getTestObject>;
+
+    beforeEach(() => {
+        changer = new AstEditor();
+        obj = getTestObject();
+    });
+
+    function getTestObject() {
+        return {
+            name: 'parent',
+            hobbies: ['gaming', 'reading', 'cycling'],
+            children: [{
+                name: 'oldest',
+                age: 15
+            }, {
+                name: 'middle',
+                age: 10
+            }, {
+                name: 'youngest',
+                age: 5
+            }],
+            jobs: [{
+                title: 'plumber',
+                annualSalary: 50000
+            }, {
+                title: 'carpenter',
+                annualSalary: 75000
+            }]
+        };
+    }
+
+    it('applies single property change', () => {
+        expect(obj.name).to.eql('parent');
+
+        changer.setProperty(obj, 'name', 'jack');
+        expect(obj.name).to.eql('jack');
+
+        changer.undoAll();
+        expect(obj.name).to.eql('parent');
+    });
+
+    it('inserts at beginning of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.addToArray(obj.hobbies, 0, 'climbing');
+        expect(obj.hobbies).to.eql(['climbing', 'gaming', 'reading', 'cycling']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('inserts at middle of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.addToArray(obj.hobbies, 1, 'climbing');
+        expect(obj.hobbies).to.eql(['gaming', 'climbing', 'reading', 'cycling']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('inserts at end of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.addToArray(obj.hobbies, 3, 'climbing');
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling', 'climbing']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('removes at beginning of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.removeFromArray(obj.hobbies, 0);
+        expect(obj.hobbies).to.eql(['reading', 'cycling']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('removes at middle of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.removeFromArray(obj.hobbies, 1);
+        expect(obj.hobbies).to.eql(['gaming', 'cycling']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('removes at middle of array', () => {
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+
+        changer.removeFromArray(obj.hobbies, 2);
+        expect(obj.hobbies).to.eql(['gaming', 'reading']);
+
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('restores array after being removed', () => {
+        changer.removeFromArray(obj.hobbies, 0);
+        changer.setProperty(obj, 'hobbies', undefined);
+        expect(obj.hobbies).to.be.undefined;
+        changer.undoAll();
+        expect(obj.hobbies).to.eql(['gaming', 'reading', 'cycling']);
+    });
+
+    it('works for many changes', () => {
+        expect(obj).to.eql(getTestObject());
+        changer.setProperty(obj, 'name', 'bob');
+        changer.setProperty(obj.children[0], 'name', 'jimmy');
+        changer.addToArray(obj.children, obj.children.length, { name: 'sally', age: 1 });
+        changer.removeFromArray(obj.jobs, 1);
+        changer.removeFromArray(obj.hobbies, 0);
+        changer.removeFromArray(obj.hobbies, 0);
+        changer.removeFromArray(obj.hobbies, 0);
+        changer.setProperty(obj, 'hobbies', undefined);
+
+        expect(obj).to.eql({
+            name: 'bob',
+            hobbies: undefined,
+            children: [{
+                name: 'jimmy',
+                age: 15
+            }, {
+                name: 'middle',
+                age: 10
+            }, {
+                name: 'youngest',
+                age: 5
+            }, {
+                name: 'sally',
+                age: 1
+            }],
+            jobs: [{
+                title: 'plumber',
+                annualSalary: 50000
+            }]
+        });
+
+        changer.undoAll();
+        expect(obj).to.eql(getTestObject());
+    });
+});