From 970bd46544ceba7f6023f3f77d873c9e7a688eb1 Mon Sep 17 00:00:00 2001 From: Ilya Ryzhenkov Date: Wed, 2 Sep 2015 00:48:01 +0300 Subject: [PATCH] Optimize siblings() --- .../jetbrains/kotlin/psi/psiUtil/psiUtils.kt | 60 ++++++++++++------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt index 808d5dc3850..2b56e048568 100644 --- a/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt +++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt @@ -22,7 +22,7 @@ import com.intellij.psi.search.PsiSearchScopeUtil import com.intellij.psi.search.SearchScope import com.intellij.psi.util.PsiTreeUtil import org.jetbrains.kotlin.diagnostics.DiagnosticUtils -import java.util.ArrayList +import java.util.* // NOTE: in this file we collect only LANGUAGE INDEPENDENT methods working with PSI and not modifying it @@ -33,10 +33,28 @@ public val PsiElement.allChildren: PsiChildRange val first = getFirstChild() return if (first != null) PsiChildRange(first, getLastChild()) else PsiChildRange.EMPTY } + public fun PsiElement.siblings(forward: Boolean = true, withItself: Boolean = true): Sequence { - val stepFun = if (forward) { e: PsiElement -> e.getNextSibling() } else { e: PsiElement -> e.getPrevSibling() } - val sequence = sequence(this, stepFun) - return if (withItself) sequence else sequence.drop(1) + return object : Sequence { + override fun iterator(): Iterator { + var next: PsiElement? = this@siblings + return object : Iterator { + init { + if (!withItself) next() + } + + override fun hasNext(): Boolean = next != null + override fun next(): PsiElement { + val result = next ?: throw NoSuchElementException() + if (forward) + next = result.nextSibling + else + next = result.prevSibling + return result + } + } + } + } } public val PsiElement.parentsWithSelf: Sequence @@ -67,14 +85,14 @@ public fun PsiElement.nextLeaf(filter: (PsiElement) -> Boolean): PsiElement? { return leaf } -public fun PsiElement.getParentOfTypesAndPredicate( - strict : Boolean = false, vararg parentClasses : Class, predicate: (T) -> Boolean -) : T? { +public fun PsiElement.getParentOfTypesAndPredicate( + strict: Boolean = false, vararg parentClasses: Class, predicate: (T) -> Boolean +): T? { var element = if (strict) getParent() else this while (element != null) { @suppress("UNCHECKED_CAST") when { - (parentClasses.isEmpty() || parentClasses.any {parentClass -> parentClass.isInstance(element)}) && predicate(element!! as T) -> + (parentClasses.isEmpty() || parentClasses.any { parentClass -> parentClass.isInstance(element) }) && predicate(element!! as T) -> return element as T element is PsiFile -> return null @@ -86,27 +104,27 @@ public fun PsiElement.getParentOfTypesAndPredicate( return null } -public fun PsiElement.getNonStrictParentOfType(parentClass : Class) : T? { +public fun PsiElement.getNonStrictParentOfType(parentClass: Class): T? { return PsiTreeUtil.getParentOfType(this, parentClass, false) } -inline public fun PsiElement.getParentOfType(strict: Boolean): T? { +inline public fun PsiElement.getParentOfType(strict: Boolean): T? { return PsiTreeUtil.getParentOfType(this, javaClass(), strict) } -inline public fun PsiElement.getStrictParentOfType(): T? { +inline public fun PsiElement.getStrictParentOfType(): T? { return PsiTreeUtil.getParentOfType(this, javaClass(), true) } -inline public fun PsiElement.getNonStrictParentOfType(): T? { +inline public fun PsiElement.getNonStrictParentOfType(): T? { return PsiTreeUtil.getParentOfType(this, javaClass(), false) } -inline public fun PsiElement.getChildOfType(): T? { +inline public fun PsiElement.getChildOfType(): T? { return PsiTreeUtil.getChildOfType(this, javaClass()) } -inline public fun PsiElement.getChildrenOfType(): Array { +inline public fun PsiElement.getChildrenOfType(): Array { return PsiTreeUtil.getChildrenOfType(this, javaClass()) ?: arrayOf() } @@ -118,11 +136,11 @@ public fun PsiElement?.isAncestor(element: PsiElement, strict: Boolean = false): return PsiTreeUtil.isAncestor(this, element, strict) } -public fun T.getIfChildIsInBranch(element: PsiElement, branch: T.() -> PsiElement?): T? { +public fun T.getIfChildIsInBranch(element: PsiElement, branch: T.() -> PsiElement?): T? { return if (branch().isAncestor(element)) this else null } -public inline fun PsiElement.getParentOfTypeAndBranch(strict: Boolean = false, noinline branch: T.() -> PsiElement?) : T? { +public inline fun PsiElement.getParentOfTypeAndBranch(strict: Boolean = false, noinline branch: T.() -> PsiElement?): T? { return getParentOfType(strict)?.getIfChildIsInBranch(this, branch) } @@ -136,11 +154,11 @@ public fun PsiElement.isInsideOf(elements: Iterable): Boolean = elem // -------------------- Recursive tree visiting -------------------------------------------------------------------------------------------- public inline fun PsiElement.forEachDescendantOfType(noinline action: (T) -> Unit) { - forEachDescendantOfType({true}, action) + forEachDescendantOfType({ true }, action) } public inline fun PsiElement.forEachDescendantOfType(inlineOptions(InlineOption.ONLY_LOCAL_RETURN) canGoInside: (PsiElement) -> Boolean, noinline action: (T) -> Unit) { - this.accept(object : PsiRecursiveElementVisitor(){ + this.accept(object : PsiRecursiveElementVisitor() { override fun visitElement(element: PsiElement) { if (canGoInside(element)) { super.visitElement(element) @@ -162,12 +180,12 @@ public inline fun PsiElement.anyDescendantOfType(inline } public inline fun PsiElement.findDescendantOfType(noinline predicate: (T) -> Boolean = { true }): T? { - return findDescendantOfType({ true}, predicate) + return findDescendantOfType({ true }, predicate) } public inline fun PsiElement.findDescendantOfType(inlineOptions(InlineOption.ONLY_LOCAL_RETURN) canGoInside: (PsiElement) -> Boolean, noinline predicate: (T) -> Boolean = { true }): T? { var result: T? = null - this.accept(object : PsiRecursiveElementVisitor(){ + this.accept(object : PsiRecursiveElementVisitor() { override fun visitElement(element: PsiElement) { if (result != null) return @@ -291,5 +309,5 @@ public fun PsiElement.getTextWithLocation(): String = "'${this.getText()}' at ${ public fun SearchScope.contains(element: PsiElement): Boolean = PsiSearchScopeUtil.isInScope(this, element) -public fun E.createSmartPointer(): SmartPsiElementPointer = +public fun E.createSmartPointer(): SmartPsiElementPointer = SmartPointerManager.getInstance(getProject()).createSmartPsiElementPointer(this) \ No newline at end of file