Optimize null-coalescing operator transpilation to if/else statements

src/astUtils/creators.ts

@@ -203,7 +203,7 @@ export function createIfStatement(options: {
         {
             if: options.if ?? createToken(TokenKind.If),
             then: options.then ?? createToken(TokenKind.Then),
-            else: options.else ?? createToken(TokenKind.Else),
+            else: options.elseBranch ? (options.else ?? createToken(TokenKind.Else)) : undefined,
             endIf: options.endIf ?? createToken(TokenKind.EndIf)
         },
         options.condition,

src/bscPlugin/transpile/BrsFilePreTranspileProcessor.ts

@@ -1,13 +1,13 @@
-import { createAssignmentStatement, createBlock, createDottedSetStatement, createIfStatement, createIndexedSetStatement, createToken } from '../../astUtils/creators';
-import { isAssignmentStatement, isBinaryExpression, isBlock, isBody, isBrsFile, isDottedGetExpression, isDottedSetStatement, isGroupingExpression, isIndexedGetExpression, isIndexedSetStatement, isLiteralExpression, isUnaryExpression, isVariableExpression } from '../../astUtils/reflection';
+import { createAssignmentStatement, createBlock, createDottedSetStatement, createIfStatement, createIndexedSetStatement, createToken, createVariableExpression, createInvalidLiteral } from '../../astUtils/creators';
+import { isAssignmentStatement, isBinaryExpression, isBlock, isBody, isBrsFile, isCallExpression, isCallfuncExpression, isDottedGetExpression, isDottedSetStatement, isGroupingExpression, isIndexedGetExpression, isIndexedSetStatement, isLiteralExpression, isUnaryExpression, isVariableExpression } from '../../astUtils/reflection';
 import { createVisitor, WalkMode } from '../../astUtils/visitors';
 import type { BrsFile } from '../../files/BrsFile';
 import type { BeforeFileTranspileEvent } from '../../interfaces';
 import type { Token } from '../../lexer/Token';
 import { TokenKind } from '../../lexer/TokenKind';
 import type { Expression, Statement } from '../../parser/AstNode';
-import type { TernaryExpression } from '../../parser/Expression';
-import { LiteralExpression } from '../../parser/Expression';
+import type { TernaryExpression, NullCoalescingExpression } from '../../parser/Expression';
+import { BinaryExpression, LiteralExpression } from '../../parser/Expression';
 import { ParseMode } from '../../parser/Parser';
 import type { IfStatement } from '../../parser/Statement';
 import type { Scope } from '../../Scope';
@@ -41,6 +41,9 @@ export class BrsFilePreTranspileProcessor {
         const visitor = createVisitor({
             TernaryExpression: (ternaryExpression) => {
                 this.processTernaryExpression(ternaryExpression, visitor, walkMode);
+            },
+            NullCoalescingExpression: (nullCoalescingExpression) => {
+                this.processNullCoalescingExpression(nullCoalescingExpression, visitor, walkMode);
             }
         });
         this.event.file.ast.walk(visitor, { walkMode: walkMode });
@@ -178,6 +181,84 @@ export class BrsFilePreTranspileProcessor {
         }
     }
 
+    private processNullCoalescingExpression(nullCoalescingExpression: NullCoalescingExpression, visitor: ReturnType<typeof createVisitor>, walkMode: WalkMode) {
+        // Check if this null coalescing expression has complex expressions that require scope protection
+        const consequentInfo = util.getExpressionInfo(nullCoalescingExpression.consequent, this.event.file);
+        const alternateInfo = util.getExpressionInfo(nullCoalescingExpression.alternate, this.event.file);
+
+        let hasComplexExpression = [
+            ...consequentInfo.expressions,
+            ...alternateInfo.expressions
+        ].find(e => isCallExpression(e) || isCallfuncExpression(e) || isDottedGetExpression(e) || isIndexedGetExpression(e));
+
+        // Only optimize if there are no complex expressions
+        if (hasComplexExpression) {
+            return;
+        }
+
+        function getOwnerAndKey(statement: Statement) {
+            const parent = statement.parent;
+            if (isBlock(parent) || isBody(parent)) {
+                let idx = parent.statements.indexOf(statement);
+                if (idx > -1) {
+                    return { owner: parent.statements, key: idx };
+                }
+            }
+        }
+
+        //if the null coalescing expression is part of a simple assignment to a local variable, rewrite it as an `IfStatement`
+        let parent = nullCoalescingExpression.findAncestor(x => !isGroupingExpression(x));
+        let ifStatement: IfStatement;
+
+        // Only support simple AssignmentStatement to local variables
+        if (isAssignmentStatement(parent)) {
+            // For simple assignments like `a = user ?? {}`, use the original logic
+            // Create condition: variableName = invalid
+            const condition = new BinaryExpression(
+                createVariableExpression(parent.name.text),
+                createToken(TokenKind.Equal, '=', nullCoalescingExpression.questionQuestionToken.range),
+                createInvalidLiteral('invalid', nullCoalescingExpression.questionQuestionToken.range)
+            );
+
+            ifStatement = createIfStatement({
+                if: createToken(TokenKind.If, 'if', nullCoalescingExpression.questionQuestionToken.range),
+                condition: condition,
+                then: createToken(TokenKind.Then, 'then', nullCoalescingExpression.questionQuestionToken.range),
+                thenBranch: createBlock({
+                    statements: [
+                        createAssignmentStatement({
+                            name: parent.name,
+                            equals: parent.equals,
+                            value: nullCoalescingExpression.alternate
+                        })
+                    ]
+                }),
+                endIf: createToken(TokenKind.EndIf, 'end if', nullCoalescingExpression.questionQuestionToken.range)
+            });
+
+            // First, we need to create the initial assignment statement
+            const initialAssignment = createAssignmentStatement({
+                name: parent.name,
+                equals: parent.equals,
+                value: nullCoalescingExpression.consequent
+            });
+
+            // Replace the parent with a sequence: first the initial assignment, then the if statement
+            let { owner, key } = getOwnerAndKey(parent as Statement) ?? {};
+            if (owner && key !== undefined) {
+                // Replace with initial assignment first
+                this.event.editor.setProperty(owner, key, initialAssignment);
+                // Insert the if statement after
+                this.event.editor.addToArray(owner, key + 1, ifStatement);
+            }
+        }
+
+        if (ifStatement) {
+            //we've injected an ifStatement, so now we need to trigger a walk to handle any nested null coalescing expressions
+            ifStatement.walk(visitor, { walkMode: walkMode });
+        }
+    }
+
     /**
      * Given a string optionally separated by dots, find an enum related to it.
      * For example, all of these would return the enum: `SomeNamespace.SomeEnum.SomeMember`, SomeEnum.SomeMember, `SomeEnum`

src/parser/tests/expression/NullCoalescenceExpression.spec.ts

@@ -197,7 +197,10 @@ describe('NullCoalescingExpression', () => {
                 end sub
             `, `
                 sub main()
-                    a = rokucommunity_bslib_coalesce(user, false)
+                    a = user
+                    if a = invalid then
+                        a = false
+                    end if
                 end sub
             `);
         });
@@ -209,9 +212,12 @@ describe('NullCoalescingExpression', () => {
                 end sub
             `, `
                 sub main()
-                    a = bslib_coalesce(user, {
-                        "id": "default"
-                    })
+                    a = user
+                    if a = invalid then
+                        a = {
+                            "id": "default"
+                        }
+                    end if
                 end sub
             `);
         });
@@ -374,5 +380,99 @@ describe('NullCoalescingExpression', () => {
                 end sub
             `);
         });
+
+        it('transpiles null coalescing in RHS of AssignmentStatement to if statement', () => {
+            testTranspile(`
+                sub main()
+                    a = user ?? {}
+                end sub
+            `, `
+                sub main()
+                    a = user
+                    if a = invalid then
+                        a = {}
+                    end if
+                end sub
+            `);
+        });
+
+        it('uses bslib_coalesce for compound assignment operators', () => {
+            testTranspile(`
+                sub main()
+                    a += user ?? 0
+                end sub
+            `, `
+                sub main()
+                    a += bslib_coalesce(user, 0)
+                end sub
+            `);
+        });
+
+        it('supports nested null coalescing in assignment', () => {
+            testTranspile(`
+                sub main()
+                    result = user ?? (fallback ?? {})
+                end sub
+            `, `
+                sub main()
+                    result = user
+                    if result = invalid then
+                        result = fallback
+                        if result = invalid then
+                            result = {}
+                        end if
+                    end if
+                end sub
+            `);
+        });
+
+        it('uses scope-captured functions for DottedSet expressions', () => {
+            testTranspile(`
+                sub main()
+                    m.a = user ?? {}
+                end sub
+            `, `
+                sub main()
+                    m.a = bslib_coalesce(user, {})
+                end sub
+            `);
+        });
+
+        it('uses scope-captured functions for IndexedSet expressions', () => {
+            testTranspile(`
+                sub main()
+                    m["a"] = user ?? {}
+                end sub
+            `, `
+                sub main()
+                    m["a"] = bslib_coalesce(user, {})
+                end sub
+            `);
+        });
+
+        it('uses scope-captured functions for complex expressions', () => {
+            testTranspile(`
+                sub main()
+                    zombie = {}
+                    result = [
+                        zombie.getName() ?? "zombie"
+                    ]
+                end sub
+            `, `
+                sub main()
+                    zombie = {}
+                    result = [
+                        (function(zombie)
+                                __bsConsequent = zombie.getName()
+                                if __bsConsequent <> invalid then
+                                    return __bsConsequent
+                                else
+                                    return "zombie"
+                                end if
+                            end function)(zombie)
+                    ]
+                end sub
+            `);
+        });
     });
 });