1/* 2 * Copyright (c) 2022-2025 Huawei Device Co., Ltd. 3 * Licensed under the Apache License, Version 2.0 (the "License"); 4 * you may not use this file except in compliance with the License. 5 * You may obtain a copy of the License at 6 * 7 * http://www.apache.org/licenses/LICENSE-2.0 8 * 9 * Unless required by applicable law or agreed to in writing, software 10 * distributed under the License is distributed on an "AS IS" BASIS, 11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 * See the License for the specific language governing permissions and 13 * limitations under the License. 14 */ 15 16import * as arkts from '@koalaui/libarkts'; 17import { factory } from './memo-factory'; 18import { AbstractVisitor } from '../common/abstract-visitor'; 19import { isSyntheticReturnStatement, ReturnTypeInfo } from './utils'; 20 21export class ReturnTransformer extends AbstractVisitor { 22 private skipNode?: arkts.ReturnStatement | arkts.BlockStatement; 23 private stableThis: boolean = false; 24 private returnTypeInfo: ReturnTypeInfo | undefined; 25 26 reset() { 27 super.reset(); 28 this.skipNode = undefined; 29 this.stableThis = false; 30 this.returnTypeInfo = undefined; 31 } 32 33 skip(syntheticReturnStatement?: arkts.ReturnStatement | arkts.BlockStatement): ReturnTransformer { 34 this.skipNode = syntheticReturnStatement; 35 return this; 36 } 37 38 rewriteThis(stableThis: boolean): ReturnTransformer { 39 this.stableThis = stableThis; 40 return this; 41 } 42 43 registerReturnTypeInfo(returnTypeInfo: ReturnTypeInfo): ReturnTransformer { 44 this.returnTypeInfo = returnTypeInfo; 45 return this; 46 } 47 48 visitor(beforeChildren: arkts.AstNode): arkts.AstNode { 49 // TODO: temporary checking skip nodes by comparison with expected skip nodes 50 // Should be fixed when update procedure implemented properly 51 if (/* beforeChildren === this.skipNode */ isSyntheticReturnStatement(beforeChildren)) { 52 return beforeChildren; 53 } 54 if (arkts.isScriptFunction(beforeChildren)) { 55 return beforeChildren; 56 } 57 const node = this.visitEachChild(beforeChildren); 58 if (arkts.isReturnStatement(node)) { 59 if (this.stableThis && node.argument && arkts.isThisExpression(node.argument)) { 60 return factory.createReturnThis(); 61 } 62 if (node.argument === undefined) { 63 return arkts.factory.createBlock([ 64 arkts.factory.createExpressionStatement(factory.createRecacheCall()), 65 node, 66 ]); 67 } 68 69 let argument = node.argument; 70 if ( 71 !!this.returnTypeInfo?.node && 72 this.returnTypeInfo.isMemo && 73 arkts.isArrowFunctionExpression(argument) 74 ) { 75 argument = arkts.factory.updateArrowFunction( 76 argument, 77 factory.updateScriptFunctionWithMemoParameters(argument.scriptFunction) 78 ); 79 } 80 81 return arkts.factory.updateReturnStatement(node, factory.createRecacheCall(argument)); 82 } 83 return node; 84 } 85} 86