001    /*
002     * Copyright 2010-2016 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.kotlin.js.inline;
018    
019    import com.intellij.psi.PsiElement;
020    import kotlin.jvm.functions.Function1;
021    import org.jetbrains.annotations.NotNull;
022    import org.jetbrains.annotations.Nullable;
023    import org.jetbrains.kotlin.backend.common.CommonCoroutineCodegenUtilKt;
024    import org.jetbrains.kotlin.descriptors.CallableDescriptor;
025    import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
026    import org.jetbrains.kotlin.descriptors.FunctionDescriptor;
027    import org.jetbrains.kotlin.diagnostics.DiagnosticSink;
028    import org.jetbrains.kotlin.diagnostics.Errors;
029    import org.jetbrains.kotlin.js.backend.ast.*;
030    import org.jetbrains.kotlin.js.backend.ast.metadata.MetadataProperties;
031    import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor;
032    import org.jetbrains.kotlin.js.inline.clean.RemoveUnusedFunctionDefinitionsKt;
033    import org.jetbrains.kotlin.js.inline.clean.RemoveUnusedLocalFunctionDeclarationsKt;
034    import org.jetbrains.kotlin.js.inline.context.FunctionContext;
035    import org.jetbrains.kotlin.js.inline.context.InliningContext;
036    import org.jetbrains.kotlin.js.inline.context.NamingContext;
037    import org.jetbrains.kotlin.js.inline.util.CollectUtilsKt;
038    import org.jetbrains.kotlin.js.inline.util.CollectionUtilsKt;
039    import org.jetbrains.kotlin.js.inline.util.NamingUtilsKt;
040    import org.jetbrains.kotlin.js.translate.context.TranslationContext;
041    import org.jetbrains.kotlin.resolve.inline.InlineStrategy;
042    
043    import java.util.*;
044    
045    import static org.jetbrains.kotlin.js.inline.FunctionInlineMutator.getInlineableCallReplacement;
046    import static org.jetbrains.kotlin.js.translate.utils.JsAstUtils.flattenStatement;
047    
048    public class JsInliner extends JsVisitorWithContextImpl {
049    
050        private final Map<JsName, JsFunction> functions;
051        private final Map<String, JsFunction> accessors;
052        private final Stack<JsInliningContext> inliningContexts = new Stack<JsInliningContext>();
053        private final Set<JsFunction> processedFunctions = CollectionUtilsKt.IdentitySet();
054        private final Set<JsFunction> inProcessFunctions = CollectionUtilsKt.IdentitySet();
055        private final FunctionReader functionReader;
056        private final DiagnosticSink trace;
057    
058        // these are needed for error reporting, when inliner detects cycle
059        private final Stack<JsFunction> namedFunctionsStack = new Stack<JsFunction>();
060        private final LinkedList<JsCallInfo> inlineCallInfos = new LinkedList<JsCallInfo>();
061        private final Function1<JsNode, Boolean> canBeExtractedByInliner = new Function1<JsNode, Boolean>() {
062            @Override
063            public Boolean invoke(JsNode node) {
064                if (!(node instanceof JsInvocation)) return false;
065    
066                JsInvocation call = (JsInvocation) node;
067                return hasToBeInlined(call);
068            }
069        };
070    
071        public static JsProgram process(@NotNull TranslationContext context) {
072            JsProgram program = context.program();
073            Map<JsName, JsFunction> functions = CollectUtilsKt.collectNamedFunctions(program);
074            Map<String, JsFunction> accessors = CollectUtilsKt.collectAccessors(program);
075            new DummyAccessorInvocationTransformer().accept(program);
076            JsInliner inliner = new JsInliner(functions, accessors, new FunctionReader(context), context.bindingTrace());
077            inliner.accept(program);
078            RemoveUnusedFunctionDefinitionsKt.removeUnusedFunctionDefinitions(program, functions);
079            return program;
080        }
081    
082        private JsInliner(
083                @NotNull Map<JsName, JsFunction> functions,
084                @NotNull Map<String, JsFunction> accessors,
085                @NotNull FunctionReader functionReader,
086                @NotNull DiagnosticSink trace
087        ) {
088            this.functions = functions;
089            this.accessors = accessors;
090            this.functionReader = functionReader;
091            this.trace = trace;
092        }
093    
094        @Override
095        public boolean visit(@NotNull JsFunction function, @NotNull JsContext context) {
096            inliningContexts.push(new JsInliningContext(function));
097            assert !inProcessFunctions.contains(function): "Inliner has revisited function";
098            inProcessFunctions.add(function);
099    
100            if (functions.containsValue(function)) {
101                namedFunctionsStack.push(function);
102            }
103    
104            return super.visit(function, context);
105        }
106    
107        @Override
108        public void endVisit(@NotNull JsFunction function, @NotNull JsContext context) {
109            super.endVisit(function, context);
110            NamingUtilsKt.refreshLabelNames(function.getBody(), function.getScope());
111    
112            RemoveUnusedLocalFunctionDeclarationsKt.removeUnusedLocalFunctionDeclarations(function);
113            processedFunctions.add(function);
114    
115            new FunctionPostProcessor(function).apply();
116    
117            assert inProcessFunctions.contains(function);
118            inProcessFunctions.remove(function);
119    
120            inliningContexts.pop();
121    
122            if (!namedFunctionsStack.empty() && namedFunctionsStack.peek() == function) {
123                namedFunctionsStack.pop();
124            }
125        }
126    
127        @Override
128        public boolean visit(@NotNull JsInvocation call, @NotNull JsContext context) {
129            if (!hasToBeInlined(call)) return true;
130    
131            JsFunction containingFunction = getCurrentNamedFunction();
132    
133            if (containingFunction != null) {
134                inlineCallInfos.add(new JsCallInfo(call, containingFunction));
135            }
136    
137            JsFunction definition = getFunctionContext().getFunctionDefinition(call);
138    
139            if (inProcessFunctions.contains(definition))  {
140                reportInlineCycle(call, definition);
141            }
142            else if (!processedFunctions.contains(definition)) {
143                accept(definition);
144            }
145    
146            return true;
147        }
148    
149        @Override
150        public void endVisit(@NotNull JsInvocation x, @NotNull JsContext ctx) {
151            if (hasToBeInlined(x)) {
152                inline(x, ctx);
153            }
154    
155            JsCallInfo lastCallInfo = null;
156    
157            if (!inlineCallInfos.isEmpty()) {
158                lastCallInfo = inlineCallInfos.getLast();
159            }
160    
161            if (lastCallInfo != null && lastCallInfo.call == x) {
162                inlineCallInfos.removeLast();
163            }
164        }
165    
166        @Override
167        protected void doAcceptStatementList(List<JsStatement> statements) {
168            // at top level of js ast, contexts stack can be empty,
169            // but there is no inline calls anyway
170            if(!inliningContexts.isEmpty()) {
171                JsScope scope = getFunctionContext().getScope();
172                int i = 0;
173    
174                while (i < statements.size()) {
175                    List<JsStatement> additionalStatements =
176                            ExpressionDecomposer.preserveEvaluationOrder(scope, statements.get(i), canBeExtractedByInliner);
177                    statements.addAll(i, additionalStatements);
178                    i += additionalStatements.size() + 1;
179                }
180            }
181    
182            super.doAcceptStatementList(statements);
183        }
184    
185        private void inline(@NotNull JsInvocation call, @NotNull JsContext context) {
186            DeclarationDescriptor callDescriptor = MetadataProperties.getDescriptor(call);
187            if (isSuspendWithCurrentContinuation(callDescriptor)) {
188                inlineSuspendWithCurrentContinuation(call, context);
189                return;
190            }
191    
192            JsInliningContext inliningContext = getInliningContext();
193            InlineableResult inlineableResult = getInlineableCallReplacement(call, inliningContext);
194    
195            JsStatement inlineableBody = inlineableResult.getInlineableBody();
196            JsExpression resultExpression = inlineableResult.getResultExpression();
197            JsContext<JsStatement> statementContext = inliningContext.getStatementContext();
198            // body of inline function can contain call to lambdas that need to be inlined
199            JsStatement inlineableBodyWithLambdasInlined = accept(inlineableBody);
200            assert inlineableBody == inlineableBodyWithLambdasInlined;
201            statementContext.addPrevious(flattenStatement(inlineableBody));
202    
203            /*
204             * Assumes, that resultExpression == null, when result is not needed.
205             * @see FunctionInlineMutator.isResultNeeded()
206             */
207            if (resultExpression == null) {
208                statementContext.removeMe();
209                return;
210            }
211    
212            resultExpression = accept(resultExpression);
213            MetadataProperties.setSynthetic(resultExpression, true);
214            context.replaceMe(resultExpression);
215        }
216    
217        private static boolean isSuspendWithCurrentContinuation(@Nullable DeclarationDescriptor descriptor) {
218            if (!(descriptor instanceof FunctionDescriptor)) return false;
219            return CommonCoroutineCodegenUtilKt.isBuiltInSuspendCoroutineOrReturn((FunctionDescriptor) descriptor.getOriginal());
220        }
221    
222        private void inlineSuspendWithCurrentContinuation(@NotNull JsInvocation call, @NotNull JsContext context) {
223            JsInliningContext inliningContext = getInliningContext();
224            JsFunction containingFunction = inliningContext.function;
225            JsExpression lambda = call.getArguments().get(0);
226            JsParameter continuationParam = containingFunction.getParameters().get(containingFunction.getParameters().size() - 1);
227    
228            JsInvocation invocation = new JsInvocation(lambda, continuationParam.getName().makeRef());
229            MetadataProperties.setSuspend(invocation, true);
230            context.replaceMe(accept(invocation));
231        }
232    
233        @NotNull
234        private JsInliningContext getInliningContext() {
235            return inliningContexts.peek();
236        }
237    
238        @NotNull
239        private FunctionContext getFunctionContext() {
240            return getInliningContext().getFunctionContext();
241        }
242    
243        @Nullable
244        private JsFunction getCurrentNamedFunction() {
245            if (namedFunctionsStack.empty()) return null;
246            return namedFunctionsStack.peek();
247        }
248    
249        private void reportInlineCycle(@NotNull JsInvocation call, @NotNull JsFunction calledFunction) {
250            MetadataProperties.setInlineStrategy(call, InlineStrategy.NOT_INLINE);
251            Iterator<JsCallInfo> it = inlineCallInfos.descendingIterator();
252    
253            while (it.hasNext()) {
254                JsCallInfo callInfo = it.next();
255                PsiElement psiElement = MetadataProperties.getPsiElement(callInfo.call);
256    
257                CallableDescriptor descriptor = MetadataProperties.getDescriptor(callInfo.call);
258                if (psiElement != null && descriptor != null) {
259                    trace.report(Errors.INLINE_CALL_CYCLE.on(psiElement, descriptor));
260                }
261    
262                if (callInfo.containingFunction == calledFunction) {
263                    break;
264                }
265            }
266        }
267    
268        private boolean hasToBeInlined(@NotNull JsInvocation call) {
269            InlineStrategy strategy = MetadataProperties.getInlineStrategy(call);
270            if (strategy == null || !strategy.isInline()) return false;
271    
272            return getFunctionContext().hasFunctionDefinition(call);
273        }
274    
275        private class JsInliningContext implements InliningContext {
276            private final FunctionContext functionContext;
277    
278            @NotNull
279            public final JsFunction function;
280    
281            JsInliningContext(@NotNull JsFunction function) {
282                this.function = function;
283                functionContext = new FunctionContext(function, functionReader) {
284                    @Nullable
285                    @Override
286                    protected JsFunction lookUpStaticFunction(@Nullable JsName functionName) {
287                        return functions.get(functionName);
288                    }
289    
290                    @Nullable
291                    @Override
292                    protected JsFunction lookUpStaticFunctionByTag(@NotNull String functionTag) {
293                        return accessors.get(functionTag);
294                    }
295                };
296            }
297    
298            @NotNull
299            @Override
300            public NamingContext newNamingContext() {
301                JsScope scope = getFunctionContext().getScope();
302                return new NamingContext(scope, getStatementContext());
303            }
304    
305            @NotNull
306            @Override
307            public JsContext<JsStatement> getStatementContext() {
308                return getLastStatementLevelContext();
309            }
310    
311            @NotNull
312            @Override
313            public FunctionContext getFunctionContext() {
314                return functionContext;
315            }
316        }
317    
318        private static class JsCallInfo {
319            @NotNull
320            public final JsInvocation call;
321    
322            @NotNull
323            public final JsFunction containingFunction;
324    
325            private JsCallInfo(@NotNull JsInvocation call, @NotNull JsFunction function) {
326                this.call = call;
327                containingFunction = function;
328            }
329        }
330    }