[aia-eval] LME-92 Optionally dump a whole dataset in a single file in Evaluation plugin

GitOrigin-RevId: e959b9c3f3dcb2ced4c7e860b9f585d58b2356ae

plugins/evaluation-plugin/core/src/com/intellij/cce/actions/ActionSerializer.kt

@@ -19,4 +19,16 @@ object ActionSerializer {
   fun getSessionsCount(json: String): Int {
     return gson.fromJson(json, FakeFileActions::class.java).sessionsCount
   }
+}
+
+object ActionArraySerializer {
+  private val gson = GsonBuilder()
+    .registerTypeAdapter(Action::class.java, Action.JsonAdapter)
+    .registerTypeAdapter(TokenProperties::class.java, TokenProperties.JsonAdapter)
+    .setPrettyPrinting()
+    .create()
+
+  fun serialize(actions: Array<FileActions>): String = gson.toJson(actions)
+
+  fun deserialize(json: String): Array<FileActions> = gson.fromJson(json, Array<FileActions>::class.java)
 }

plugins/evaluation-plugin/core/src/com/intellij/cce/core/TokenProperties.kt

@@ -134,6 +134,18 @@ class SimpleTokenProperties private constructor(
 
   override fun withFeatures(features: Set<String>): TokenProperties =
     SimpleTokenProperties(tokenType, location, this.features.apply { addAll(features) }, additional)
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is SimpleTokenProperties) return false
+
+    if (tokenType != other.tokenType) return false
+    if (location != other.location) return false
+    if (features != other.features) return false
+    if (additional != other.additional) return false
+
+    return true
+  }
 }
 
 class DocumentationProperties(val docComment: String, val startOffset: Int, val endOffset: Int, val docStartOffset: Int, val docEndOffset: Int, val nameIdentifierOffset: Int) : TokenProperties {

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/EvaluationWorkspace.kt

@@ -6,6 +6,9 @@ import com.intellij.cce.evaluable.EvaluationStrategy
 import com.intellij.cce.evaluable.StrategySerializer
 import com.intellij.cce.fus.FusLogsSaver
 import com.intellij.cce.workspace.storages.*
+import com.intellij.cce.workspace.storages.storage.ActionsStorage
+import com.intellij.cce.workspace.storages.storage.ActionsStorageFactory
+import com.intellij.cce.workspace.storages.storage.getActionsStorageTypeFromEnv
 import java.io.FileWriter
 import java.nio.file.Files
 import java.nio.file.Path
@@ -42,7 +45,7 @@ class EvaluationWorkspace private constructor(private val basePath: Path,
 
   val sessionsStorage: CompositeSessionsStorage = CompositeSessionsStorage(sessionsDir.toString())
 
-  val actionsStorage: ActionsStorage = ActionsStorage(actionsDir.toString())
+  val actionsStorage: ActionsStorage = ActionsStorageFactory.create(actionsDir.toString(), getActionsStorageTypeFromEnv())
 
   val errorsStorage: FileErrorsStorage = FileErrorsStorage(errorsDir.toString())
 

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/FileArchivesStorage.kt

@@ -12,8 +12,7 @@ class FileArchivesStorage(private val storageDir: String) : KeyValueStorage<Stri
   val fileExtension: String = ".gz"
 
   init {
-    val storagePath = Paths.get(storageDir)
-    if (!Files.exists(storagePath)) Files.createDirectories(storagePath)
+    ensureDirExists(storageDir)
   }
 
   override fun get(key: String): String {

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/storage/ActionsMultiplyFilesStorage.kt

@@ -1,31 +1,32 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package com.intellij.cce.workspace.storages
+package com.intellij.cce.workspace.storages.storage
 
 import com.intellij.cce.actions.ActionSerializer
 import com.intellij.cce.actions.FileActions
+import com.intellij.cce.workspace.storages.FileArchivesStorage
 import java.nio.file.Paths
 
-class ActionsStorage(storageDir: String) {
+internal class ActionsMultiplyFilesStorage(storageDir: String) : ActionsStorage {
   private val keyValueStorage = FileArchivesStorage(storageDir)
   private var filesCounter = 0
 
-  fun saveActions(actions: FileActions) {
+  override fun saveActions(actions: FileActions) {
     filesCounter++
     keyValueStorage.save("${Paths.get(actions.path).fileName}($filesCounter).json", ActionSerializer.serialize(actions))
   }
 
-  fun computeSessionsCount(): Int {
+  override fun computeSessionsCount(): Int {
     var count = 0
     for (file in getActionFiles())
       count += ActionSerializer.getSessionsCount(keyValueStorage.get(file))
     return count
   }
 
-  fun getActionFiles(): List<String> = keyValueStorage.getKeys().sortedBy {
+  override fun getActionFiles(): List<String> = keyValueStorage.getKeys().sortedBy {
     it.substringAfterLast('(').substringBefore(')').toInt()
   }
 
-  fun getActions(path: String): FileActions {
+  override fun getActions(path: String): FileActions {
     return ActionSerializer.deserialize(keyValueStorage.get(path))
   }
 }

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/storage/ActionsSingleFileStorage.kt

@@ -0,0 +1,54 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.cce.workspace.storages.storage
+
+import com.intellij.cce.actions.ActionArraySerializer
+import com.intellij.cce.actions.FileActions
+import com.intellij.openapi.diagnostic.logger
+import java.io.File
+import java.nio.file.Files
+import java.nio.file.Paths
+import kotlin.io.path.absolutePathString
+
+class ActionsSingleFileStorage(storageDir: String) : ActionsStorage {
+  private val LOG = logger<ActionsSingleFileStorage>()
+
+  private val file: File
+
+  init {
+    val filePath = Paths.get(storageDir, "actions")
+    file = if (Files.exists(filePath)) {
+      File(filePath.absolutePathString())
+    } else {
+      Files.createFile(filePath).toFile()
+    }
+  }
+
+  override fun saveActions(actions: FileActions) {
+    val array = getSavedActions()
+    val newVal = ActionArraySerializer.serialize(array.toMutableList().also { it.add(actions) }.sortedBy { it.path }.toTypedArray())
+    file.writeText(newVal)
+  }
+
+  override fun computeSessionsCount(): Int {
+    return getSavedActions().sumOf { it.sessionsCount }
+  }
+
+  override fun getActionFiles(): List<String> {
+    return getSavedActions().map { it.path }
+  }
+
+  override fun getActions(path: String): FileActions {
+    return getSavedActions().singleOrNull { it.path == path } ?: error("there's no actions for file with path $path")
+  }
+
+  private fun getSavedActions(): Array<FileActions> {
+    val text = file.readText()
+    if (text.isEmpty()) return emptyArray()
+    return try {
+      ActionArraySerializer.deserialize(text)
+    } catch (t: Throwable) {
+      LOG.error("failed to deserialize actions", t)
+      throw t
+    }
+  }
+}

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/storage/ActionsStorage.kt

@@ -0,0 +1,33 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.cce.workspace.storages.storage
+
+import com.intellij.cce.actions.FileActions
+import com.intellij.cce.workspace.storages.ensureDirExists
+
+interface ActionsStorage {
+  fun saveActions(actions: FileActions)
+  fun computeSessionsCount(): Int
+  fun getActionFiles(): List<String>
+  fun getActions(path: String): FileActions
+}
+
+enum class ActionsStorageType {
+  MULTIPLY_FILES,
+  SINGLE_FILE,
+}
+
+fun getActionsStorageTypeFromEnv(): ActionsStorageType {
+  val envVar = System.getenv("AIA_EVALUATION_ACTIONS_STORAGE_TYPE")?.lowercase() ?: ""
+  return ActionsStorageType.entries.firstOrNull { it.name.lowercase() == envVar } ?: ActionsStorageType.MULTIPLY_FILES
+}
+
+object ActionsStorageFactory {
+  fun create(storageDir: String, type: ActionsStorageType): ActionsStorage {
+    ensureDirExists(storageDir)
+
+    return when (type) {
+      ActionsStorageType.MULTIPLY_FILES -> ActionsMultiplyFilesStorage(storageDir)
+      ActionsStorageType.SINGLE_FILE -> ActionsSingleFileStorage(storageDir)
+    }
+  }
+}

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/utils.kt

@@ -0,0 +1,10 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.cce.workspace.storages
+
+import java.nio.file.Files
+import java.nio.file.Paths
+
+fun ensureDirExists(dir: String) {
+  val path = Paths.get(dir)
+  if (!Files.exists(path)) Files.createDirectories(path)
+}

plugins/evaluation-plugin/test/com/intellij/cce/actions/ActionsStorageTest.kt

@@ -0,0 +1,80 @@
+package com.intellij.cce.actions
+
+import com.intellij.cce.core.SimpleTokenProperties
+import com.intellij.cce.core.SymbolLocation
+import com.intellij.cce.core.TypeProperty
+import com.intellij.cce.workspace.storages.storage.ActionsStorageFactory
+import com.intellij.cce.workspace.storages.storage.ActionsStorageType
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+
+private const val file1 = "file1"
+private const val file2 = "file2"
+
+class ActionsMultiplyFilesStorageTest : ActionsStorageTestBase(ActionsStorageType.MULTIPLY_FILES)
+
+class ActionsSingleFileStorageTest : ActionsStorageTestBase(ActionsStorageType.SINGLE_FILE)
+
+abstract class ActionsStorageTestBase(private val storageType: ActionsStorageType)  {
+  @field:TempDir
+  lateinit var tempDir: File
+
+  @Test
+  fun `test actions for 2 files`() {
+    val storage = ActionsStorageFactory.create(tempDir.absolutePath, storageType)
+    val actions = createFileActionsCollection()
+    actions.forEach {
+      storage.saveActions(it)
+    }
+
+    val actionFiles = storage.getActionFiles()
+
+    assertEquals(listOf(file1, file2).count(), actionFiles.count())
+    repeat(2) { index ->
+      assertEquals(actions[index], storage.getActions(actionFiles[index]))
+    }
+  }
+
+  @Test
+  fun `compute session files`() {
+    val storage = ActionsStorageFactory.create(tempDir.absolutePath, storageType)
+    val actions = createFileActionsCollection()
+    actions.forEach {
+      storage.saveActions(it)
+    }
+
+    assertEquals(3, storage.computeSessionsCount())
+  }
+
+  private fun createFileActionsCollection(): List<FileActions> {
+    return listOf(
+      createSingleFileActions(
+        file1,
+        ActionsBuilder().also {
+          it.session {
+            moveCaret(5)
+            callFeature("expectedText1", 300, SimpleTokenProperties.create(TypeProperty.TOKEN, SymbolLocation.UNKNOWN) {})
+          }
+        }.build()
+      ),
+      createSingleFileActions(
+        file2,
+        ActionsBuilder().also {
+          it.session {
+            moveCaret(15)
+            callFeature("expectedText2", 350, SimpleTokenProperties.create(TypeProperty.TOKEN, SymbolLocation.UNKNOWN) {})
+          }
+          it.session {
+            printText("some text")
+            callFeature("expectedText3", 550, SimpleTokenProperties.create(TypeProperty.TOKEN, SymbolLocation.UNKNOWN) {})
+          }
+        }.build()
+      ),
+    )
+  }
+
+  private fun createSingleFileActions(path: String, actions: List<Action>) =
+    FileActions(path, "checksum_for_$path", actions.count { it is CallFeature}, actions)
+}