001 /*
002 * Copyright 2010-2015 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.google.dart.compiler.backend.js.ast.*;
020 import com.google.dart.compiler.backend.js.ast.metadata.MetadataPackage;
021 import org.jetbrains.annotations.NotNull;
022 import org.jetbrains.annotations.Nullable;
023 import org.jetbrains.kotlin.builtins.InlineStrategy;
024 import org.jetbrains.kotlin.js.inline.context.*;
025 import org.jetbrains.kotlin.js.inline.exception.InlineRecursionException;
026
027 import java.util.IdentityHashMap;
028 import java.util.Set;
029 import java.util.Stack;
030
031 import static org.jetbrains.kotlin.js.inline.FunctionInlineMutator.getInlineableCallReplacement;
032 import static org.jetbrains.kotlin.js.inline.clean.CleanPackage.removeUnusedFunctionDefinitions;
033 import static org.jetbrains.kotlin.js.inline.clean.CleanPackage.removeUnusedLocalFunctionDeclarations;
034 import static org.jetbrains.kotlin.js.inline.util.UtilPackage.*;
035 import static org.jetbrains.kotlin.js.translate.utils.JsAstUtils.flattenStatement;
036
037 public class JsInliner extends JsVisitorWithContextImpl {
038
039 private final IdentityHashMap<JsName, JsFunction> functions;
040 private final Stack<JsInliningContext> inliningContexts = new Stack<JsInliningContext>();
041 private final Set<JsFunction> processedFunctions = IdentitySet();
042 private final Set<JsFunction> inProcessFunctions = IdentitySet();
043
044 /**
045 * A statement can contain more, than one inlineable sub-expressions.
046 * When inline call is expanded, current statement is shifted forward,
047 * but still has same statement context with same index on stack.
048 *
049 * The shifting is intentional, because there could be function literals,
050 * that need to be inlined, after expansion.
051 *
052 * After shifting following inline expansion in the same statement could be
053 * incorrect, because wrong statement index is used.
054 *
055 * To prevent this, after every shift this flag is set to true,
056 * so that visitor wont go deeper until statement is visited.
057 *
058 * Example:
059 * inline fun f(g: () -> Int): Int { val a = g(); return a }
060 * inline fun Int.abs(): Int = if (this < 0) -this else this
061 *
062 * val g = { 10 }
063 * >> val h = f(g).abs() // last statement context index
064 *
065 * val g = { 10 } // after inline
066 * >> val f$result // statement index was not changed
067 * val a = g()
068 * f$result = a
069 * val h = f$result.abs() // current expression still here; incorrect to inline abs(),
070 * // because statement context on stack point to different statement
071 */
072 private boolean lastStatementWasShifted = false;
073
074 public static JsProgram process(JsProgram program) {
075 IdentityHashMap<JsName, JsFunction> functions = collectNamedFunctions(program);
076 JsInliner inliner = new JsInliner(functions);
077 inliner.accept(program);
078 removeUnusedFunctionDefinitions(program, functions);
079 return program;
080 }
081
082 JsInliner(IdentityHashMap<JsName, JsFunction> functions) {
083 this.functions = functions;
084 }
085
086 @Override
087 public boolean visit(JsFunction function, JsContext context) {
088 inliningContexts.push(new JsInliningContext(function));
089
090 if (inProcessFunctions.contains(function)) throw new InlineRecursionException();
091 inProcessFunctions.add(function);
092
093 return super.visit(function, context);
094 }
095
096 @Override
097 public void endVisit(JsFunction function, JsContext context) {
098 super.endVisit(function, context);
099 refreshLabelNames(getInliningContext().newNamingContext(), function);
100
101 removeUnusedLocalFunctionDeclarations(function);
102 processedFunctions.add(function);
103
104 assert inProcessFunctions.contains(function);
105 inProcessFunctions.remove(function);
106
107 inliningContexts.pop();
108 }
109
110 @Override
111 public boolean visit(JsInvocation call, JsContext context) {
112 if (call == null) {
113 return false;
114 }
115
116 if (shouldInline(call) && canInline(call)) {
117 JsFunction definition = getFunctionContext().getFunctionDefinition(call);
118 if (!processedFunctions.contains(definition)) {
119 accept(definition);
120 }
121
122 inline(call, context);
123 }
124
125 return !lastStatementWasShifted;
126 }
127
128 private void inline(@NotNull JsInvocation call, @NotNull JsContext context) {
129 JsInliningContext inliningContext = getInliningContext();
130 FunctionContext functionContext = getFunctionContext();
131 functionContext.declareFunctionConstructorCalls(call.getArguments());
132 InlineableResult inlineableResult = getInlineableCallReplacement(call, inliningContext);
133
134 JsStatement inlineableBody = inlineableResult.getInlineableBody();
135 JsExpression resultExpression = inlineableResult.getResultExpression();
136 StatementContext statementContext = inliningContext.getStatementContext();
137
138 /**
139 * Assumes, that resultExpression == null, when result is not needed.
140 * @see FunctionInlineMutator.isResultNeeded()
141 */
142 if (resultExpression == null) {
143 statementContext.removeCurrentStatement();
144 } else {
145 context.replaceMe(resultExpression);
146 }
147
148 /** @see #lastStatementWasShifted */
149 statementContext.shiftCurrentStatementForward();
150 InsertionPoint<JsStatement> insertionPoint = statementContext.getInsertionPoint();
151 insertionPoint.insertAllAfter(flattenStatement(inlineableBody));
152 }
153
154 /**
155 * Prevents JsInliner from traversing sub-expressions,
156 * when current statement was shifted forward.
157 */
158 @Override
159 protected <T extends JsNode> void doTraverse(T node, JsContext ctx) {
160 if (node instanceof JsStatement) {
161 /** @see #lastStatementWasShifted */
162 lastStatementWasShifted = false;
163 }
164
165 if (!lastStatementWasShifted) {
166 super.doTraverse(node, ctx);
167 }
168 }
169
170 @NotNull
171 private JsInliningContext getInliningContext() {
172 return inliningContexts.peek();
173 }
174
175 @NotNull FunctionContext getFunctionContext() {
176 return getInliningContext().getFunctionContext();
177 }
178
179 private boolean canInline(@NotNull JsInvocation call) {
180 FunctionContext functionContext = getFunctionContext();
181 return functionContext.hasFunctionDefinition(call);
182 }
183
184 private static boolean shouldInline(@NotNull JsInvocation call) {
185 InlineStrategy strategy = MetadataPackage.getInlineStrategy(call);
186 return strategy != null && strategy.isInline();
187 }
188
189
190 private class JsInliningContext implements InliningContext {
191 private final FunctionContext functionContext;
192
193 JsInliningContext(JsFunction function) {
194 functionContext = new FunctionContext(function, this) {
195 @Nullable
196 @Override
197 protected JsFunction lookUpStaticFunction(@Nullable JsName functionName) {
198 return functions.get(functionName);
199 }
200 };
201 }
202
203 @NotNull
204 @Override
205 public NamingContext newNamingContext() {
206 JsScope scope = getFunctionContext().getScope();
207 InsertionPoint<JsStatement> insertionPoint = getStatementContext().getInsertionPoint();
208 return new NamingContext(scope, insertionPoint);
209 }
210
211 @NotNull
212 @Override
213 public StatementContext getStatementContext() {
214 return new StatementContext() {
215 @NotNull
216 @Override
217 public JsContext getCurrentStatementContext() {
218 return getLastStatementLevelContext();
219 }
220
221 @NotNull
222 @Override
223 protected JsStatement getEmptyStatement() {
224 return getFunctionContext().getEmpty();
225 }
226
227 @Override
228 public void shiftCurrentStatementForward() {
229 super.shiftCurrentStatementForward();
230 lastStatementWasShifted = true;
231 }
232 };
233 }
234
235 @NotNull
236 @Override
237 public FunctionContext getFunctionContext() {
238 return functionContext;
239 }
240 }
241 }