1 /*
<lambda>null2  * Copyright (C) 2020 The Dagger Authors.
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  * http://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 
17 package dagger.hilt.android.processor.internal.viewmodel
18 
19 import com.google.auto.common.GeneratedAnnotationSpecs
20 import com.squareup.javapoet.AnnotationSpec
21 import com.squareup.javapoet.ClassName
22 import com.squareup.javapoet.JavaFile
23 import com.squareup.javapoet.MethodSpec
24 import com.squareup.javapoet.TypeSpec
25 import dagger.hilt.android.processor.internal.AndroidClassNames
26 import dagger.hilt.processor.internal.ClassNames
27 import javax.annotation.processing.ProcessingEnvironment
28 import javax.lang.model.SourceVersion
29 import javax.lang.model.element.Modifier
30 import javax.lang.model.util.Elements
31 
32 /**
33  * Source generator to support Hilt injection of ViewModels.
34  *
35  * Should generate:
36  * ```
37  * public final class $_HiltModules {
38  *   @Module
39  *   @InstallIn(ViewModelComponent.class)
40  *   public static abstract class BindsModule {
41  *     @Binds
42  *     @IntoMap
43  *     @StringKey("pkg.$")
44  *     @HiltViewModelMap
45  *     public abstract ViewModel bind($ vm)
46  *   }
47  *   @Module
48  *   @InstallIn(ActivityRetainedComponent.class)
49  *   public static final class KeyModule {
50  *     @Provides
51  *     @IntoSet
52  *     @HiltViewModelMap.KeySet
53  *     public static String provide() {
54  *      return "pkg.$";
55  *     }
56  *   }
57  * }
58  * ```
59  */
60 internal class ViewModelModuleGenerator(
61   private val processingEnv: ProcessingEnvironment,
62   private val injectedViewModel: ViewModelMetadata
63 ) {
64   fun generate() {
65     val modulesTypeSpec = TypeSpec.classBuilder(injectedViewModel.modulesClassName)
66       .addOriginatingElement(injectedViewModel.typeElement)
67       .addGeneratedAnnotation(processingEnv.elementUtils, processingEnv.sourceVersion)
68       .addAnnotation(
69         AnnotationSpec.builder(ClassNames.ORIGINATING_ELEMENT)
70           .addMember(
71             "topLevelClass",
72             "$T.class",
73             injectedViewModel.className.topLevelClassName()
74           )
75           .build()
76       )
77       .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
78       .addType(getBindsModuleTypeSpec())
79       .addType(getKeyModuleTypeSpec())
80       .addMethod(
81         MethodSpec.constructorBuilder()
82           .addModifiers(Modifier.PRIVATE)
83           .build()
84       )
85       .build()
86     JavaFile.builder(injectedViewModel.modulesClassName.packageName(), modulesTypeSpec)
87       .build()
88       .writeTo(processingEnv.filer)
89   }
90 
91   private fun getBindsModuleTypeSpec() = createModuleTypeSpec(
92     className = "BindsModule",
93     component = AndroidClassNames.VIEW_MODEL_COMPONENT
94   )
95     .addModifiers(Modifier.ABSTRACT)
96     .addMethod(getViewModelBindsMethod())
97     .build()
98 
99   private fun getViewModelBindsMethod() =
100     MethodSpec.methodBuilder("binds")
101       .addAnnotation(ClassNames.BINDS)
102       .addAnnotation(ClassNames.INTO_MAP)
103       .addAnnotation(
104         AnnotationSpec.builder(ClassNames.STRING_KEY)
105           .addMember("value", S, injectedViewModel.className.reflectionName())
106           .build()
107       )
108       .addAnnotation(AndroidClassNames.HILT_VIEW_MODEL_MAP_QUALIFIER)
109       .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
110       .returns(AndroidClassNames.VIEW_MODEL)
111       .addParameter(injectedViewModel.className, "vm")
112       .build()
113 
114   private fun getKeyModuleTypeSpec() = createModuleTypeSpec(
115     className = "KeyModule",
116     component = AndroidClassNames.ACTIVITY_RETAINED_COMPONENT
117   )
118     .addModifiers(Modifier.FINAL)
119     .addMethod(
120       MethodSpec.constructorBuilder()
121         .addModifiers(Modifier.PRIVATE)
122         .build()
123     )
124     .addMethod(getViewModelKeyProvidesMethod())
125     .build()
126 
127   private fun getViewModelKeyProvidesMethod() =
128     MethodSpec.methodBuilder("provide")
129       .addAnnotation(ClassNames.PROVIDES)
130       .addAnnotation(ClassNames.INTO_SET)
131       .addAnnotation(AndroidClassNames.HILT_VIEW_MODEL_KEYS_QUALIFIER)
132       .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
133       .returns(String::class.java)
134       .addStatement("return $S", injectedViewModel.className.reflectionName())
135       .build()
136 
137   private fun createModuleTypeSpec(className: String, component: ClassName) =
138     TypeSpec.classBuilder(className)
139       .addOriginatingElement(injectedViewModel.typeElement)
140       .addAnnotation(ClassNames.MODULE)
141       .addAnnotation(
142         AnnotationSpec.builder(ClassNames.INSTALL_IN)
143           .addMember("value", "$T.class", component)
144           .build()
145       )
146       .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
147 
148   companion object {
149 
150     const val L = "\$L"
151     const val T = "\$T"
152     const val N = "\$N"
153     const val S = "\$S"
154     const val W = "\$W"
155 
156     private fun TypeSpec.Builder.addGeneratedAnnotation(
157       elements: Elements,
158       sourceVersion: SourceVersion
159     ) = apply {
160       GeneratedAnnotationSpecs.generatedAnnotationSpec(
161         elements,
162         sourceVersion,
163         ViewModelProcessor::class.java
164       ).ifPresent { generatedAnnotation ->
165         addAnnotation(generatedAnnotation)
166       }
167     }
168   }
169 }
170