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