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