From 2a03d402333185c8eeb3571e431c17ae35617644 Mon Sep 17 00:00:00 2001
From: Andrey Breslav <andrey.breslav@jetbrains.com>
Date: Wed, 27 Nov 2013 15:27:59 +0400
Subject: [PATCH] TCO: Optimize all function calls, not only JetCallExpressions

---
 .../jet/codegen/ExpressionCodegen.java        | 19 ++++---------------
 .../codegen/TailRecursionGeneratorUtil.java   | 11 +++++------
 2 files changed, 9 insertions(+), 21 deletions(-)

diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
index 3b4ee88248e..5a142c5d341 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
@@ -1593,17 +1593,6 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
     public StackValue visitReturnExpression(@NotNull JetReturnExpression expression, StackValue receiver) {
         JetExpression returnedExpression = expression.getReturnedExpression();
         if (returnedExpression != null) {
-            if (returnedExpression instanceof JetCallExpression) {
-                JetCallExpression callExpression = (JetCallExpression) returnedExpression;
-                JetExpression calleeExpression = callExpression.getCalleeExpression();
-                if (calleeExpression != null) {
-                    ResolvedCall<? extends CallableDescriptor> resolvedCall = bindingContext.get(BindingContext.RESOLVED_CALL, calleeExpression);
-                    if (resolvedCall != null && tailRecursionGeneratorUtil.isTailRecursion(resolvedCall)) {
-                        return tailRecursionGeneratorUtil.generateTailRecursion(resolvedCall, callExpression);
-                    }
-                }
-            }
-
             gen(returnedExpression, returnType);
             doFinallyOnReturn();
             v.areturn(returnType);
@@ -1933,10 +1922,6 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             }
         }
 
-        if (tailRecursionGeneratorUtil.isTailRecursion(resolvedCall)) {
-            return tailRecursionGeneratorUtil.generateTailRecursion(resolvedCall, expression);
-        }
-
         return invokeFunction(call, receiver, resolvedCall);
     }
 
@@ -2011,6 +1996,10 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             StackValue receiver,
             ResolvedCall<? extends CallableDescriptor> resolvedCall
     ) {
+        if (tailRecursionGeneratorUtil.isTailRecursion(resolvedCall)) {
+            return tailRecursionGeneratorUtil.generateTailRecursion(resolvedCall, call);
+        }
+
         if (resolvedCall instanceof VariableAsFunctionResolvedCall) {
             VariableAsFunctionResolvedCall variableAsFunctionResolvedCall = (VariableAsFunctionResolvedCall) resolvedCall;
             ResolvedCallWithTrace<FunctionDescriptor> functionCall = variableAsFunctionResolvedCall.getFunctionCall();
diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/TailRecursionGeneratorUtil.java b/compiler/backend/src/org/jetbrains/jet/codegen/TailRecursionGeneratorUtil.java
index 78801994a80..fc79258b82d 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/TailRecursionGeneratorUtil.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/TailRecursionGeneratorUtil.java
@@ -17,7 +17,6 @@
 package org.jetbrains.jet.codegen;
 
 import com.google.common.collect.Lists;
-import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.asm4.Type;
 import org.jetbrains.asm4.commons.InstructionAdapter;
@@ -26,7 +25,7 @@ import org.jetbrains.jet.codegen.state.GenerationState;
 import org.jetbrains.jet.lang.descriptors.CallableDescriptor;
 import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
-import org.jetbrains.jet.lang.psi.JetCallExpression;
+import org.jetbrains.jet.lang.psi.Call;
 import org.jetbrains.jet.lang.psi.JetExpression;
 import org.jetbrains.jet.lang.psi.JetSimpleNameExpression;
 import org.jetbrains.jet.lang.psi.ValueArgument;
@@ -66,7 +65,7 @@ public class TailRecursionGeneratorUtil {
         return status != null && status.isDoGenerateTailRecursion();
     }
 
-    public StackValue generateTailRecursion(ResolvedCall<? extends CallableDescriptor> resolvedCall, JetCallExpression callExpression) {
+    public StackValue generateTailRecursion(ResolvedCall<? extends CallableDescriptor> resolvedCall, Call call) {
         CallableDescriptor fd = resolvedCall.getResultingDescriptor();
         assert fd instanceof FunctionDescriptor : "the resolved call is not refer to the function descriptor so why do we use generateTailRecursion for something strange?";
         CallableMethod callable = (CallableMethod) codegen.resolveToCallable((FunctionDescriptor) fd, false);
@@ -76,7 +75,7 @@ public class TailRecursionGeneratorUtil {
         // we can't store values to the variables in the loop above because it will affect expressions evaluation
         for (ValueParameterDescriptor parameterDescriptor : Lists.reverse(parametersStored)) {
             Type asmType = types.get(parameterDescriptor.getIndex());
-            int index = getParameterVariableIndex(parameterDescriptor, callExpression);
+            int index = getParameterVariableIndex(parameterDescriptor, call);
 
             v.store(index, asmType);
         }
@@ -127,14 +126,14 @@ public class TailRecursionGeneratorUtil {
         return descriptorsStored;
     }
 
-    private int getParameterVariableIndex(ValueParameterDescriptor parameterDescriptor, PsiElement node) {
+    private int getParameterVariableIndex(ValueParameterDescriptor parameterDescriptor, Call call) {
         int index = codegen.lookupLocalIndex(parameterDescriptor);
         if (index == -1) {
             index = codegen.lookupLocalIndex(parameterDescriptor.getOriginal());
         }
 
         if (index == -1) {
-            throw new CompilationException("Failed to obtain parameter index: " + parameterDescriptor.getName(), null, node);
+            throw new CompilationException("Failed to obtain parameter index: " + parameterDescriptor.getName(), null, call.getCallElement());
         }
 
         return index;