1 /*
2 * Copyright (C) 2021 Square, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * https://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 package com.squareup.kotlinpoet.ksp
17
18 import com.google.devtools.ksp.processing.CodeGenerator
19 import com.google.devtools.ksp.processing.Dependencies
20 import com.google.devtools.ksp.symbol.KSFile
21 import com.squareup.kotlinpoet.FileSpec
22 import com.squareup.kotlinpoet.FunSpec
23 import com.squareup.kotlinpoet.PropertySpec
24 import com.squareup.kotlinpoet.Taggable
25 import com.squareup.kotlinpoet.TypeAliasSpec
26 import com.squareup.kotlinpoet.TypeSpec
27 import com.squareup.kotlinpoet.tag
28 import java.io.OutputStreamWriter
29 import java.nio.charset.StandardCharsets
30
31 /**
32 * A simple holder class for containing originating [KSFiles][KSFile], which are used by KSP to
33 * inform its incremental processing.
34 *
35 * See [the docs](https://kotlinlang.org/docs/ksp-incremental.html) for more information.
36 */
37 public interface OriginatingKSFiles {
38 public val files: List<KSFile>
39 }
40
41 /** Returns this spec's originating [KSFiles][KSFile] for use with incremental processing. */
originatingKSFilesnull42 public fun TypeSpec.originatingKSFiles(): List<KSFile> = getKSFilesTag()
43
44 /** Returns this spec's originating [KSFiles][KSFile] for use with incremental processing. */
45 public fun FunSpec.originatingKSFiles(): List<KSFile> = getKSFilesTag()
46
47 /** Returns this spec's originating [KSFiles][KSFile] for use with incremental processing. */
48 public fun PropertySpec.originatingKSFiles(): List<KSFile> = getKSFilesTag()
49
50 /** Returns this spec's originating [KSFiles][KSFile] for use with incremental processing. */
51 public fun TypeAliasSpec.originatingKSFiles(): List<KSFile> = getKSFilesTag()
52
53 /**
54 * Returns the list of all files added to the contained
55 * [TypeSpecs][TypeSpec], [PropertySpecs][PropertySpec], [FunSpecs][FunSpec], or
56 * [TypeAliasSpecs][TypeAliasSpec] contained in this spec.
57 */
58 public fun FileSpec.originatingKSFiles(): List<KSFile> {
59 return members
60 .flatMap {
61 when (it) {
62 is FunSpec -> it.originatingKSFiles()
63 is PropertySpec -> it.originatingKSFiles()
64 is TypeSpec -> it.originatingKSFiles()
65 is TypeAliasSpec -> it.originatingKSFiles()
66 else -> emptyList()
67 }
68 }
69 .distinct()
70 }
71
72 /** Adds the given [ksFile] to this builder's tags for use with [originatingKSFiles]. */
addOriginatingKSFilenull73 public fun TypeAliasSpec.Builder.addOriginatingKSFile(ksFile: KSFile): TypeAliasSpec.Builder = apply {
74 getOrCreateKSFilesTag().add(ksFile)
75 }
76
77 /** Adds the given [ksFile] to this builder's tags for use with [originatingKSFiles]. */
addOriginatingKSFilenull78 public fun PropertySpec.Builder.addOriginatingKSFile(ksFile: KSFile): PropertySpec.Builder = apply {
79 getOrCreateKSFilesTag().add(ksFile)
80 }
81
82 /** Adds the given [ksFile] to this builder's tags for use with [originatingKSFiles]. */
addOriginatingKSFilenull83 public fun FunSpec.Builder.addOriginatingKSFile(ksFile: KSFile): FunSpec.Builder = apply {
84 getOrCreateKSFilesTag().add(ksFile)
85 }
86
87 /** Adds the given [ksFile] to this builder's tags for use with [originatingKSFiles]. */
addOriginatingKSFilenull88 public fun TypeSpec.Builder.addOriginatingKSFile(ksFile: KSFile): TypeSpec.Builder = apply {
89 getOrCreateKSFilesTag().add(ksFile)
90 }
91
92 /**
93 * Writes this [FileSpec] to a given [codeGenerator] with the given [originatingKSFiles].
94 *
95 * Note that if none are specified, the [originatingKSFiles] argument defaults to using
96 * [FileSpec.originatingKSFiles], which will automatically resolve any files added to the
97 * contained declarations.
98 *
99 * See [the docs](https://github.com/google/ksp/blob/main/docs/incremental.md) for more information.
100 *
101 * @see FileSpec.originatingKSFiles
102 * @param codeGenerator the [CodeGenerator] to write to.
103 * @param aggregating flag indicating if this is an aggregating symbol processor.
104 */
writeTonull105 public fun FileSpec.writeTo(
106 codeGenerator: CodeGenerator,
107 aggregating: Boolean,
108 originatingKSFiles: Iterable<KSFile> = originatingKSFiles(),
109 ) {
110 val dependencies = kspDependencies(aggregating, originatingKSFiles)
111 writeTo(codeGenerator, dependencies)
112 }
113
114 /**
115 * Writes this [FileSpec] to a given [codeGenerator] with the given [dependencies].
116 *
117 * See [the docs](https://github.com/google/ksp/blob/main/docs/incremental.md) for more information.
118 *
119 * @see FileSpec.originatingKSFiles
120 * @see kspDependencies
121 * @param codeGenerator the [CodeGenerator] to write to.
122 * @param dependencies the [Dependencies] to create a new file with.
123 */
writeTonull124 public fun FileSpec.writeTo(
125 codeGenerator: CodeGenerator,
126 dependencies: Dependencies,
127 ) {
128 val file = codeGenerator.createNewFile(dependencies, packageName, name)
129 // Don't use writeTo(file) because that tries to handle directories under the hood
130 OutputStreamWriter(file, StandardCharsets.UTF_8)
131 .use(::writeTo)
132 }
133
134 /**
135 * Returns a KSP [Dependencies] component of this [FileSpec] with the given [originatingKSFiles],
136 * intended to be used in tandem with [writeTo].
137 *
138 * Note that if no [originatingKSFiles] are specified, the [originatingKSFiles] argument defaults
139 * to using [FileSpec.originatingKSFiles], which will automatically resolve any files added to the
140 * contained declarations.
141 *
142 * See [the docs](https://github.com/google/ksp/blob/main/docs/incremental.md) for more information.
143 *
144 * @see FileSpec.originatingKSFiles
145 * @see FileSpec.writeTo
146 * @param aggregating flag indicating if this is an aggregating symbol processor.
147 */
kspDependenciesnull148 public fun FileSpec.kspDependencies(
149 aggregating: Boolean,
150 originatingKSFiles: Iterable<KSFile> = originatingKSFiles(),
151 ): Dependencies = Dependencies(aggregating, *originatingKSFiles.toList().toTypedArray())
152
153 /**
154 * A mutable [OriginatingKSFiles] instance for use with KotlinPoet Builders via [Taggable.Builder].
155 */
156 private interface MutableOriginatingKSFiles : OriginatingKSFiles {
157 override val files: MutableList<KSFile>
158 }
159
160 private data class MutableOriginatingKSFilesImpl(override val files: MutableList<KSFile> = mutableListOf()) : MutableOriginatingKSFiles
161
getKSFilesTagnull162 private fun Taggable.getKSFilesTag(): List<KSFile> {
163 return tag<OriginatingKSFiles>()?.files.orEmpty()
164 }
165
getOrCreateKSFilesTagnull166 private fun Taggable.Builder<*>.getOrCreateKSFilesTag(): MutableList<KSFile> {
167 val holder = tags.getOrPut(
168 OriginatingKSFiles::class,
169 ::MutableOriginatingKSFilesImpl,
170 ) as MutableOriginatingKSFiles
171 return holder.files
172 }
173