A new CodeGenerator.createFile method with more flexibility

api/src/main/kotlin/com/google/devtools/ksp/processing/CodeGenerator.kt

@@ -63,6 +63,38 @@ interface CodeGenerator {
         extensionName: String = "kt"
     ): OutputStream
 
+    /**
+     * Creates a file which is managed by [CodeGenerator]
+     *
+     * Sources of corresponding [KSNode]s which are obtained directly from [Resolver] need to be specified.
+     * Namely, the containing files of those [KSNode]s who are obtained from:
+     *   * [Resolver.getAllFiles]
+     *   * [Resolver.getSymbolsWithAnnotation]
+     *   * [Resolver.getClassDeclarationByName]
+     *
+     * Instead of requiring processors to specify all source files which are relevant in generating the given output,
+     * KSP traces dependencies automatically and only needs to know those sources that only processors know what they
+     * are for. If a [KSFile] is indirectly obtained through other [KSNode]s, it hasn't to be specified for the given
+     * output, even if its contents contribute to the generation of the output.
+     *
+     * For example, a processor generates an output `O` after reading class `A` in `A.kt` and class `B` in `B.kt`,
+     * where `A` extends `B`. The processor got `A` by [Resolver.getSymbolsWithAnnotation] and then got `B` by
+     * [KSClassDeclaration.superTypes] from `A`. Because the inclusion of `B` is due to `A`, `B.kt` needn't to be
+     * specified in [dependencies] for `O`. Note that specifying `B.kt` in this case doesn't hurt, it is only unnecessary.
+     *
+     * @param dependencies are [KSFile]s from which this output is built. Only those that are obtained directly
+     *                     from [Resolver] are required.
+     * @param path corresponds to the relative path of the generated file; includes the full file name
+     * @param fileType determines the target directory to store the file
+     * @return OutputStream for writing into files.
+     * @see [CodeGenerator] for more details.
+     */
+    fun createNewFileByPath(
+        dependencies: Dependencies,
+        path: String,
+        extensionName: String = "kt"
+    ): OutputStream
+
     /**
      * Associate [sources] to an output file.
      *
@@ -76,6 +108,17 @@ interface CodeGenerator {
      */
     fun associate(sources: List<KSFile>, packageName: String, fileName: String, extensionName: String = "kt")
 
+    /**
+     * Associate [sources] to an output file.
+     *
+     * @param sources are [KSFile]s from which this output is built. Only those that are obtained directly
+     *                     from [Resolver] are required.
+     * @param path corresponds to the relative path of the generated file; includes the full file name
+     * @param fileType determines the target directory where the file should exist
+     * @see [CodeGenerator] for more details.
+     */
+    fun associate(sources: List<KSFile>, path: String, extensionName: String = "kt")
+
     /**
      * Associate [classes] to an output file.
      *

api/src/main/kotlin/com/google/devtools/ksp/processing/FileType.kt

@@ -0,0 +1,12 @@
+package com.google.devtools.ksp.processing
+
+/**
+ * Used to determine where files should be stored when being
+ * created through the [CodeGenerator].
+ */
+enum class FileType {
+    CLASS,
+    JAVA_SOURCE,
+    KOTLIN_SOURCE,
+    RESOURCE
+}

common-util/src/main/kotlin/com/google/devtools/ksp/processing/impl/CodeGeneratorImpl.kt

@@ -20,11 +20,14 @@ package com.google.devtools.ksp.processing.impl
 import com.google.devtools.ksp.NoSourceFile
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.FileType
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSFile
 import java.io.File
 import java.io.FileOutputStream
+import java.io.IOException
 import java.io.OutputStream
