ability to infer types of groovy vars in certain place

java/java-impl/src/com/intellij/codeInsight/template/macro/ArrayVariableMacro.java

@@ -43,7 +43,7 @@ public class ArrayVariableMacro extends VariableTypeMacroBase {
     PsiElement place = file.findElementAt(offset);
     PsiVariable[] variables = MacroUtil.getVariablesVisibleAt(place, "");
     for (PsiVariable variable : variables) {
-      PsiType type = variable.getType();
+      PsiType type = VariableTypeCalculator.getVarTypeAt(variable, place);
       if (type instanceof PsiArrayType) {
         array.add(variable);
       }

java/java-impl/src/com/intellij/codeInsight/template/macro/IterableComponentTypeMacro.java

@@ -50,6 +50,8 @@ public class IterableComponentTypeMacro implements Macro {
     PsiExpression expr = MacroUtil.resultToPsiExpression(result, context);
     if (expr == null) return null;
     PsiType type = expr.getType();
+
+
     if (type instanceof PsiArrayType) {
       return new PsiTypeResult(((PsiArrayType)type).getComponentType(), project);
     }

java/java-impl/src/com/intellij/codeInsight/template/macro/IterableVariableMacro.java

@@ -68,7 +68,7 @@ public class IterableVariableMacro extends VariableTypeMacroBase {
       final PsiElement parent = var.getParent();
       if (parent instanceof PsiForeachStatement && parent == PsiTreeUtil.getParentOfType(place, PsiForeachStatement.class)) continue;
 
-      PsiType type = var.getType();
+      PsiType type = VariableTypeCalculator.getVarTypeAt(var, place);
       if (type instanceof PsiArrayType || iterableType.isAssignableFrom(type)) {
         result.add(var);
       }

java/java-impl/src/com/intellij/codeInsight/template/macro/VariableOfTypeMacro.java

@@ -97,7 +97,7 @@ public class VariableOfTypeMacro implements Macro {
         }
       }
 
-      PsiType type1 = var.getType();
+      PsiType type1 = VariableTypeCalculator.getVarTypeAt(var, place);
       if (type == null || type.isAssignableFrom(type1)) {
         array.add(var);
       }

java/java-impl/src/com/intellij/codeInsight/template/macro/VariableTypeCalculator.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2000-2011 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 com.intellij.codeInsight.template.macro;
+
+import com.intellij.openapi.extensions.ExtensionPointName;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiType;
+import com.intellij.psi.PsiVariable;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+* @author Max Medvedev
+*/
+public abstract class VariableTypeCalculator {
+  public static final ExtensionPointName<VariableTypeCalculator> EP_NAME =
+    ExtensionPointName.create("com.intellij.variableTypeCalculator");
+
+  @Nullable
+  public abstract PsiType inferVarTypeAt(@NotNull PsiVariable var, @NotNull PsiElement place);
+
+  /**
+   * @return inferred type of variable in the context of place
+   */
+  @NotNull
+  public static PsiType getVarTypeAt(@NotNull PsiVariable var, @NotNull PsiElement place) {
+    for (VariableTypeCalculator calculator : EP_NAME.getExtensions()) {
+      final PsiType type = calculator.inferVarTypeAt(var, place);
+      if (type != null) return type;
+    }
+
+    return var.getType();
+  }
+}

platform/platform-resources/src/META-INF/LangExtensionPoints.xml

@@ -193,6 +193,7 @@
   <extensionPoint name="liveTemplateOptionalProcessor" interface="com.intellij.codeInsight.template.impl.TemplateOptionalProcessor"/>
   <extensionPoint name="liveTemplatePreprocessor" interface="com.intellij.codeInsight.template.impl.TemplatePreprocessor"/>
   <extensionPoint name="customLiveTemplate" interface="com.intellij.codeInsight.template.CustomLiveTemplate"/>
+  <extensionPoint name="variableTypeCalculator" interface="com.intellij.codeInsight.template.macro.VariableTypeCalculator"/>
 
   <extensionPoint name="fileTemplateGroup"
                   interface="com.intellij.ide.fileTemplates.FileTemplateGroupDescriptorFactory"/>

plugins/groovy/src/META-INF/plugin.xml

@@ -193,6 +193,7 @@
 
     <liveTemplateContext implementation="org.jetbrains.plugins.groovy.template.GroovyTemplateContextType"/>
     <liveTemplateOptionalProcessor implementation="org.jetbrains.plugins.groovy.template.GroovyShortenFQNamesProcessor"/>
+    <variableTypeCalculator implementation="org.jetbrains.plugins.groovy.template.GroovyVariableTypeCalculator"/>
 
     <gotoSymbolContributor implementation="org.jetbrains.plugins.groovy.gotoclass.GroovyGoToSymbolContributor"/>
     <lang.refactoringSupport language="Groovy"

plugins/groovy/src/org/jetbrains/plugins/groovy/codeInspection/utils/ControlFlowUtils.java

@@ -634,13 +634,14 @@ public class ControlFlowUtils {
 
   @Nullable
   public static GrControlFlowOwner findControlFlowOwner(PsiElement place) {
-    while (place.getParent() != null) {
+    while (true) {
       place = place.getParent();
+      if (place == null) return null;
       if (place instanceof GrClosableBlock) return (GrClosableBlock)place;
       if (place instanceof GrMethod) return ((GrMethod)place).getBlock();
       if (place instanceof GroovyFile) return (GroovyFile)place;
+      if (place instanceof GrClassInitializer) return ((GrClassInitializer)place).getBlock();
     }
-    return null;
   }
 
   /**

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/TypeInferenceHelper.java

@@ -18,6 +18,7 @@ package org.jetbrains.plugins.groovy.lang.psi.impl;
 import com.intellij.openapi.util.Computable;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.RecursionManager;
+import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiType;
 import com.intellij.psi.util.*;
@@ -25,16 +26,12 @@ import gnu.trove.TIntHashSet;
 import gnu.trove.TIntObjectHashMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.groovy.codeInspection.utils.ControlFlowUtils;
 import org.jetbrains.plugins.groovy.lang.lexer.TokenSets;
 import org.jetbrains.plugins.groovy.lang.psi.GrControlFlowOwner;
-import org.jetbrains.plugins.groovy.lang.psi.GroovyFileBase;
-import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElement;
-import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrClassInitializer;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrVariable;
-import org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks.GrClosableBlock;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.*;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.params.GrParameter;
-import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.members.GrMethod;
 import org.jetbrains.plugins.groovy.lang.psi.api.types.GrTypeElement;
 import org.jetbrains.plugins.groovy.lang.psi.controlFlow.AssertionInstruction;
 import org.jetbrains.plugins.groovy.lang.psi.controlFlow.Instruction;
@@ -44,6 +41,9 @@ import org.jetbrains.plugins.groovy.lang.psi.dataFlow.reachingDefs.ReachingDefin
 import org.jetbrains.plugins.groovy.lang.psi.dataFlow.reachingDefs.ReachingDefinitionsSemilattice;
 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions.TypesUtil;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 
 /**
@@ -57,47 +57,87 @@ public class TypeInferenceHelper {
     return RecursionManager.doPreventingRecursion(refExpr, true, new Computable<PsiType>() {
       @Override
       public PsiType compute() {
-        @SuppressWarnings("unchecked") GroovyPsiElement scope =
-          PsiTreeUtil.getParentOfType(refExpr, GrMethod.class, GrClosableBlock.class, GrClassInitializer.class, GroovyFileBase.class);
-        if (scope instanceof GrMethod) {
-          scope = ((GrMethod)scope).getBlock();
-        }
-        else if (scope instanceof GrClassInitializer) {
-          scope = ((GrClassInitializer)scope).getBlock();
+        final GrControlFlowOwner scope = ControlFlowUtils.findControlFlowOwner(refExpr);
+        if (scope == null) return null;
+
+        final Instruction[] flow = scope.getControlFlow();
+        ReadWriteVariableInstruction instruction = findInstruction(refExpr, flow);
+        if (instruction == null) return null;
+        
+        if (instruction.isWrite()) {
+          return getInitializerType(refExpr);
         }
 
-        if (scope != null) {
-          final Instruction[] flow = ((GrControlFlowOwner)scope).getControlFlow();
-          ReadWriteVariableInstruction instruction = findInstruction(refExpr, flow);
-          if (instruction == null) {
-            return null;
-          }
-          if (instruction.isWrite()) {
-            return getInitializerType(refExpr);
-          }
-
-          final Pair<ReachingDefinitionsDfaInstance, List<TIntObjectHashMap<TIntHashSet>>> pair = getDefUseMaps((GrControlFlowOwner)scope);
-
-          final int varIndex = pair.first.getVarIndex(refExpr.getReferenceName());
-          final TIntObjectHashMap<TIntHashSet> allDefs = pair.second.get(instruction.num());
-          final TIntHashSet varDefs = allDefs.get(varIndex);
-          if (varDefs != null) {
-            PsiType result = null;
-            for (int defIndex : varDefs.toArray()) {
-              PsiType defType = getDefinitionType(flow[defIndex]);
-              if (defType != null) {
-                defType = TypesUtil.boxPrimitiveType(defType, scope.getManager(), scope.getResolveScope());
-                result = result == null ? defType : TypesUtil.getLeastUpperBound(result, defType, scope.getManager());
-              }
-            }
-            return result;
-          }
-        }
-        return null;
+        return getInferredType(refExpr.getReferenceName(), instruction, flow, scope);
       }
 
     });
   }
+  
+  @Nullable
+  public static PsiType getInferredType(@NotNull PsiElement place, String variableName) {
+    final GrControlFlowOwner scope = ControlFlowUtils.findControlFlowOwner(place);
+    if (scope == null) return null;
+
+    final Instruction[] flow = scope.getControlFlow();
+    Instruction instruction = findInstructionAt(place, flow);
+    if (instruction == null) return null;
+
+    return getInferredType(variableName, instruction, flow, scope);
+  }
+
+  @Nullable
+  private static Instruction findInstructionAt(PsiElement place, Instruction[] flow) {
+    List<Instruction> applicable = new ArrayList<Instruction>();
+    for (Instruction instruction : flow) {
+      final PsiElement element = instruction.getElement();
+      if (element == null) continue;
+
+      if (element == place) return instruction;
+
+      if (PsiTreeUtil.isAncestor(element, place, true)) {
+        applicable.add(instruction);
+      }
+    }
+    if (applicable.size() == 0) return null;
+
+    Collections.sort(applicable, new Comparator<Instruction>() {
+      @Override
+      public int compare(Instruction o1, Instruction o2) {
+        final TextRange t1 = o1.getElement().getTextRange();
+        final TextRange t2 = o2.getElement().getTextRange();
+        final int s1 = t1.getStartOffset();
+        final int s2 = t2.getStartOffset();
+
+        if (s1 == s2) {
+          return t1.getEndOffset() - t2.getEndOffset();
+        }
+        return s2 - s1;
+      }
+    });
+
+    return applicable.get(0);
+  }
+
+  @Nullable
+  private static PsiType getInferredType(String varName, Instruction instruction, Instruction[] flow, GrControlFlowOwner scope) {
+    final Pair<ReachingDefinitionsDfaInstance, List<TIntObjectHashMap<TIntHashSet>>> pair = getDefUseMaps(scope);
+
+    final int varIndex = pair.first.getVarIndex(varName);
+    final TIntObjectHashMap<TIntHashSet> allDefs = pair.second.get(instruction.num());
+    final TIntHashSet varDefs = allDefs.get(varIndex);
+    if (varDefs == null) return null;
+
+    PsiType result = null;
+    for (int defIndex : varDefs.toArray()) {
+      PsiType defType = getDefinitionType(flow[defIndex]);
+      if (defType != null) {
+        defType = TypesUtil.boxPrimitiveType(defType, scope.getManager(), scope.getResolveScope());
+        result = result == null ? defType : TypesUtil.getLeastUpperBound(result, defType, scope.getManager());
+      }
+    }
+    return result;
+  }
 
   private static Pair<ReachingDefinitionsDfaInstance, List<TIntObjectHashMap<TIntHashSet>>> getDefUseMaps(final GrControlFlowOwner scope) {
     return CachedValuesManager.getManager(scope.getProject()).getCachedValue(scope, new CachedValueProvider<Pair<ReachingDefinitionsDfaInstance, List<TIntObjectHashMap<TIntHashSet>>>>() {
@@ -156,7 +196,7 @@ public class TypeInferenceHelper {
 
 
   @Nullable
-  public static ReadWriteVariableInstruction findInstruction(final GrReferenceExpression refExpr, final Instruction[] flow) {
+  private static ReadWriteVariableInstruction findInstruction(final GrReferenceExpression refExpr, final Instruction[] flow) {
     for (Instruction instruction : flow) {
       if (instruction instanceof ReadWriteVariableInstruction && instruction.getElement() == refExpr) {
         return (ReadWriteVariableInstruction)instruction;
@@ -164,7 +204,7 @@ public class TypeInferenceHelper {
     }
     return null;
   }
-
+  
   @Nullable
   public static PsiType getInitializerType(final PsiElement element) {
     if (element instanceof GrReferenceExpression && ((GrReferenceExpression) element).getQualifierExpression() == null) {

plugins/groovy/src/org/jetbrains/plugins/groovy/template/GroovyVariableTypeCalculator.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2000-2011 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.plugins.groovy.template;
+
+import com.intellij.codeInsight.template.macro.VariableTypeCalculator;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiType;
+import com.intellij.psi.PsiVariable;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.plugins.groovy.GroovyFileType;
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrField;
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrVariable;
+import org.jetbrains.plugins.groovy.lang.psi.impl.TypeInferenceHelper;
+
+/**
+ * @author Max Medvedev
+ */
+public class GroovyVariableTypeCalculator extends VariableTypeCalculator {
+  @Override
+  public PsiType inferVarTypeAt(@NotNull PsiVariable var, @NotNull PsiElement place) {
+    if (!(var instanceof GrVariable) || !(place.getLanguage() == GroovyFileType.GROOVY_LANGUAGE)) return null;
+    if (var instanceof GrField) return var.getType();
+
+    return TypeInferenceHelper.getInferredType(place, var.getName());
+  }
+}