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,138 @@ 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 operator: Token;
+        //operators like `+=` will cause the RHS to be a BinaryExpression due to how the parser handles this. let's do a little magic to detect this situation
+        if (
+            //parent is a binary expression
+            isBinaryExpression(parent) &&
+            isAssignmentStatement(parent.parent) && isVariableExpression(parent.left) && parent.left.name === parent.parent.name
+        ) {
+            //keep the correct operator (i.e. `+=`)
+            operator = parent.operator;
+            //use the outer parent and skip this BinaryExpression
+            parent = parent.parent;
+        }
+        let ifStatement: IfStatement;
+
+        // Only support AssignmentStatement to local variables
+        if (isAssignmentStatement(parent)) {
+            if (operator) {
+                // For compound assignments like `a += user ?? 0`, we need to check the left side value first
+                // Create condition: user <> invalid
+                const condition = new BinaryExpression(
+                    nullCoalescingExpression.consequent,
+                    createToken(TokenKind.LessGreater, '<>', 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: operator,
+                                value: nullCoalescingExpression.consequent
+                            })
+                        ]
+                    }),
+                    else: createToken(TokenKind.Else, 'else', nullCoalescingExpression.questionQuestionToken.range),
+                    elseBranch: createBlock({
+                        statements: [
+                            createAssignmentStatement({
+                                name: parent.name,
+                                equals: operator,
+                                value: nullCoalescingExpression.alternate
+                            })
+                        ]
+                    }),
+                    endIf: createToken(TokenKind.EndIf, 'end if', nullCoalescingExpression.questionQuestionToken.range)
+                });
+
+                // Replace the parent statement with the if statement
+                let { owner, key } = getOwnerAndKey(parent as Statement) ?? {};
+                if (owner && key !== undefined) {
+                    this.event.editor.setProperty(owner, key, ifStatement);
+                }
+            } else {
+                // 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,103 @@ 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('transpiles null coalescing in RHS of incrementor AssignmentStatement to if statement', () => {
+            testTranspile(`
+                sub main()
+                    a += user ?? 0
+                end sub
+            `, `
+                sub main()
+                    if user <> invalid then
+                        a += user
+                    else
+                        a += 0
+                    end if
+                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
+            `);
+        });
     });
 });