+import java.nio.file.Path
 
 class CodeGeneratorImpl(
     private val classDir: File,
@@ -54,13 +57,7 @@ class CodeGeneratorImpl(
     fun pathOf(packageName: String, fileName: String, extensionName: String): String {
         val packageDirs = if (packageName != "") "${packageName.split(".").joinToString(separator)}$separator" else ""
         val extension = if (extensionName != "") ".$extensionName" else ""
-        val typeRoot = when (extensionName) {
-            "class" -> classDir
-            "java" -> javaDir
-            "kt" -> kotlinDir
-            else -> resourcesDir
-        }.path
-        return "$typeRoot$separator$packageDirs$fileName$extension"
+        return "$packageDirs$fileName$extension"
     }
 
     override fun createNewFile(
@@ -69,17 +66,70 @@ class CodeGeneratorImpl(
         fileName: String,
         extensionName: String
     ): OutputStream {
+        return createNewFile(dependencies, pathOf(packageName, fileName, extensionName), extensionToType(extensionName))
+    }
+
+    override fun createNewFileByPath(dependencies: Dependencies, path: String, extensionName: String): OutputStream {
+        val extension = if (extensionName != "") ".$extensionName" else ""
+        return createNewFile(dependencies, path + extension, extensionToType(extensionName))
+    }
+
+    override fun associate(sources: List<KSFile>, packageName: String, fileName: String, extensionName: String) {
+        associate(sources, pathOf(packageName, fileName, extensionName), extensionToType(extensionName))
+    }
+
+    override fun associate(sources: List<KSFile>, path: String, extensionName: String) {
+        val extension = if (extensionName != "") ".$extensionName" else ""
+        associate(sources, path + extension, extensionToType(extensionName))
+    }
+
+    override fun associateWithClasses(
+        classes: List<KSClassDeclaration>,
+        packageName: String,
+        fileName: String,
+        extensionName: String
+    ) {
         val path = pathOf(packageName, fileName, extensionName)
-        val file = File(path)
-        if (path in fileMap) {
+        val files = classes.map {
+            it.containingFile ?: NoSourceFile(projectBase, it.qualifiedName?.asString().toString())
+        }
+        associate(files, path, extensionToType(extensionName))
+    }
+
+    private fun extensionToType(extensionName: String): FileType {
+        return when (extensionName) {
+            "class" -> FileType.CLASS
+            "java" -> FileType.JAVA_SOURCE
+            "kt" -> FileType.KOTLIN_SOURCE
+            else -> FileType.RESOURCE
+        }
+    }
+
+    private fun baseDirOf(fileType: FileType): File {
+        return when (fileType) {
+            FileType.CLASS -> classDir
+            FileType.JAVA_SOURCE -> javaDir
+            FileType.KOTLIN_SOURCE -> kotlinDir
+            FileType.RESOURCE -> resourcesDir
+        }
+    }
+
+    private fun createNewFile(dependencies: Dependencies, path: String, fileType: FileType): OutputStream {
+        val baseDir = baseDirOf(fileType)
+        val file = File(baseDir, path)
+        if (!isWithinBaseDir(baseDir, file)) {
+            throw IllegalStateException("requested path is outside the bounds of the required directory")
+        }
+        val absolutePath = file.absolutePath
+        if (absolutePath in fileMap) {
             throw FileAlreadyExistsException(file)
         }
         val parentFile = file.parentFile
         if (!parentFile.exists() && !parentFile.mkdirs()) {
             throw IllegalStateException("failed to make parent directories.")
         }
         file.writeText("")
-        fileMap[path] = file
+        fileMap[absolutePath] = file
         val sources = if (dependencies.isAllSources) {
             allSources + anyChangesWildcard
         } else {
@@ -89,41 +139,42 @@ class CodeGeneratorImpl(
                 dependencies.originatingFiles
             }
         }
-        associate(sources, path)
-        fileOutputStreamMap[path] = fileMap[path]!!.outputStream()
-        return fileOutputStreamMap[path]!!
+        associate(sources, file)
+        fileOutputStreamMap[absolutePath] = fileMap[absolutePath]!!.outputStream()
+        return fileOutputStreamMap[absolutePath]!!
     }
 
-    override fun associate(sources: List<KSFile>, packageName: String, fileName: String, extensionName: String) {
-        val path = pathOf(packageName, fileName, extensionName)
-        associate(sources, path)
+    private fun isWithinBaseDir(baseDir: File, file: File): Boolean {
+        val base = baseDir.toPath().normalize()
+        return try {
+            val relativePath = file.toPath().normalize()
+            relativePath.startsWith(base)
+        } catch (e: IOException) {
+            false
+        }
     }
 
-    override fun associateWithClasses(
-        classes: List<KSClassDeclaration>,
-        packageName: String,
-        fileName: String,
-        extensionName: String
-    ) {
-        val path = pathOf(packageName, fileName, extensionName)
-        val files = classes.map {
-            it.containingFile ?: NoSourceFile(projectBase, it.qualifiedName?.asString().toString())
+    private fun associate(sources: List<KSFile>, path: String, fileType: FileType) {
+        val baseDir = baseDirOf(fileType)
+        val file = File(baseDir, path)
+        if (!isWithinBaseDir(baseDir, file)) {
+            throw IllegalStateException("requested path is outside the bounds of the required directory")
         }
-        associate(files, path)
+        associate(sources, file)
     }
 
-    private fun associate(sources: List<KSFile>, outputPath: String) {
+    private fun associate(sources: List<KSFile>, outputPath: File) {
         if (!isIncremental)
             return
 
-        val output = File(outputPath).relativeTo(projectBase)
+        val output = outputPath.relativeTo(projectBase)
         sources.forEach { source ->
             sourceToOutputs.getOrPut(File(source.filePath).relativeTo(projectBase)) { mutableSetOf() }.add(output)
         }
     }
 
     val outputs: Set<File>
-        get() = fileMap.keys.mapTo(mutableSetOf()) { File(it).relativeTo(projectBase) }
+        get() = fileMap.values.mapTo(mutableSetOf()) { it.relativeTo(projectBase) }
 
     override val generatedFile: Collection<File>
         get() = fileOutputStreamMap.keys.map { fileMap[it]!! }

common-util/src/test/kotlin/com/google/devtools/ksp/processing/impl/CodeGeneratorImplTest.kt

@@ -0,0 +1,128 @@
+package com.google.devtools.ksp.processing.impl
+
+import com.google.devtools.ksp.AnyChanges
+import com.google.devtools.ksp.processing.Dependencies
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Test
+import java.io.File
+import java.nio.file.Files
+
+class CodeGeneratorImplTest {
+
+    lateinit var codeGenerator: CodeGeneratorImpl
+    lateinit var baseDir: File
+
+    @Before
+    fun setup() {
+        baseDir = Files.createTempDirectory("project").toFile()
+        val classesDir = File(baseDir, "classes")
+        classesDir.mkdir()
+        val javaDir = File(baseDir, "java")
+        javaDir.mkdir()
+        val kotlinDir = File(baseDir, "kotlin")
+        kotlinDir.mkdir()
+        val resourcesDir = File(baseDir, "resources")
+        resourcesDir.mkdir()
+        codeGenerator = CodeGeneratorImpl(
+            classesDir,
+            javaDir,
+            kotlinDir,
+            resourcesDir,
+            baseDir,
+            AnyChanges(baseDir),
+            emptyList(),
+            true
+        )
+    }
+
+    @Test
+    fun testCreatingAFile() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "java")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "kt")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "class")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+
+        try {
+            codeGenerator.outputs
+        } catch (e: Exception) {
+            Assert.fail("Failed to get outputs: ${e.message}")
+        }
+    }
+
+    @Test
+    fun testCreatingAFileWithSlash() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "java")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "kt")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "class")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+
+        try {
+            codeGenerator.outputs
+        } catch (e: Exception) {
+            Assert.fail("Failed to get outputs: ${e.message}")
+        }
+    }
+
+    @Test
+    fun testCreatingAFileWithPath() {
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/Test", "java")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/Test")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/Test", "class")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+
+        try {
+            codeGenerator.outputs
+        } catch (e: Exception) {
+            Assert.fail("Failed to get outputs: ${e.message}")
+        }
+    }
+
+    @Test
+    fun testCreatingAFileWithPathAndDots() {
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test", "java")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test", "class")
+        codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/dir.with.dot/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/dir.with.dot/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/dir.with.dot/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/dir.with.dot/Test"), files[3])
+
+        try {
+            codeGenerator.outputs
+        } catch (e: Exception) {
+            Assert.fail("Failed to get outputs: ${e.message}")
+        }
+    }
+
+    @Test
+    fun testCreatingAFileByPathWithInvalidPath() {
+        try {
+            codeGenerator.createNewFileByPath(Dependencies.ALL_FILES, "../../b/c/Test", "java")
+            Assert.fail()
+        } catch (e: java.lang.IllegalStateException) {
+            Assert.assertEquals(e.message, "requested path is outside the bounds of the required directory")
+        }
+    }
+}