diff --git a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuggestVariableNameMacro.kt b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuggestVariableNameMacro.kt index c1f10391851..b9b4674ce31 100644 --- a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuggestVariableNameMacro.kt +++ b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuggestVariableNameMacro.kt @@ -16,15 +16,73 @@ package org.jetbrains.kotlin.idea.liveTemplates.macro -import com.intellij.openapi.project.Project -import com.intellij.openapi.util.UserDataHolder +import com.intellij.codeInsight.lookup.LookupElement +import com.intellij.codeInsight.lookup.LookupElementBuilder +import com.intellij.codeInsight.template.* +import com.intellij.psi.PsiDocumentManager import org.jetbrains.kotlin.descriptors.VariableDescriptor +import org.jetbrains.kotlin.idea.caches.resolve.analyze +import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade +import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor +import org.jetbrains.kotlin.idea.core.IterableTypesDetection +import org.jetbrains.kotlin.idea.core.KotlinNameSuggester +import org.jetbrains.kotlin.idea.resolve.ideService +import org.jetbrains.kotlin.idea.util.getResolutionScope +import org.jetbrains.kotlin.psi.KtCallableDeclaration +import org.jetbrains.kotlin.psi.KtFile +import org.jetbrains.kotlin.psi.KtForExpression +import org.jetbrains.kotlin.psi.KtWithExpressionInitializer +import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode -class SuggestVariableNameMacro : BaseKotlinVariableMacro() { +class SuggestVariableNameMacro : Macro() { override fun getName() = "kotlinSuggestVariableName" override fun getPresentableName() = "kotlinSuggestVariableName()" - override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder): Boolean { - return variableDescriptor.type.isMarkedNullable + override fun calculateResult(params: Array, context: ExpressionContext): Result? { + return nameSuggestions(context).firstOrNull()?.let { TextResult(it) } + } + + override fun calculateLookupItems(params: Array, context: ExpressionContext): Array? { + val suggestions = nameSuggestions(context) + if (suggestions.size < 2) return null + return suggestions.map { LookupElementBuilder.create(it) }.toTypedArray() + } + + private fun nameSuggestions(context: ExpressionContext): Collection { + val project = context.project + val psiDocumentManager = PsiDocumentManager.getInstance(project) + psiDocumentManager.commitAllDocuments() + + val psiFile = psiDocumentManager.getPsiFile(context.editor!!.document) as? KtFile ?: return emptyList() + val token = psiFile.findElementAt(context.startOffset) ?: return emptyList() + val declaration = token.parent as? KtCallableDeclaration ?: return emptyList() + if (token != declaration.nameIdentifier) return emptyList() + + val nameValidator: (String) -> Boolean = { true } + + val initializer = (declaration as? KtWithExpressionInitializer)?.initializer + if (initializer != null) { + val bindingContext = initializer.analyze(BodyResolveMode.PARTIAL) + return KotlinNameSuggester.suggestNamesByExpressionAndType(initializer, bindingContext, nameValidator, null) + } + + val parent = declaration.parent + if (parent is KtForExpression && declaration == parent.loopParameter) { + iterationVariableNameSuggestions(parent, nameValidator)?.let { return it } + } + + val descriptor = declaration.resolveToDescriptor() as? VariableDescriptor ?: return emptyList() + return KotlinNameSuggester.suggestNamesByType(descriptor.type, nameValidator, null) + } + + private fun iterationVariableNameSuggestions(forExpression: KtForExpression, nameValidator: (String) -> Boolean): Collection? { + val loopRange = forExpression.loopRange ?: return null + val resolutionFacade = forExpression.getResolutionFacade() + val bindingContext = resolutionFacade.analyze(loopRange, BodyResolveMode.PARTIAL) + val type = bindingContext.getType(loopRange) ?: return null + val scope = loopRange.getResolutionScope(bindingContext, resolutionFacade) + val detector = resolutionFacade.ideService().createDetector(scope) + val elementType = detector.elementType(type)?.type ?: return null + return KotlinNameSuggester.suggestIterationVariableNames(loopRange, elementType, bindingContext, nameValidator, null) } } diff --git a/idea/testData/templates/iter.exp.kt b/idea/testData/templates/iter.exp.kt index d3848f4b245..21552c04c9e 100644 --- a/idea/testData/templates/iter.exp.kt +++ b/idea/testData/templates/iter.exp.kt @@ -12,7 +12,7 @@ class MyClass { val myList = ArrayList() val stream = FileInputStream(".") if (o is String) { - for (collection in args) { + for (arg in args) { } }