Simplify anchor determination

intellij-rust · Jul 6, 2019 · a6cc04a · a6cc04a
1 parent c1c59cc
commit a6cc04a
Showing 1 changed file with 16 additions and 34 deletions.
diff --git a/src/main/kotlin/org/rust/ide/utils/StructFieldsExpander.kt b/src/main/kotlin/org/rust/ide/utils/StructFieldsExpander.kt
@@ -16,20 +16,7 @@ fun expand(factory: RsPsiFactory, rsPat: RsPat, amount: Int, hasTrailingComma: B
             val declaration = rsPat.path.reference.deepResolve() as? RsFieldsOwner ?: return
             val bodyFields = rsPat.patFieldList
             val missingFields = calculateMissingFields(bodyFields.map { it.text }.toSet(), declaration)
-            var comma = rsPat.childrenOfType<LeafPsiElement>().lastOrNull { it.elementType == RsElementTypes.COMMA }
-            var anchor = comma
-
-            if (comma == null) {
-                anchor = rsPat.lbrace as LeafPsiElement
-                comma = factory.createComma() as LeafPsiElement
-            }
-            if (comma.nextSibling.elementType == RsElementTypes.PAT_FIELD) {
-                rsPat.addAfter(comma, comma.nextSibling)
-                anchor = comma.nextSibling.nextSibling as LeafPsiElement
-            } else if (comma.nextSibling?.nextSibling?.elementType == RsElementTypes.PAT_FIELD) {
-                rsPat.addAfter(comma, comma.nextSibling.nextSibling)
-                anchor = comma.nextSibling.nextSibling.nextSibling as LeafPsiElement
-            }
+            var anchor = determineOrCreateAnchor(factory, rsPat)
             var amount = amount
             if (amount == -1) {
                 amount = missingFields.size
@@ -39,7 +26,7 @@ fun expand(factory: RsPsiFactory, rsPat: RsPat, amount: Int, hasTrailingComma: B
                 .take(amount)
                 .forEach { missingField ->
                     rsPat.addAfter(missingField, anchor)
-                    rsPat.addAfter(comma, anchor!!.nextSibling)
+                    rsPat.addAfter(factory.createComma(), anchor!!.nextSibling)
                     anchor = anchor!!.nextSibling.nextSibling as LeafPsiElement
                 }
             if (!hasTrailingComma) {
@@ -48,25 +35,7 @@ fun expand(factory: RsPsiFactory, rsPat: RsPat, amount: Int, hasTrailingComma: B
         }
         is RsPatTupleStruct -> {
             var tempPlaceholder = 'a'
-            var comma1 = rsPat.childrenOfType<LeafPsiElement>().lastOrNull { it.elementType == RsElementTypes.COMMA }
-            var anchor = comma1
-
-            if (comma1 == null) {
-                comma1 = factory.createComma() as LeafPsiElement
-                anchor = if (rsPat.lparen.getNextNonCommentSibling()?.elementType == RsElementTypes.PAT_IDENT) {
-                    rsPat.addAfter(comma1, rsPat.lparen.getNextNonCommentSibling())
-                    rsPat.lparen.getNextNonCommentSibling()?.nextSibling as LeafPsiElement
-                } else {
-                    rsPat.lparen as LeafPsiElement
-                }
-            }
-            if (comma1.nextSibling.elementType == RsElementTypes.PAT_IDENT) {
-                rsPat.addAfter(comma1, comma1.nextSibling)
-                anchor = comma1.nextSibling.nextSibling as LeafPsiElement
-            } else if (comma1.nextSibling?.nextSibling?.elementType == RsElementTypes.PAT_IDENT) {
-                rsPat.addAfter(comma1, comma1.nextSibling.nextSibling)
-                anchor = comma1.nextSibling.nextSibling.nextSibling as LeafPsiElement
-            }
+            var anchor = determineOrCreateAnchor(factory, rsPat)
             for (i in 0 until amount) {
                 rsPat.addAfter(factory.createPatBinding(tempPlaceholder.toString()), anchor)
                 rsPat.addAfter(factory.createComma(), anchor!!.nextSibling)
@@ -80,5 +49,18 @@ fun expand(factory: RsPsiFactory, rsPat: RsPat, amount: Int, hasTrailingComma: B
     }
 }
 
+private fun determineOrCreateAnchor(factory: RsPsiFactory, rsPatBody: RsPat) : LeafPsiElement? {
+    var lastElementInBody = rsPatBody.lastChild.getPrevNonCommentSibling()
+    if (lastElementInBody?.elementType == RsElementTypes.DOTDOT) {
+        lastElementInBody = lastElementInBody.getPrevNonCommentSibling()
+    }
+    return if (lastElementInBody !is LeafPsiElement) {
+        rsPatBody.addAfter(factory.createComma(), lastElementInBody)
+        lastElementInBody?.getNextNonCommentSibling() as LeafPsiElement
+    } else {
+        lastElementInBody
+    }
+}
+
 private fun calculateMissingFields(bodyFields: Set<String>, declaration: RsFieldsOwner)
     = declaration.fields.filter { it.name !in bodyFields && !it.queryAttributes.hasCfgAttr() }