1 /*
2  * Copyright 2020 The Android Open Source Project
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 androidx.appsearch.compiler;
18 
19 import static androidx.appsearch.compiler.IntrospectionHelper.APPSEARCH_EXCEPTION_CLASS;
20 import static androidx.appsearch.compiler.IntrospectionHelper.APPSEARCH_SCHEMA_CLASS;
21 import static androidx.appsearch.compiler.IntrospectionHelper.PROPERTY_CONFIG_CLASS;
22 import static androidx.appsearch.compiler.IntrospectionHelper.getDocumentClassFactoryForClass;
23 
24 import static com.google.auto.common.MoreTypes.asTypeElement;
25 
26 import static javax.lang.model.type.TypeKind.DECLARED;
27 
28 import androidx.appsearch.compiler.annotationwrapper.DataPropertyAnnotation;
29 import androidx.appsearch.compiler.annotationwrapper.DocumentPropertyAnnotation;
30 import androidx.appsearch.compiler.annotationwrapper.EmbeddingPropertyAnnotation;
31 import androidx.appsearch.compiler.annotationwrapper.LongPropertyAnnotation;
32 import androidx.appsearch.compiler.annotationwrapper.StringPropertyAnnotation;
33 
34 import com.squareup.javapoet.ClassName;
35 import com.squareup.javapoet.CodeBlock;
36 import com.squareup.javapoet.FieldSpec;
37 import com.squareup.javapoet.MethodSpec;
38 import com.squareup.javapoet.ParameterizedTypeName;
39 import com.squareup.javapoet.TypeName;
40 import com.squareup.javapoet.TypeSpec;
41 import com.squareup.javapoet.WildcardTypeName;
42 
43 import org.jspecify.annotations.NonNull;
44 
45 import java.util.ArrayDeque;
46 import java.util.ArrayList;
47 import java.util.Collections;
48 import java.util.HashSet;
49 import java.util.LinkedHashSet;
50 import java.util.List;
51 import java.util.Queue;
52 import java.util.Set;
53 
54 import javax.annotation.processing.ProcessingEnvironment;
55 import javax.lang.model.element.Modifier;
56 import javax.lang.model.element.TypeElement;
57 import javax.lang.model.type.TypeMirror;
58 
59 /** Generates java code for an {@link androidx.appsearch.app.AppSearchSchema}. */
60 class SchemaCodeGenerator {
61     private final DocumentModel mModel;
62     private final IntrospectionHelper mHelper;
63     private final LinkedHashSet<TypeElement> mDependencyDocumentClasses;
64 
generate( @onNull ProcessingEnvironment env, @NonNull DocumentModel model, TypeSpec.@NonNull Builder classBuilder)65     public static void generate(
66             @NonNull ProcessingEnvironment env,
67             @NonNull DocumentModel model,
68             TypeSpec.@NonNull Builder classBuilder) throws ProcessingException {
69         new SchemaCodeGenerator(model, env).generate(classBuilder);
70     }
71 
SchemaCodeGenerator(@onNull DocumentModel model, @NonNull ProcessingEnvironment env)72     private SchemaCodeGenerator(@NonNull DocumentModel model, @NonNull ProcessingEnvironment env) {
73         mModel = model;
74         mHelper = new IntrospectionHelper(env);
75         mDependencyDocumentClasses = computeDependencyClasses(model, env);
76     }
77 
computeDependencyClasses( @onNull DocumentModel model, @NonNull ProcessingEnvironment env)78     private static @NonNull LinkedHashSet<TypeElement> computeDependencyClasses(
79             @NonNull DocumentModel model,
80             @NonNull ProcessingEnvironment env) {
81         LinkedHashSet<TypeElement> dependencies = new LinkedHashSet<>(model.getParentTypes());
82         for (AnnotatedGetterOrField getterOrField : model.getAnnotatedGettersAndFields()) {
83             if (!(getterOrField.getAnnotation() instanceof DocumentPropertyAnnotation)) {
84                 continue;
85             }
86 
87             TypeMirror documentClass = getterOrField.getComponentType();
88             dependencies.add((TypeElement) env.getTypeUtils().asElement(documentClass));
89         }
90         return dependencies;
91     }
92 
generate(TypeSpec.@onNull Builder classBuilder)93     private void generate(TypeSpec.@NonNull Builder classBuilder) throws ProcessingException {
94         classBuilder.addField(
95                 FieldSpec.builder(String.class, "SCHEMA_NAME")
96                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
97                         .initializer("$S", mModel.getSchemaName())
98                         .build());
99 
100         classBuilder.addMethod(
101                 MethodSpec.methodBuilder("getSchemaName")
102                         .addModifiers(Modifier.PUBLIC)
103                         .returns(String.class)
104                         .addAnnotation(Override.class)
105                         .addStatement("return SCHEMA_NAME")
106                         .build());
107 
108         classBuilder.addMethod(
109                 MethodSpec.methodBuilder("getSchema")
110                         .addModifiers(Modifier.PUBLIC)
111                         .returns(APPSEARCH_SCHEMA_CLASS)
112                         .addAnnotation(Override.class)
113                         .addException(APPSEARCH_EXCEPTION_CLASS)
114                         .addStatement("return $L", createSchemaInitializerGetDocumentTypes())
115                         .build());
116 
117         classBuilder.addMethod(createDependencyClassesMethod());
118     }
119 
createDependencyClassesMethod()120     private @NonNull MethodSpec createDependencyClassesMethod() {
121         TypeName listOfClasses = ParameterizedTypeName.get(ClassName.get("java.util", "List"),
122                 ParameterizedTypeName.get(ClassName.get(Class.class),
123                         WildcardTypeName.subtypeOf(Object.class)));
124 
125         TypeName arrayListOfClasses =
126                 ParameterizedTypeName.get(ClassName.get("java.util", "ArrayList"),
127                         ParameterizedTypeName.get(ClassName.get(Class.class),
128                                 WildcardTypeName.subtypeOf(Object.class)));
129 
130         MethodSpec.Builder methodBuilder = MethodSpec.methodBuilder("getDependencyDocumentClasses")
131                 .addModifiers(Modifier.PUBLIC)
132                 .returns(listOfClasses)
133                 .addAnnotation(Override.class)
134                 .addException(APPSEARCH_EXCEPTION_CLASS);
135 
136         if (mDependencyDocumentClasses.isEmpty()) {
137             methodBuilder.addStatement("return $T.emptyList()", ClassName.get(Collections.class));
138         } else {
139             methodBuilder.addStatement("$T classSet = new $T()", listOfClasses, arrayListOfClasses);
140             for (TypeElement dependencyType : mDependencyDocumentClasses) {
141                 methodBuilder.addStatement("classSet.add($T.class)", ClassName.get(dependencyType));
142             }
143             methodBuilder.addStatement("return classSet").build();
144         }
145 
146         return methodBuilder.build();
147     }
148 
149     /**
150      * Creates an expr of type {@link androidx.appsearch.app.AppSearchSchema}.
151      *
152      * <p>The AppSearchSchema has parent types and various Document.*Properties set.
153      */
createSchemaInitializerGetDocumentTypes()154     private CodeBlock createSchemaInitializerGetDocumentTypes() throws ProcessingException {
155         CodeBlock.Builder codeBlock = CodeBlock.builder()
156                 .add("new $T(SCHEMA_NAME)", APPSEARCH_SCHEMA_CLASS.nestedClass("Builder"))
157                 .indent();
158         for (TypeElement parentType : mModel.getParentTypes()) {
159             ClassName parentDocumentFactoryClass =
160                     getDocumentClassFactoryForClass(ClassName.get(parentType));
161             codeBlock.add("\n.addParentType($T.SCHEMA_NAME)", parentDocumentFactoryClass);
162         }
163 
164         for (AnnotatedGetterOrField getterOrField : mModel.getAnnotatedGettersAndFields()) {
165             if (!(getterOrField.getAnnotation() instanceof DataPropertyAnnotation)) {
166                 continue;
167             }
168 
169             CodeBlock propertyConfigExpr = createPropertyConfig(
170                     (DataPropertyAnnotation) getterOrField.getAnnotation(), getterOrField);
171             codeBlock.add("\n.addProperty($L)", propertyConfigExpr);
172         }
173 
174         codeBlock.add("\n.build()").unindent();
175         return codeBlock.build();
176     }
177 
178     /**
179      * Produces an expr for the creating the property's config e.g.
180      *
181      * <pre>
182      * {@code
183      * new StringPropertyConfig.Builder("someProp")
184      *   .setCardinality(StringPropertyConfig.CARDINALITY_REPEATED)
185      *   .setTokenizerType(StringPropertyConfig.TOKENIZER_TYPE_PLAIN)
186      *   .build()
187      * }
188      * </pre>
189      */
createPropertyConfig( @onNull DataPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)190     private CodeBlock createPropertyConfig(
191             @NonNull DataPropertyAnnotation annotation,
192             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
193         CodeBlock.Builder codeBlock = CodeBlock.builder();
194         if (annotation.getDataPropertyKind() == DataPropertyAnnotation.Kind.DOCUMENT_PROPERTY) {
195             ClassName documentClass = (ClassName) ClassName.get(getterOrField.getComponentType());
196             ClassName documentFactoryClass = getDocumentClassFactoryForClass(documentClass);
197             codeBlock.add("new $T.Builder($S, $T.SCHEMA_NAME)",
198                     DocumentPropertyAnnotation.CONFIG_CLASS,
199                     annotation.getName(),
200                     documentFactoryClass);
201         } else {
202             // All other property configs have a single param constructor that just takes the
203             // property's serialized name as input
204             codeBlock.add("new $T.Builder($S)",
205                     annotation.getConfigClassName(), annotation.getName());
206         }
207         codeBlock.indent().add(createSetCardinalityExpr(annotation, getterOrField));
208         switch (annotation.getDataPropertyKind()) {
209             case STRING_PROPERTY:
210                 StringPropertyAnnotation stringPropertyAnnotation =
211                         (StringPropertyAnnotation) annotation;
212                 codeBlock.add(createSetTokenizerTypeExpr(stringPropertyAnnotation, getterOrField))
213                         .add(createSetIndexingTypeExpr(stringPropertyAnnotation, getterOrField))
214                         .add(createSetJoinableValueTypeExpr(
215                                 stringPropertyAnnotation, getterOrField));
216                 break;
217             case DOCUMENT_PROPERTY:
218                 DocumentPropertyAnnotation documentPropertyAnnotation =
219                         (DocumentPropertyAnnotation) annotation;
220                 codeBlock.add(createSetShouldIndexNestedPropertiesExpr(documentPropertyAnnotation));
221                 Set<String> indexableNestedProperties = getAllIndexableNestedProperties(
222                         documentPropertyAnnotation);
223                 for (String propertyPath : indexableNestedProperties) {
224                     codeBlock.add(
225                             CodeBlock.of("\n.addIndexableNestedProperties($L)", propertyPath));
226                 }
227                 break;
228             case LONG_PROPERTY:
229                 LongPropertyAnnotation longPropertyAnnotation = (LongPropertyAnnotation) annotation;
230                 codeBlock.add(createSetIndexingTypeExpr(longPropertyAnnotation, getterOrField));
231                 break;
232             case EMBEDDING_PROPERTY:
233                 EmbeddingPropertyAnnotation embeddingPropertyAnnotation =
234                         (EmbeddingPropertyAnnotation) annotation;
235                 codeBlock
236                         .add(createSetIndexingTypeExpr(embeddingPropertyAnnotation, getterOrField))
237                         .add(createSetQuantizationTypeExpr(embeddingPropertyAnnotation,
238                                 getterOrField));
239                 break;
240             case DOUBLE_PROPERTY: // fall-through
241             case BOOLEAN_PROPERTY: // fall-through
242             case BYTES_PROPERTY: // fall-through
243             case BLOB_HANDLE_PROPERTY:
244                 break;
245             default:
246                 throw new IllegalStateException("Unhandled annotation: " + annotation);
247         }
248         return codeBlock.add("\n.build()")
249                 .unindent()
250                 .build();
251     }
252 
253 
254     /**
255      * Finds all indexable nested properties for the given type class and document property
256      * annotation. This includes indexable nested properties that should be inherited from the
257      * type's parent.
258      */
getAllIndexableNestedProperties( @onNull DocumentPropertyAnnotation documentPropertyAnnotation)259     private Set<String> getAllIndexableNestedProperties(
260             @NonNull DocumentPropertyAnnotation documentPropertyAnnotation)
261             throws ProcessingException {
262         Set<String> indexableNestedProperties = new HashSet<>(
263                 documentPropertyAnnotation.getIndexableNestedPropertiesList());
264 
265         if (documentPropertyAnnotation.getShouldInheritIndexableNestedPropertiesFromSuperClass()) {
266             // List of classes to expand into parent classes to search for the property annotation
267             Queue<TypeElement> classesToExpand = new ArrayDeque<>();
268             Set<TypeElement> visited = new HashSet<>();
269             classesToExpand.add(mModel.getClassElement());
270             while (!classesToExpand.isEmpty()) {
271                 TypeElement currentClass = classesToExpand.poll();
272                 if (visited.contains(currentClass)) {
273                     continue;
274                 }
275                 visited.add(currentClass);
276                 // Look for the document property annotation in the class's parent classes
277                 List<TypeMirror> parentTypes = new ArrayList<>();
278                 parentTypes.add(currentClass.getSuperclass());
279                 parentTypes.addAll(currentClass.getInterfaces());
280                 for (TypeMirror parent : parentTypes) {
281                     if (!parent.getKind().equals(DECLARED)) {
282                         continue;
283                     }
284                     TypeElement parentElement = asTypeElement(parent);
285                     DocumentPropertyAnnotation annotation = mHelper.getDocumentPropertyAnnotation(
286                             parentElement, documentPropertyAnnotation.getName());
287                     if (annotation == null) {
288                         // The property is not found in this level. Continue searching in one level
289                         // above as the property could still be defined for this level by class
290                         // inheritance.
291                         classesToExpand.add(parentElement);
292                     } else {
293                         indexableNestedProperties.addAll(
294                                 annotation.getIndexableNestedPropertiesList());
295                         if (annotation.getShouldInheritIndexableNestedPropertiesFromSuperClass()) {
296                             // Continue searching in the parent class's parents
297                             classesToExpand.add(parentElement);
298                         }
299                     }
300                 }
301             }
302         }
303         return indexableNestedProperties;
304     }
305 
306     /**
307      * Creates an expr like {@code .setCardinality(PropertyConfig.CARDINALITY_REPEATED)}.
308      */
createSetCardinalityExpr( @onNull DataPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)309     private static @NonNull CodeBlock createSetCardinalityExpr(
310             @NonNull DataPropertyAnnotation annotation,
311             @NonNull AnnotatedGetterOrField getterOrField) {
312         AnnotatedGetterOrField.ElementTypeCategory typeCategory =
313                 getterOrField.getElementTypeCategory();
314         String enumName;
315         switch (typeCategory) {
316             case COLLECTION: // fall-through
317             case ARRAY:
318                 enumName = "CARDINALITY_REPEATED";
319                 break;
320             case SINGLE:
321                 enumName = annotation.isRequired()
322                         ? "CARDINALITY_REQUIRED"
323                         : "CARDINALITY_OPTIONAL";
324                 break;
325             default:
326                 throw new IllegalStateException("Unhandled type category: " + typeCategory);
327         }
328         return CodeBlock.of("\n.setCardinality($T.$N)", PROPERTY_CONFIG_CLASS, enumName);
329     }
330 
331     /**
332      * Creates an expr like {@code .setTokenizerType(StringPropertyConfig.TOKENIZER_TYPE_PLAIN)}.
333      */
createSetTokenizerTypeExpr( @onNull StringPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)334     private static @NonNull CodeBlock createSetTokenizerTypeExpr(
335             @NonNull StringPropertyAnnotation annotation,
336             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
337         String enumName;
338         if (annotation.getIndexingType() == 0) { // INDEXING_TYPE_NONE
339             //TODO(b/171857731) remove this hack after apply to Icing lib's change.
340             enumName = "TOKENIZER_TYPE_NONE";
341         } else {
342             switch (annotation.getTokenizerType()) {
343                 case 0:
344                     enumName = "TOKENIZER_TYPE_NONE";
345                     break;
346                 case 1:
347                     enumName = "TOKENIZER_TYPE_PLAIN";
348                     break;
349                 case 2:
350                     enumName = "TOKENIZER_TYPE_VERBATIM";
351                     break;
352                 case 3:
353                     enumName = "TOKENIZER_TYPE_RFC822";
354                     break;
355                 default:
356                     throw new ProcessingException(
357                             "Unknown tokenizer type " + annotation.getTokenizerType(),
358                             getterOrField.getElement());
359             }
360         }
361         return CodeBlock.of("\n.setTokenizerType($T.$N)",
362                 StringPropertyAnnotation.CONFIG_CLASS, enumName);
363     }
364 
365     /**
366      * Creates an expr like {@code .setIndexingType(StringPropertyConfig.INDEXING_TYPE_PREFIXES)}.
367      */
createSetIndexingTypeExpr( @onNull StringPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)368     private static @NonNull CodeBlock createSetIndexingTypeExpr(
369             @NonNull StringPropertyAnnotation annotation,
370             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
371         String enumName;
372         switch (annotation.getIndexingType()) {
373             case 0:
374                 enumName = "INDEXING_TYPE_NONE";
375                 break;
376             case 1:
377                 enumName = "INDEXING_TYPE_EXACT_TERMS";
378                 break;
379             case 2:
380                 enumName = "INDEXING_TYPE_PREFIXES";
381                 break;
382             default:
383                 throw new ProcessingException(
384                         "Unknown indexing type " + annotation.getIndexingType(),
385                         getterOrField.getElement());
386         }
387         return CodeBlock.of("\n.setIndexingType($T.$N)",
388                 StringPropertyAnnotation.CONFIG_CLASS, enumName);
389     }
390 
391     /**
392      * Creates an expr like {@code .setShouldIndexNestedProperties(true)}.
393      */
createSetShouldIndexNestedPropertiesExpr( @onNull DocumentPropertyAnnotation annotation)394     private static @NonNull CodeBlock createSetShouldIndexNestedPropertiesExpr(
395             @NonNull DocumentPropertyAnnotation annotation) {
396         return CodeBlock.of("\n.setShouldIndexNestedProperties($L)",
397                 annotation.getShouldIndexNestedProperties());
398     }
399 
400     /**
401      * Creates an expr like {@code .setIndexingType(LongPropertyConfig.INDEXING_TYPE_RANGE)}.
402      */
createSetIndexingTypeExpr( @onNull LongPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)403     private static @NonNull CodeBlock createSetIndexingTypeExpr(
404             @NonNull LongPropertyAnnotation annotation,
405             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
406         String enumName;
407         switch (annotation.getIndexingType()) {
408             case 0:
409                 enumName = "INDEXING_TYPE_NONE";
410                 break;
411             case 1:
412                 enumName = "INDEXING_TYPE_RANGE";
413                 break;
414             default:
415                 throw new ProcessingException(
416                         "Unknown indexing type " + annotation.getIndexingType(),
417                         getterOrField.getElement());
418         }
419         return CodeBlock.of("\n.setIndexingType($T.$N)",
420                 LongPropertyAnnotation.CONFIG_CLASS, enumName);
421     }
422 
423     /**
424      * Creates an expr like
425      * {@code .setIndexingType(EmbeddingPropertyConfig.INDEXING_TYPE_SIMILARITY)}.
426      */
createSetIndexingTypeExpr( @onNull EmbeddingPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)427     private static @NonNull CodeBlock createSetIndexingTypeExpr(
428             @NonNull EmbeddingPropertyAnnotation annotation,
429             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
430         String enumName;
431         switch (annotation.getIndexingType()) {
432             case 0:
433                 enumName = "INDEXING_TYPE_NONE";
434                 break;
435             case 1:
436                 enumName = "INDEXING_TYPE_SIMILARITY";
437                 break;
438             default:
439                 throw new ProcessingException(
440                         "Unknown indexing type " + annotation.getIndexingType(),
441                         getterOrField.getElement());
442         }
443         return CodeBlock.of("\n.setIndexingType($T.$N)",
444                 EmbeddingPropertyAnnotation.CONFIG_CLASS, enumName);
445     }
446 
447     /**
448      * Creates an expr like
449      * {@code .setQuantizationType(EmbeddingPropertyConfig.QUANTIZATION_TYPE_8_BIT)}.
450      */
createSetQuantizationTypeExpr( @onNull EmbeddingPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)451     private static @NonNull CodeBlock createSetQuantizationTypeExpr(
452             @NonNull EmbeddingPropertyAnnotation annotation,
453             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
454         String enumName;
455         switch (annotation.getQuantizationType()) {
456             case 0:
457                 enumName = "QUANTIZATION_TYPE_NONE";
458                 break;
459             case 1:
460                 enumName = "QUANTIZATION_TYPE_8_BIT";
461                 break;
462             default:
463                 throw new ProcessingException(
464                         "Unknown quantization type " + annotation.getQuantizationType(),
465                         getterOrField.getElement());
466         }
467         return CodeBlock.of("\n.setQuantizationType($T.$N)",
468                 EmbeddingPropertyAnnotation.CONFIG_CLASS, enumName);
469     }
470 
471     /**
472      * Creates an expr like
473      * {@code .setJoinableValueType(StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID)}.
474      */
createSetJoinableValueTypeExpr( @onNull StringPropertyAnnotation annotation, @NonNull AnnotatedGetterOrField getterOrField)475     private static @NonNull CodeBlock createSetJoinableValueTypeExpr(
476             @NonNull StringPropertyAnnotation annotation,
477             @NonNull AnnotatedGetterOrField getterOrField) throws ProcessingException {
478         String enumName;
479         AnnotatedGetterOrField.ElementTypeCategory typeCategory =
480                 getterOrField.getElementTypeCategory();
481         switch (annotation.getJoinableValueType()) {
482             case 0:
483                 enumName = "JOINABLE_VALUE_TYPE_NONE";
484                 break;
485             case 1:
486                 switch (typeCategory) {
487                     case COLLECTION: // fall-through
488                     case ARRAY:
489                         throw new ProcessingException(
490                                 "Joinable value type 1 not allowed on repeated properties.",
491                                 getterOrField.getElement());
492                     case SINGLE: // fall-through
493                         break;
494                     default:
495                         throw new IllegalStateException("Unhandled cardinality: " + typeCategory);
496                 }
497                 enumName = "JOINABLE_VALUE_TYPE_QUALIFIED_ID";
498                 break;
499             default:
500                 throw new ProcessingException(
501                         "Unknown joinable value type " + annotation.getJoinableValueType(),
502                         getterOrField.getElement());
503         }
504         return CodeBlock.of("\n.setJoinableValueType($T.$N)",
505                 StringPropertyAnnotation.CONFIG_CLASS, enumName);
506     }
507 }
508