feat(custom): fix new line issues

unit-mesh · Aug 9, 2023 · 6da2821 · 6da2821
1 parent 6758260
commit 6da2821
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/context/JavaMethodContextBuilder.kt b/java/src/main/kotlin/cc/unitmesh/idea/context/JavaMethodContextBuilder.kt
@@ -1,12 +1,11 @@
 package cc.unitmesh.idea.context
 
-import cc.unitmesh.devti.context.builder.MethodContextBuilder
 import cc.unitmesh.devti.context.MethodContext
+import cc.unitmesh.devti.context.builder.MethodContextBuilder
 import cc.unitmesh.idea.service.JavaTypeUtil
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiNameIdentifierOwner
-import com.intellij.psi.impl.source.PsiClassReferenceType
 
 class JavaMethodContextBuilder : MethodContextBuilder {
     override fun getMethodContext(

diff --git a/java/src/main/kotlin/cc/unitmesh/idea/service/JavaTypeUtil.kt b/java/src/main/kotlin/cc/unitmesh/idea/service/JavaTypeUtil.kt
@@ -35,6 +35,11 @@ object JavaTypeUtil {
         return resolvedClasses
     }
 
+    /**
+     * The resolved classes include all the classes in the method signature. For example, if the method signature is
+     * Int, will return Int, but if the method signature is List<Int>, will return List and Int.
+     * So, remember to filter out the classes that are not needed.
+     */
     fun resolveByMethod(element: PsiElement): MutableMap<String, PsiClass?> {
         val resolvedClasses = mutableMapOf<String, PsiClass?>()
         if (element is PsiMethod) {

diff --git a/src/main/kotlin/cc/unitmesh/devti/custom/variable/MethodInputOutputVariableResolver.kt b/src/main/kotlin/cc/unitmesh/devti/custom/variable/MethodInputOutputVariableResolver.kt
@@ -3,12 +3,14 @@ package cc.unitmesh.devti.custom.variable
 import cc.unitmesh.devti.context.ClassContextProvider
 import cc.unitmesh.devti.context.MethodContextProvider
 import com.intellij.lang.LanguageCommenters
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 
 class MethodInputOutputVariableResolver(val element: PsiElement) : VariableResolver {
     override val type: CustomIntentionVariableType = CustomIntentionVariableType.METHOD_INPUT_OUTPUT
     private val commenter = LanguageCommenters.INSTANCE.forLanguage(element.language) ?: null
-    val commentPrefix = commenter?.lineCommentPrefix ?: ""
+    private val commentPrefix = commenter?.lineCommentPrefix ?: ""
+    private val project: Project = element.project
 
     override fun resolve(): String {
         var result = ""
@@ -18,10 +20,19 @@ class MethodInputOutputVariableResolver(val element: PsiElement) : VariableResol
         }
 
         methodContext.inputOutputClasses.forEach {
-            ClassContextProvider(false).from(it).let { classContext ->
+            val context = ClassContextProvider(false).from(it)
+            val element = context.root
+            val basePath = project.basePath ?: return@forEach
+
+            if (element.containingFile?.virtualFile?.path?.contains(basePath) != true) {
+                return@forEach
+            }
+
+            context.let { classContext ->
                 result += classContext.toQuery().lines().joinToString(separator = "\n") { line ->
-                    "$commentPrefix  $line"
+                    "$commentPrefix $line"
                 }
+                result += "\n"
             }
         }
 

diff --git a/src/main/kotlin/com/intellij/temporary/similar/chunks/SimilarChunksWithPaths.kt b/src/main/kotlin/com/intellij/temporary/similar/chunks/SimilarChunksWithPaths.kt
@@ -1,6 +1,7 @@
 // Copyright 2000-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.temporary.similar.chunks
 
+import cc.unitmesh.devti.llms.tokenizer.TokenizerImpl
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.fileEditor.impl.EditorHistoryManager
 import com.intellij.openapi.fileTypes.FileType
@@ -13,6 +14,8 @@ import com.intellij.psi.PsiManager
 import java.io.File
 
 class SimilarChunksWithPaths(private var chunkSize: Int = 60, private var maxRelevantFiles: Int = 20) {
+    private val tokenizer = TokenizerImpl.INSTANCE
+
     companion object {
         val INSTANCE: SimilarChunksWithPaths = SimilarChunksWithPaths()
 
@@ -56,10 +59,10 @@ class SimilarChunksWithPaths(private var chunkSize: Int = 60, private var maxRel
     }
 
     private fun tokenLevelJaccardSimilarity(chunks: List<List<String>>, element: PsiElement): List<List<Double>> {
-        val currentFileTokens: Set<String> = tokenize(element.containingFile.text).toSet()
+        val currentFileTokens: List<Int> = tokenize(element.containingFile.text)
         return chunks.map { list ->
             list.map {
-                val tokenizedFile: Set<String> = tokenize(it).toSet()
+                val tokenizedFile: List<Int> = tokenize(it)
                 jaccardSimilarity(currentFileTokens, tokenizedFile)
             }
         }
@@ -78,14 +81,13 @@ class SimilarChunksWithPaths(private var chunkSize: Int = 60, private var maxRel
         return contentRoot?.let { VfsUtilCore.getRelativePath(relativeFile, it, File.separatorChar) }
     }
 
-    private fun tokenize(chunk: String): List<String> {
-        return chunk.split(Regex("[^a-zA-Z0-9]"))
-            .filter { it.isNotBlank() }
+    private fun tokenize(chunk: String): List<Int> {
+        return tokenizer.tokenize(chunk)
     }
 
-    private fun jaccardSimilarity(set1: Set<String>, set2: Set<String>): Double {
-        val intersectionSize: Int = set1.intersect(set2).size
-        val unionSize: Int = set1.union(set2).size
+    private fun jaccardSimilarity(left: List<Int>, right: List<Int>): Double {
+        val intersectionSize: Int = left.intersect(right.toSet()).size
+        val unionSize: Int = left.union(right).size
         return intersectionSize.toDouble() / unionSize.toDouble()
     }