Refactoring: getDeclaringScopes() in BodiesResolveContext used to be a map, and now it is a Function.

compiler/frontend/src/org/jetbrains/jet/lang/resolve/BodiesResolveContext.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.lang.resolve;
 
+import com.google.common.base.Function;
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.descriptors.*;
@@ -33,7 +34,7 @@ public interface BodiesResolveContext {
     Map<JetObjectDeclaration, MutableClassDescriptor> getObjects();
     Map<JetProperty, PropertyDescriptor> getProperties();
     Map<JetNamedFunction, SimpleFunctionDescriptor> getFunctions();
-    Map<JetDeclaration, JetScope> getDeclaringScopes();
+    Function<JetDeclaration, JetScope> getDeclaringScopes();
     Map<JetScript, ScriptDescriptor> getScripts();
     Map<JetScript, WritableScope> getScriptScopes();
 

compiler/frontend/src/org/jetbrains/jet/lang/resolve/BodyResolver.java

@@ -399,7 +399,7 @@ public class BodyResolver {
                 if (initializer != null) {
                     ConstructorDescriptor primaryConstructor = classDescriptor.getUnsubstitutedPrimaryConstructor();
                     if (primaryConstructor != null) {
-                        JetScope declaringScopeForPropertyInitializer = this.context.getDeclaringScopes().get(property);
+                        JetScope declaringScopeForPropertyInitializer = this.context.getDeclaringScopes().apply(property);
                         resolvePropertyInitializer(property, propertyDescriptor, initializer, declaringScopeForPropertyInitializer);
                     }
                 }
@@ -419,7 +419,7 @@ public class BodyResolver {
 
             computeDeferredType(propertyDescriptor.getReturnType());
 
-            JetScope declaringScope = this.context.getDeclaringScopes().get(property);
+            JetScope declaringScope = this.context.getDeclaringScopes().apply(property);
 
             JetExpression initializer = property.getInitializer();
             if (initializer != null) {
@@ -431,7 +431,7 @@ public class BodyResolver {
     }
 
     private JetScope makeScopeForPropertyAccessor(@NotNull JetPropertyAccessor accessor, PropertyDescriptor propertyDescriptor) {
-        JetScope declaringScope = context.getDeclaringScopes().get(accessor);
+        JetScope declaringScope = context.getDeclaringScopes().apply(accessor);
 
         JetScope propertyDeclarationInnerScope = descriptorResolver.getPropertyDeclarationInnerScope(
                 declaringScope, propertyDescriptor.getTypeParameters(), propertyDescriptor.getReceiverParameter(), trace);
@@ -490,7 +490,7 @@ public class BodyResolver {
 
             computeDeferredType(descriptor.getReturnType());
 
-            JetScope declaringScope = this.context.getDeclaringScopes().get(declaration);
+            JetScope declaringScope = this.context.getDeclaringScopes().apply(declaration);
             assert declaringScope != null;
 
             resolveFunctionBody(trace, declaration, descriptor, declaringScope);

compiler/frontend/src/org/jetbrains/jet/lang/resolve/CachedBodiesResolveContext.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.lang.resolve;
 
+import com.google.common.base.Function;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import org.jetbrains.annotations.NotNull;
@@ -38,7 +39,7 @@ public class CachedBodiesResolveContext implements BodiesResolveContext {
     private final Map<JetObjectDeclaration, MutableClassDescriptor> objects;
     private final Map<JetProperty, PropertyDescriptor> properties;
     private final Map<JetNamedFunction, SimpleFunctionDescriptor> functions;
-    private final Map<JetDeclaration, JetScope> declaringScopes;
+    private final Function<JetDeclaration, JetScope> declaringScopes;
     private final Map<JetScript, ScriptDescriptor> scripts;
     private final Map<JetScript, WritableScope> scriptScopes;
 
@@ -49,7 +50,7 @@ public class CachedBodiesResolveContext implements BodiesResolveContext {
         objects = Collections.unmodifiableMap(context.getObjects());
         properties = Collections.unmodifiableMap(context.getProperties());
         functions = Collections.unmodifiableMap(context.getFunctions());
-        declaringScopes = Collections.unmodifiableMap(context.getDeclaringScopes());
+        declaringScopes = context.getDeclaringScopes();
         scripts = Collections.unmodifiableMap(context.getScripts());
         scriptScopes = Collections.unmodifiableMap(context.getScriptScopes());
 
@@ -77,7 +78,7 @@ public class CachedBodiesResolveContext implements BodiesResolveContext {
     }
 
     @Override
-    public Map<JetDeclaration, JetScope> getDeclaringScopes() {
+    public Function<JetDeclaration, JetScope> getDeclaringScopes() {
         return declaringScopes;
     }
 

compiler/frontend/src/org/jetbrains/jet/lang/resolve/DeclarationResolver.java

@@ -180,7 +180,7 @@ public class DeclarationResolver {
                     SimpleFunctionDescriptor functionDescriptor = descriptorResolver.resolveFunctionDescriptor(namespaceLike.getOwnerForChildren(), scopeForFunctions, function, trace);
                     namespaceLike.addFunctionDescriptor(functionDescriptor);
                     context.getFunctions().put(function, functionDescriptor);
-                    context.getDeclaringScopes().put(function, scopeForFunctions);
+                    context.registerDeclaringScope(function, scopeForFunctions);
                 }
 
                 @Override
@@ -188,12 +188,14 @@ public class DeclarationResolver {
                     PropertyDescriptor propertyDescriptor = descriptorResolver.resolvePropertyDescriptor(namespaceLike.getOwnerForChildren(), scopeForPropertyInitializers, property, trace);
                     namespaceLike.addPropertyDescriptor(propertyDescriptor);
                     context.getProperties().put(property, propertyDescriptor);
-                    context.getDeclaringScopes().put(property, scopeForPropertyInitializers);
+                    context.registerDeclaringScope(property, scopeForPropertyInitializers);
-                    if (property.getGetter() != null) {
+                    JetPropertyAccessor getter = property.getGetter();
-                        context.getDeclaringScopes().put(property.getGetter(), scopeForPropertyAccessors);
+                    if (getter != null) {
+                        context.registerDeclaringScope(getter, scopeForPropertyAccessors);
                     }
-                    if (property.getSetter() != null) {
+                    JetPropertyAccessor setter = property.getSetter();
-                        context.getDeclaringScopes().put(property.getSetter(), scopeForPropertyAccessors);
+                    if (setter != null) {
+                        context.registerDeclaringScope(setter, scopeForPropertyAccessors);
                     }
                 }
 

compiler/frontend/src/org/jetbrains/jet/lang/resolve/TopDownAnalysisContext.java

@@ -16,6 +16,8 @@
 
 package org.jetbrains.jet.lang.resolve;
 
+import com.google.common.base.Function;
+import com.google.common.base.Functions;
 import com.google.common.collect.Maps;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
@@ -138,8 +140,12 @@ public class TopDownAnalysisContext implements BodiesResolveContext {
     }
 
     @Override
-    public Map<JetDeclaration, JetScope> getDeclaringScopes() {
+    public Function<JetDeclaration, JetScope> getDeclaringScopes() {
-        return declaringScopes;
+        return Functions.forMap(declaringScopes);
+    }
+
+    public void registerDeclaringScope(@NotNull JetDeclaration declaration, @NotNull JetScope scope) {
+        declaringScopes.put(declaration, scope);
     }
 
     @Override

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/ResolveSessionUtils.java

@@ -16,8 +16,9 @@
 
 package org.jetbrains.jet.lang.resolve.lazy;
 
+import com.google.common.base.Function;
+import com.google.common.base.Functions;
 import com.google.common.base.Predicates;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
@@ -40,12 +41,20 @@ import java.util.*;
  * @author Nikolay Krasko
  */
 public class ResolveSessionUtils {
+
     private ResolveSessionUtils() {
     }
 
-    private static class EmptyBodyResolveContext implements BodiesResolveContext {
+    @SuppressWarnings("unchecked")
+    private static final BodyResolveContextForLazy EMPTY_CONTEXT = new BodyResolveContextForLazy((Function) Functions.constant(null));
 
-        private Map<JetDeclaration, JetScope> declarationScopesMap = Collections.emptyMap();
+    private static class BodyResolveContextForLazy implements BodiesResolveContext {
+
+        private final Function<JetDeclaration, JetScope> declaringScopes;
+
+        private BodyResolveContextForLazy(@NotNull Function<JetDeclaration, JetScope> declaringScopes) {
+            this.declaringScopes = declaringScopes;
+        }
 
         @Override
         public Map<JetClass, MutableClassDescriptor> getClasses() {
@@ -68,12 +77,8 @@ public class ResolveSessionUtils {
         }
 
         @Override
-        public Map<JetDeclaration, JetScope> getDeclaringScopes() {
+        public Function<JetDeclaration, JetScope> getDeclaringScopes() {
-            return declarationScopesMap;
+            return declaringScopes;
-        }
-
-        public void setDeclaringScopes(Map<JetDeclaration, JetScope> declarationScopes) {
-            declarationScopesMap = declarationScopes;
         }
 
         @Override
@@ -142,7 +147,7 @@ public class ResolveSessionUtils {
 
     private static void delegationSpecifierAdditionalResolve(final ResolveSession resolveSession,
             final JetDelegationSpecifierList specifier, DelegatingBindingTrace trace, JetFile file) {
-        BodyResolver bodyResolver = createBodyResolver(trace, file, resolveSession.getModuleConfiguration());
+        BodyResolver bodyResolver = createBodyResolverWithEmptyContext(trace, file, resolveSession.getModuleConfiguration());
 
         JetClassOrObject classOrObject = (JetClassOrObject) specifier.getParent();
         LazyClassDescriptor descriptor = (LazyClassDescriptor) resolveSession.resolveToDescriptor(classOrObject);
@@ -157,42 +162,35 @@ public class ResolveSessionUtils {
     }
 
     private static void propertyAdditionalResolve(final ResolveSession resolveSession, final JetProperty jetProperty, DelegatingBindingTrace trace, JetFile file) {
-        EmptyBodyResolveContext bodyResolveContext = new EmptyBodyResolveContext();
+        final JetScope propertyResolutionScope = resolveSession.getInjector().getScopeProvider().getResolutionScopeForDeclaration(
+                jetProperty);
+
+        BodyResolveContextForLazy bodyResolveContext = new BodyResolveContextForLazy(new Function<JetDeclaration, JetScope>() {
+            @Override
+            public JetScope apply(JetDeclaration declaration) {
+                assert declaration.getParent() == jetProperty : "Must be called only for property accessors, but called for " + declaration;
+                return propertyResolutionScope;
+            }
+        });
         BodyResolver bodyResolver = createBodyResolver(trace, file, bodyResolveContext, resolveSession.getModuleConfiguration());
         PropertyDescriptor descriptor = (PropertyDescriptor) resolveSession.resolveToDescriptor(jetProperty);
 
         JetExpression propertyInitializer = jetProperty.getInitializer();
 
-        final JetScope propertyResolutionScope = resolveSession.getInjector().getScopeProvider().getResolutionScopeForDeclaration(
-                jetProperty);
-
         if (propertyInitializer != null) {
             bodyResolver.resolvePropertyInitializer(jetProperty, descriptor, propertyInitializer, propertyResolutionScope);
         }
 
-        addAccessorDeclaringScopes(jetProperty, bodyResolveContext, propertyResolutionScope);
         bodyResolver.resolvePropertyAccessors(jetProperty, descriptor);
     }
 
-    private static void addAccessorDeclaringScopes(
-            JetProperty jetProperty,
-            EmptyBodyResolveContext bodyResolveContext,
-            JetScope propertyResolutionScope
-    ) {
-        ImmutableMap.Builder<JetDeclaration, JetScope> builder = ImmutableMap.builder();
-        for (final JetPropertyAccessor propertyAccessor : jetProperty.getAccessors()) {
-            builder.put(propertyAccessor, propertyResolutionScope);
-        }
-        bodyResolveContext.setDeclaringScopes(builder.build());
-    }
-
     private static void functionAdditionalResolve(
             ResolveSession resolveSession,
             JetNamedFunction namedFunction,
             DelegatingBindingTrace trace,
             JetFile file
     ) {
-        BodyResolver bodyResolver = createBodyResolver(trace, file, resolveSession.getModuleConfiguration());
+        BodyResolver bodyResolver = createBodyResolverWithEmptyContext(trace, file, resolveSession.getModuleConfiguration());
         JetScope scope = resolveSession.getInjector().getScopeProvider().getResolutionScopeForDeclaration(namedFunction);
         FunctionDescriptor functionDescriptor = (FunctionDescriptor) resolveSession.resolveToDescriptor(namedFunction);
         bodyResolver.resolveFunctionBody(trace, namedFunction, functionDescriptor, scope);
@@ -204,7 +202,7 @@ public class ResolveSessionUtils {
             DelegatingBindingTrace trace,
             JetFile file
     ) {
-        BodyResolver bodyResolver = createBodyResolver(trace, file, resolveSession.getModuleConfiguration());
+        BodyResolver bodyResolver = createBodyResolverWithEmptyContext(trace, file, resolveSession.getModuleConfiguration());
         JetClassOrObject classOrObject = PsiTreeUtil.getParentOfType(classInitializer, JetClassOrObject.class);
         LazyClassDescriptor classOrObjectDescriptor = (LazyClassDescriptor) resolveSession.resolveToDescriptor(classOrObject);
         bodyResolver.resolveAnonymousInitializers(classOrObject, classOrObjectDescriptor.getUnsubstitutedPrimaryConstructor(),
@@ -213,7 +211,7 @@ public class ResolveSessionUtils {
         return true;
     }
 
-    private static BodyResolver createBodyResolver(DelegatingBindingTrace trace, JetFile file, EmptyBodyResolveContext bodyResolveContext,
+    private static BodyResolver createBodyResolver(DelegatingBindingTrace trace, JetFile file, BodyResolveContextForLazy bodyResolveContext,
             ModuleConfiguration moduleConfiguration) {
         TopDownAnalysisParameters parameters = new TopDownAnalysisParameters(
                 Predicates.<PsiFile>alwaysTrue(), false, true, Collections.<AnalyzerScriptParameter>emptyList());
@@ -221,8 +219,12 @@ public class ResolveSessionUtils {
         return bodyResolve.getBodyResolver();
     }
 
-    private static BodyResolver createBodyResolver(DelegatingBindingTrace trace, JetFile file, ModuleConfiguration moduleConfiguration) {
+    private static BodyResolver createBodyResolverWithEmptyContext(
-        return createBodyResolver(trace, file, new EmptyBodyResolveContext(), moduleConfiguration);
+            DelegatingBindingTrace trace,
+            JetFile file,
+            ModuleConfiguration moduleConfiguration
+    ) {
+        return createBodyResolver(trace, file, EMPTY_CONTEXT, moduleConfiguration);
     }
 
     private static JetScope getExpressionResolutionScope(@NotNull ResolveSession resolveSession, @NotNull JetExpression expression) {