@@ -30,7 +30,7 @@ export function getPrefixFunctionsTransformer(): ts.TransformerFactory<ts.Source
3030 }
3131
3232 // Add pure function comment to top level functions.
33- if ( topLevelFunctions . indexOf ( node ) !== - 1 ) {
33+ if ( topLevelFunctions . has ( node ) ) {
3434 const newNode = ts . addSyntheticLeadingComment (
3535 node , ts . SyntaxKind . MultiLineCommentTrivia , pureFunctionComment , false ) ;
3636
@@ -49,13 +49,12 @@ export function getPrefixFunctionsTransformer(): ts.TransformerFactory<ts.Source
4949 } ;
5050}
5151
52- export function findTopLevelFunctions ( parentNode : ts . Node ) : ts . Node [ ] {
53- const topLevelFunctions : ts . Node [ ] = [ ] ;
52+ export function findTopLevelFunctions ( parentNode : ts . Node ) : Set < ts . Node > {
53+ const topLevelFunctions = new Set < ts . Node > ( ) ;
5454
5555 function cb ( node : ts . Node ) {
5656 // Stop recursing into this branch if it's a function expression or declaration
57- if ( node . kind === ts . SyntaxKind . FunctionDeclaration
58- || node . kind === ts . SyntaxKind . FunctionExpression ) {
57+ if ( ts . isFunctionDeclaration ( node ) || ts . isFunctionExpression ( node ) ) {
5958 return ;
6059 }
6160
@@ -71,9 +70,24 @@ export function findTopLevelFunctions(parentNode: ts.Node): ts.Node[] {
7170 }
7271
7372 if ( noPureComment ) {
74- if ( innerNode . kind === ts . SyntaxKind . CallExpression
75- || innerNode . kind === ts . SyntaxKind . NewExpression ) {
76- topLevelFunctions . push ( node ) ;
73+ if ( ts . isNewExpression ( innerNode ) ) {
74+ topLevelFunctions . add ( node ) ;
75+ } else if ( ts . isCallExpression ( innerNode ) ) {
76+ let expression : ts . Expression = innerNode . expression ;
77+ while ( expression && ts . isParenthesizedExpression ( expression ) ) {
78+ expression = expression . expression ;
79+ }
80+ if ( expression ) {
81+ if ( ts . isFunctionExpression ( expression ) ) {
82+ // Skip IIFE's with arguments
83+ // This could be improved to check if there are any references to variables
84+ if ( innerNode . arguments . length === 0 ) {
85+ topLevelFunctions . add ( node ) ;
86+ }
87+ } else {
88+ topLevelFunctions . add ( node ) ;
89+ }
90+ }
7791 }
7892 }
7993
0 commit comments