Introduce BodyResolveCache for storing resolve function body task in IDE

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/di/injection.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.java.lazy.SingleModuleClassResolver
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmCheckerProvider
 import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.BodyResolveCache
 import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzerForTopLevel
 import org.jetbrains.kotlin.resolve.jvm.JavaClassFinderPostConstruct
 import org.jetbrains.kotlin.resolve.jvm.JavaDescriptorResolver
@@ -39,9 +40,9 @@ public fun createContainerForReplWithJava(
     configureModule(moduleContext, KotlinJvmCheckerProvider, bindingTrace)
     configureJavaTopDownAnalysis(moduleContentScope, moduleContext.project)
 
-
     useInstance(additionalFileScopeProvider)
     useInstance(declarationProviderFactory)
+    useInstance(BodyResolveCache.ThrowException)
 
     useImpl<ScopeProvider>()
     useImpl<SingleModuleClassResolver>()

compiler/frontend.java/src/org/jetbrains/kotlin/frontend/java/di/injection.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.load.kotlin.DeserializationComponentsForJava
 import org.jetbrains.kotlin.load.kotlin.JvmVirtualFileFinderFactory
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmCheckerProvider
 import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.BodyResolveCache
 import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzer
 import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzerForTopLevel
 import org.jetbrains.kotlin.resolve.jvm.JavaClassFinderPostConstruct
@@ -76,6 +77,7 @@ public fun createContainerForLazyResolveWithJava(
     useInstance(moduleClassResolver)
 
     useInstance(declarationProviderFactory)
+    useInstance(BodyResolveCache.ThrowException)
     useImpl<ScopeProvider>()
     useImpl<LazyResolveToken>()
 }.let {
@@ -93,6 +95,7 @@ public fun createContainerForTopDownAnalyzerForJvm(
     configureModule(moduleContext, KotlinJvmCheckerProvider, bindingTrace)
     configureJavaTopDownAnalysis(moduleContentScope, moduleContext.project)
     useInstance(declarationProviderFactory)
+    useInstance(BodyResolveCache.ThrowException)
 
     useImpl<SingleModuleClassResolver>()
     useImpl<ScopeProvider>()

compiler/frontend/src/org/jetbrains/kotlin/frontend/di/injection.kt

@@ -53,24 +53,29 @@ public fun StorageComponentContainer.configureModule(
 
 public fun createContainerForBodyResolve(
         moduleContext: ModuleContext, bindingTrace: BindingTrace,
-        additionalCheckerProvider: AdditionalCheckerProvider, statementFilter: StatementFilter
+        additionalCheckerProvider: AdditionalCheckerProvider, statementFilter: StatementFilter,
+        dynamicTypesSettings: DynamicTypesSettings
 ): StorageComponentContainer = createContainer("BodyResolve") {
     configureModule(moduleContext, additionalCheckerProvider, bindingTrace)
 
     useInstance(statementFilter)
+    useInstance(dynamicTypesSettings)
+    useInstance(BodyResolveCache.ThrowException)
     useImpl<BodyResolver>()
 }
 
 public fun createContainerForLazyBodyResolve(
         moduleContext: ModuleContext, kotlinCodeAnalyzer: KotlinCodeAnalyzer,
         bindingTrace: BindingTrace, additionalCheckerProvider: AdditionalCheckerProvider,
-        dynamicTypesSettings: DynamicTypesSettings
+        dynamicTypesSettings: DynamicTypesSettings,
+        bodyResolveCache: BodyResolveCache
 ): StorageComponentContainer = createContainer("LazyBodyResolve") {
     configureModule(moduleContext, additionalCheckerProvider, bindingTrace)
 
     useInstance(kotlinCodeAnalyzer)
     useInstance(kotlinCodeAnalyzer.getScopeProvider())
     useInstance(dynamicTypesSettings)
+    useInstance(bodyResolveCache)
     useImpl<LazyTopDownAnalyzerForTopLevel>()
 }
 
@@ -90,6 +95,7 @@ public fun createContainerForLazyLocalClassifierAnalyzer(
 
     useInstance(NoTopLevelDescriptorProvider)
     useInstance(NoFileScopeProvider)
+    useInstance(BodyResolveCache.ThrowException)
 
     useImpl<DeclarationScopeProviderForLocalClassifierAnalyzer>()
     useImpl<LocalLazyDeclarationResolver>()

compiler/frontend/src/org/jetbrains/kotlin/resolve/BodyResolveCache.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.resolve
+
+import org.jetbrains.kotlin.psi.JetNamedFunction
+
+public interface BodyResolveCache {
+    public open fun resolveFunctionBody(function: JetNamedFunction): BindingContext
+
+    public object ThrowException : BodyResolveCache {
+        public override fun resolveFunctionBody(function: JetNamedFunction): BindingContext {
+            throw UnsupportedOperationException()
+        }
+    }
+}
+

compiler/frontend/src/org/jetbrains/kotlin/resolve/BodyResolver.java

@@ -26,6 +26,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.lexer.JetTokens;
 import org.jetbrains.kotlin.psi.*;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 import org.jetbrains.kotlin.resolve.calls.CallResolver;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
@@ -65,6 +66,7 @@ public class BodyResolver {
     private FunctionAnalyzerExtension functionAnalyzerExtension;
     private AdditionalCheckerProvider additionalCheckerProvider;
     private ValueParameterResolver valueParameterResolver;
+    private BodyResolveCache bodyResolveCache;
 
     //<editor-fold desc="Injector Setters">
     @Inject
@@ -121,6 +123,11 @@ public class BodyResolver {
     public void setValueParameterResolver(ValueParameterResolver valueParameterResolver) {
         this.valueParameterResolver = valueParameterResolver;
     }
+
+    @Inject
+    public void setBodyResolveCache(BodyResolveCache bodyResolveCache) {
+        this.bodyResolveCache = bodyResolveCache;
+    }
     //</editor-fold>
 
     private void resolveBehaviorDeclarationBodies(@NotNull BodiesResolveContext c) {
@@ -745,16 +752,17 @@ public class BodyResolver {
     private void resolveFunctionBodies(@NotNull BodiesResolveContext c) {
         for (Map.Entry<JetNamedFunction, SimpleFunctionDescriptor> entry : c.getFunctions().entrySet()) {
             JetNamedFunction declaration = entry.getKey();
-            SimpleFunctionDescriptor descriptor = entry.getValue();
 
-            computeDeferredType(descriptor.getReturnType());
+            JetScope scope = c.getDeclaringScope(declaration);
+            assert scope != null : "Scope is null: " + PsiUtilPackage.getElementTextWithContext(declaration);
 
-            JetScope declaringScope = c.getDeclaringScope(declaration);
-            assert declaringScope != null;
-
-            resolveFunctionBody(c.getOuterDataFlowInfo(), trace, declaration, descriptor, declaringScope);
-
-            assert descriptor.getReturnType() != null;
+            if (!c.getTopDownAnalysisMode().getIsLocalDeclarations() && !(bodyResolveCache instanceof BodyResolveCache.ThrowException) &&
+                expressionTypingServices.getStatementFilter() != StatementFilter.NONE) {
+                bodyResolveCache.resolveFunctionBody(declaration).addOwnDataTo(trace, true);
+            }
+            else {
+                resolveFunctionBody(c.getOuterDataFlowInfo(), trace, declaration, entry.getValue(), scope);
+            }
         }
     }
 
@@ -765,7 +773,11 @@ public class BodyResolver {
             @NotNull FunctionDescriptor functionDescriptor,
             @NotNull JetScope declaringScope
     ) {
+        computeDeferredType(functionDescriptor.getReturnType());
+
         resolveFunctionBody(outerDataFlowInfo, trace, function, functionDescriptor, declaringScope, null, null);
+
+        assert functionDescriptor.getReturnType() != null;
     }
 
     public void resolveFunctionBody(

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java

@@ -62,6 +62,10 @@ public class ExpressionTypingServices {
         this.expressionTypingFacade = ExpressionTypingVisitorDispatcher.create(components);
     }
 
+    public StatementFilter getStatementFilter() {
+        return statementFilter;
+    }
+
     @NotNull
     public JetType safeGetType(
             @NotNull JetScope scope,

idea/ide-common/src/org/jetbrains/kotlin/resolve/lazy/ElementResolver.kt

@@ -40,6 +40,9 @@ import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyClassDescriptor
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyPackageDescriptor
 import org.jetbrains.kotlin.resolve.scopes.ChainedScope
 import org.jetbrains.kotlin.resolve.scopes.JetScope
+import org.jetbrains.kotlin.storage.LockBasedLazyResolveStorageManager
+import org.jetbrains.kotlin.storage.StorageManager
+import org.jetbrains.kotlin.types.DynamicTypesSettings
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
@@ -160,7 +163,9 @@ public abstract class ElementResolver protected constructor(
             }
         }
 
-        JetFlowInformationProvider(resolveElement, trace).checkDeclaration()
+        val controlFlowTrace = DelegatingBindingTrace(trace.getBindingContext(), "Element control flow resolve", resolveElement)
+        JetFlowInformationProvider(resolveElement, controlFlowTrace).checkDeclaration()
+        controlFlowTrace.addOwnDataTo(trace, null, false)
 
         return trace.getBindingContext()
     }
@@ -404,7 +409,7 @@ public abstract class ElementResolver protected constructor(
         val globalContext = SimpleGlobalContext(resolveSession.getStorageManager(), resolveSession.getExceptionTracker())
         return createContainerForBodyResolve(
                 globalContext.withProject(file.getProject()).withModule(resolveSession.getModuleDescriptor()),
-                trace, getAdditionalCheckerProvider(file), statementFilter
+                trace, getAdditionalCheckerProvider(file), statementFilter, getDynamicTypesSettings(file)
         ).get<BodyResolver>()
     }
 
@@ -478,6 +483,7 @@ public abstract class ElementResolver protected constructor(
     }
 
     protected abstract fun getAdditionalCheckerProvider(jetFile: JetFile): AdditionalCheckerProvider
+    protected abstract fun getDynamicTypesSettings(jetFile: JetFile): DynamicTypesSettings
 
     private class BodyResolveContextForLazy(
             private val topDownAnalysisMode: TopDownAnalysisMode,
@@ -504,3 +510,4 @@ public abstract class ElementResolver protected constructor(
         override fun getTopDownAnalysisMode() = topDownAnalysisMode
     }
 }
+

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinResolveCache.kt

@@ -238,7 +238,8 @@ private object KotlinResolveDataProvider {
                     resolveSession,
                     trace,
                     targetPlatform.getAdditionalCheckerProvider(),
-                    targetPlatform.getDynamicTypesSettings()
+                    targetPlatform.getDynamicTypesSettings(),
+                    resolveSession.getBodyResolveCache()
             ).get<LazyTopDownAnalyzerForTopLevel>()
 
             lazyTopDownAnalyzer.analyzeDeclarations(

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/ResolveElementCache.java

@@ -28,18 +28,21 @@ import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingFunctionShortNameIn
 import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingPropertyShortNameIndex;
 import org.jetbrains.kotlin.psi.JetElement;
 import org.jetbrains.kotlin.psi.JetFile;
+import org.jetbrains.kotlin.psi.JetNamedFunction;
 import org.jetbrains.kotlin.resolve.AdditionalCheckerProvider;
 import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.BodyResolveCache;
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.resolve.lazy.ElementResolver;
 import org.jetbrains.kotlin.resolve.lazy.ProbablyNothingCallableNames;
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession;
 import org.jetbrains.kotlin.storage.LazyResolveStorageManager;
 import org.jetbrains.kotlin.storage.MemoizedFunctionToNotNull;
+import org.jetbrains.kotlin.types.DynamicTypesSettings;
 
 import java.util.Collection;
 
-public class ResolveElementCache extends ElementResolver {
+public class ResolveElementCache extends ElementResolver implements BodyResolveCache {
     private final Project project;
     private final CachedValue<MemoizedFunctionToNotNull<JetElement, BindingContext>> additionalResolveCache;
 
@@ -89,6 +92,12 @@ public class ResolveElementCache extends ElementResolver {
         return TargetPlatformDetector.getPlatform(jetFile).getAdditionalCheckerProvider();
     }
 
+    @NotNull
+    @Override
+    public DynamicTypesSettings getDynamicTypesSettings(@NotNull JetFile jetFile) {
+        return TargetPlatformDetector.getPlatform(jetFile).getDynamicTypesSettings();
+    }
+
     @NotNull
     @Override
     protected ProbablyNothingCallableNames probablyNothingCallableNames() {
@@ -107,4 +116,10 @@ public class ResolveElementCache extends ElementResolver {
             }
         };
     }
+
+    @NotNull
+    @Override
+    public BindingContext resolveFunctionBody(@NotNull JetNamedFunction function) {
+        return getElementAdditionalResolve(function);
+    }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/ResolveSessionForBodies.java

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.BindingContextUtils;
 import org.jetbrains.kotlin.resolve.QualifiedExpressionResolver;
+import org.jetbrains.kotlin.resolve.BodyResolveCache;
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.resolve.lazy.KotlinCodeAnalyzer;
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession;
@@ -134,6 +135,11 @@ public class ResolveSessionForBodies implements KotlinCodeAnalyzer {
         return resolveSession.getExceptionTracker();
     }
 
+    @NotNull
+    public BodyResolveCache getBodyResolveCache() {
+        return resolveElementCache;
+    }
+
     @NotNull
     @Override
     public PackageFragmentProvider getPackageFragmentProvider() {

js/js.frontend/src/org/jetbrains/kotlin/frontend/js/di/injection.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.context.ModuleContext
 import org.jetbrains.kotlin.frontend.di.configureModule
 import org.jetbrains.kotlin.js.resolve.KotlinJsCheckerProvider
 import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.BodyResolveCache
 import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzerForTopLevel
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession
 import org.jetbrains.kotlin.resolve.lazy.ScopeProvider
@@ -39,6 +40,7 @@ public fun createTopDownAnalyzerForJs(
         configureModule(moduleContext, KotlinJsCheckerProvider, bindingTrace)
 
         useInstance(declarationProviderFactory)
+        useInstance(BodyResolveCache.ThrowException)
         useImpl<ScopeProvider>()
         useImpl<ResolveSession>()
         useImpl<LazyTopDownAnalyzerForTopLevel>()