fix(codewhisperer): show full reference in reference log (#3326)

Author: zixlin7

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceManager.kt

@@ -23,6 +23,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.wm.ToolWindowManager
 import com.intellij.ui.awt.RelativePoint
+import software.amazon.awssdk.services.codewhisperer.model.Recommendation
 import software.amazon.awssdk.services.codewhisperer.model.Reference
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.getPopupPositionAboveText
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.getRelativePathToContentRoot
@@ -66,34 +67,17 @@ class CodeWhispererCodeReferenceManager(private val project: Project) {
         val (requestContext, _, recommendationContext) = states
         val (_, editor, _, caretPosition) = requestContext
         val (_, detail, reformattedDetail) = recommendationContext.details[selectedIndex]
-        val userInput = recommendationContext.userInputSinceInvocation
         val startOffset = caretPosition.offset
         val relativePath = getRelativePathToContentRoot(editor)
         reformattedDetail.references().forEachIndexed { i, reference ->
-            // start and end could intersect with the userInput, we do not want to show reference for the
-            // userInput part, so we truncate the range to exclude userInput here if there's an overlap.
-            if (reference.recommendationContentSpan().end() <= userInput.length) return@forEachIndexed
-
-            // Now there's at least some valid range, truncate start if applicable
-            val spanStartExcludeUserInput = maxOf(reference.recommendationContentSpan().start(), userInput.length)
-            val start = startOffset + spanStartExcludeUserInput
+            val start = startOffset + reference.recommendationContentSpan().start()
             val end = startOffset + reference.recommendationContentSpan().end()
-            val startLine = editor.document.getLineNumber(start)
-            val endLine = editor.document.getLineNumber(end)
-            val lineNums = if (startLine == endLine) {
-                (startLine + 1).toString()
-            } else {
-                "${startLine + 1} to ${endLine + 1}"
-            }
-
-            val originalSpan = detail.references()[i].recommendationContentSpan()
+            val lineNums = getReferenceLineNums(editor, start, end)
 
             // There is an unformatted recommendation(directly from response) and reformatted one. We want to get
             // the line number, start/end offset of the reformatted one because it's the one inserted to the editor.
             // However, the one that shows in the tool window record should show the original recommendation, as below.
-            val originalContentLines = detail.content().substring(
-                maxOf(originalSpan.start(), userInput.length), originalSpan.end()
-            ).split("\n")
+            val originalContentLines = getOriginalContentLines(detail, i)
 
             codeReferenceComponents.contentPanel.apply {
                 add(
@@ -115,6 +99,24 @@ class CodeWhispererCodeReferenceManager(private val project: Project) {
         }
     }
 
+    fun getReferenceLineNums(editor: Editor, start: Int, end: Int): String {
+        val startLine = editor.document.getLineNumber(start)
+        val endLine = editor.document.getLineNumber(end)
+        val lineNums = if (startLine == endLine) {
+            (startLine + 1).toString()
+        } else {
+            "${startLine + 1} to ${endLine + 1}"
+        }
+        return lineNums
+    }
+
+    fun getOriginalContentLines(detail: Recommendation, i: Int): List<String> {
+        val originalSpan = detail.references()[i].recommendationContentSpan()
+        return detail.content()
+            .substring(originalSpan.start(), originalSpan.end())
+            .split("\n")
+    }
+
     private fun insertHighLightContext(editor: Editor, start: Int, end: Int, reference: Reference) {
         val codeContent = editor.document.getText(TextRange.create(start, end))
         val referenceContent = message(

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererReferenceManagerTest.kt

@@ -0,0 +1,69 @@
+// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.codewhisperer
+
+import com.intellij.openapi.project.Project
+import com.intellij.testFramework.DisposableRule
+import com.intellij.testFramework.fixtures.CodeInsightTestFixture
+import com.intellij.testFramework.runInEdtAndWait
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import software.amazon.awssdk.services.codewhisperer.model.Recommendation
+import software.amazon.awssdk.services.codewhisperer.model.Reference
+import software.amazon.awssdk.services.codewhisperer.model.Span
+import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
+import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
+
+class CodeWhispererReferenceManagerTest {
+    @Rule
+    @JvmField
+    var projectRule = PythonCodeInsightTestFixtureRule()
+
+    @Rule
+    @JvmField
+    val disposableRule = DisposableRule()
+
+    private val documentContentContent = "012345678\n9"
+    private lateinit var fixture: CodeInsightTestFixture
+    private lateinit var project: Project
+    private val originalReference = Reference.builder()
+        .licenseName("test_license")
+        .repository("test_repo")
+        .recommendationContentSpan(
+            Span.builder().start(0).end(14).build()
+        )
+        .build()
+
+    private val recommendation = Recommendation.builder()
+        .references(originalReference)
+        .content("test\nreference")
+        .build()
+
+    @Before
+    fun setup() {
+        fixture = projectRule.fixture
+        project = projectRule.project
+
+        fixture.configureByText("test.py", documentContentContent)
+        runInEdtAndWait {
+            fixture.editor.caretModel.moveToOffset(documentContentContent.length)
+        }
+    }
+
+    @Test
+    fun `test getReferenceLineNums return expected line numbers`() {
+        val referenceManager = CodeWhispererCodeReferenceManager(project)
+        assertThat(referenceManager.getReferenceLineNums(fixture.editor, 0, 1)).isEqualTo("1")
+        assertThat(referenceManager.getReferenceLineNums(fixture.editor, 0, 10)).isEqualTo("1 to 2")
+    }
+
+    @Test
+    fun `test getOriginalContent lines returns full reference lines`() {
+        val referenceManager = CodeWhispererCodeReferenceManager(project)
+        val expectedRecommendation = listOf<String>("test", "reference")
+        assertThat(referenceManager.getOriginalContentLines(recommendation, 0)).isEqualTo(expectedRecommendation)
+    }
+}