fix(devin-lang): Improve code block parsing and injection logic #101

Author: phodal

exts/devin-lang/src/main/kotlin/cc/unitmesh/language/DevInLanguageInjector.kt

@@ -25,16 +25,16 @@ class DevInLanguageInjector : LanguageInjector {
         val text = languageIdentifier?.text ?: return
         val language = CodeFenceLanguageGuesser.guessLanguageForInjection(text) ?: return
 
-        val elements = host.getContents()
-        if (elements.size < 2) {
+        val contentList = CodeBlockElement.obtainFenceContent(host) ?: return
+        if (contentList.isEmpty()) {
             return
         }
 
         injectAsOnePlace(host, language, registrar)
     }
 
     private fun injectAsOnePlace(host: CodeBlockElement, language: Language, registrar: InjectedLanguagePlaces) {
-        val elements = CodeBlockElement.obtainFenceContent(host, withWhitespaces = true) ?: return
+        val elements = CodeBlockElement.obtainFenceContent(host) ?: return
 
         val first = elements.first()
         val last = elements.last()

exts/devin-lang/src/main/kotlin/cc/unitmesh/language/parser/CodeBlockElement.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.LiteralTextEscaper
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiLanguageInjectionHost
 import com.intellij.psi.impl.source.tree.injected.InjectionBackgroundSuppressor
+import com.intellij.psi.templateLanguages.OuterLanguageElement
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.util.*
 
@@ -32,36 +33,43 @@ class CodeBlockElement(node: ASTNode) : ASTWrapperPsiElement(node), PsiLanguageI
         return findChildByType(DevInTypes.LANGUAGE_ID)
     }
 
-    fun getContents(): List<PsiElement> {
-        return obtainFenceContent(this, withWhitespaces = false) ?: return emptyList()
+    fun getContents(): PsiElement? {
+        return findChildByType(DevInTypes.CODE_CONTENTS)
     }
 
     companion object {
-        fun obtainFenceContent(element: CodeBlockElement, withWhitespaces: Boolean): List<PsiElement>? {
-            return when {
-                withWhitespaces -> CachedValuesManager.getCachedValue(element) {
-                    CachedValueProvider.Result.create(getContent(element, true), element)
-                }
-                else -> CachedValuesManager.getCachedValue(element) {
-                    CachedValueProvider.Result.create(getContent(element, false), element)
-                }
+        fun obtainFenceContent(element: CodeBlockElement): List<PsiElement>? {
+            return CachedValuesManager.getCachedValue(element) {
+                CachedValueProvider.Result.create(getContent(element), element)
             }
         }
 
-        private fun getContent(element: CodeBlockElement, withWhitespaces: Boolean): List<PsiElement>? {
-            val codeContents = element.children.firstOrNull { it.elementType == DevInTypes.CODE_CONTENTS } ?: return null
+        private fun getContent(host: CodeBlockElement): List<PsiElement>? {
+            val children = host.firstChild
+                ?.siblings(forward = true, withSelf = true) ?: return null
+
+            val elements =
+                children.filter {
+                    it !is OuterLanguageElement
+                            && (it.node.elementType == DevInTypes.CODE_CONTENTS || it == DevInTypes.NEWLINE)
+                }
+                    .toList()
 
-            val psiElements = PsiTreeUtil.collectElements(codeContents) {
-                it.elementType == DevInTypes.CODE_CONTENT
-            }.toMutableList()
+            if (elements.isNotEmpty() && elements.first() == DevInTypes.NEWLINE) {
+                elements.drop(1)
+            }
+            if (elements.isNotEmpty() && elements.last() == DevInTypes.NEWLINE) {
+                elements.dropLast(1)
+            }
 
-            return psiElements.toList()
+            return elements.takeIf { it.isNotEmpty() }
         }
 
         fun obtainRelevantTextRange(element: CodeBlockElement): TextRange {
-            val elements = obtainFenceContent(element, withWhitespaces = true) ?: return getEmptyRange(element)
+            val elements = obtainFenceContent(element) ?: return getEmptyRange(element)
             val first = elements.first()
             val last = elements.last()
+
             return TextRange.create(first.startOffsetInParent, last.startOffsetInParent + last.textLength)
         }
 
@@ -82,27 +90,26 @@ class CodeBlockLiteralTextEscaper(host: CodeBlockElement) : LiteralTextEscaper<C
     override fun isOneLine(): Boolean = false;
 
     override fun decode(rangeInsideHost: TextRange, outChars: StringBuilder): Boolean {
-        val elements = CodeBlockElement.obtainFenceContent(myHost, withWhitespaces = false) ?: return true
+        val elements = CodeBlockElement.obtainFenceContent(myHost) ?: return true
         for (element in elements) {
             val intersected = rangeInsideHost.intersection(element.textRangeInParent) ?: continue
             outChars.append(intersected.substring(myHost.text))
         }
+
         return true
     }
 
     override fun getOffsetInHost(offsetInDecoded: Int, rangeInsideHost: TextRange): Int {
-        val elements = CodeBlockElement.obtainFenceContent(myHost, withWhitespaces = false) ?: return -1
+        val elements = CodeBlockElement.obtainFenceContent(myHost) ?: return -1
         var cur = 0
         for (element in elements) {
             val intersected = rangeInsideHost.intersection(element.textRangeInParent)
             if (intersected == null || intersected.isEmpty) continue
             if (cur + intersected.length == offsetInDecoded) {
                 return intersected.startOffset + intersected.length
-            }
-            else if (cur == offsetInDecoded) {
+            } else if (cur == offsetInDecoded) {
                 return intersected.startOffset
-            }
-            else if (cur < offsetInDecoded && cur + intersected.length > offsetInDecoded) {
+            } else if (cur < offsetInDecoded && cur + intersected.length > offsetInDecoded) {
                 return intersected.startOffset + (offsetInDecoded - cur)
             }
             cur += intersected.length
@@ -114,8 +121,7 @@ class CodeBlockLiteralTextEscaper(host: CodeBlockElement) : LiteralTextEscaper<C
         val result = intersected.startOffset + (offsetInDecoded - (cur - intersected.length))
         return if (rangeInsideHost.startOffset <= result && result <= rangeInsideHost.endOffset) {
             result
-        }
-        else -1
+        } else -1
     }
 
     override fun getRelevantTextRange(): TextRange {

exts/devin-lang/src/main/kotlin/cc/unitmesh/language/parser/DevInParserDefinition.kt

@@ -4,6 +4,7 @@ import cc.unitmesh.language.DevInLanguage
 import cc.unitmesh.language.lexer.DevInLexerAdapter
 import cc.unitmesh.language.psi.DevInFile
 import cc.unitmesh.language.psi.DevInTypes
+import com.intellij.extapi.psi.ASTWrapperPsiElement
 import com.intellij.lang.ASTNode
 import com.intellij.lang.ParserDefinition
 import com.intellij.lang.PsiParser
@@ -19,9 +20,7 @@ import org.jetbrains.annotations.NotNull
 
 internal class DevInParserDefinition : ParserDefinition {
     @NotNull
-    override fun createLexer(project: Project?): Lexer {
-        return DevInLexerAdapter()
-    }
+    override fun createLexer(project: Project?): Lexer = DevInLexerAdapter()
 
     @NotNull
     override fun getCommentTokens(): TokenSet = TokenSet.EMPTY
@@ -30,19 +29,13 @@ internal class DevInParserDefinition : ParserDefinition {
     override fun getStringLiteralElements(): TokenSet = TokenSet.EMPTY
 
     @NotNull
-    override fun createParser(project: Project?): PsiParser {
-        return DevInParser()
-    }
+    override fun createParser(project: Project?): PsiParser = DevInParser()
 
     @NotNull
-    override fun getFileNodeType(): IFileElementType {
-        return FILE
-    }
+    override fun getFileNodeType(): IFileElementType = FILE
 
     @NotNull
-    override fun createFile(@NotNull viewProvider: FileViewProvider): PsiFile {
-        return DevInFile(viewProvider)
-    }
+    override fun createFile(@NotNull viewProvider: FileViewProvider): PsiFile = DevInFile(viewProvider)
 
     @NotNull
     override fun createElement(node: ASTNode?): PsiElement {
@@ -51,6 +44,10 @@ internal class DevInParserDefinition : ParserDefinition {
             return CodeBlockElement(node)
         }
 
+        if (elementType == DevInTypes.CODE_CONTENTS) {
+            return ASTWrapperPsiElement(node)
+        }
+
         return DevInTypes.Factory.createElement(node)
     }