1/*
2 * Copyright (C) 2016 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
17package androidx.room
18
19import androidx.room.processor.Context
20import androidx.room.processor.DatabaseProcessor
21import androidx.room.processor.ProcessorErrors
22import androidx.room.vo.DaoMethod
23import androidx.room.vo.Warning
24import androidx.room.writer.DaoWriter
25import androidx.room.writer.DatabaseWriter
26import com.google.auto.common.BasicAnnotationProcessor
27import com.google.auto.common.MoreElements
28import com.google.common.collect.SetMultimap
29import java.io.File
30import javax.lang.model.SourceVersion
31import javax.lang.model.element.Element
32
33/**
34 * The annotation processor for Room.
35 */
36class RoomProcessor : BasicAnnotationProcessor() {
37    override fun initSteps(): MutableIterable<ProcessingStep>? {
38        val context = Context(processingEnv)
39        return arrayListOf(DatabaseProcessingStep(context))
40    }
41
42    override fun getSupportedOptions(): MutableSet<String> {
43        return Context.ARG_OPTIONS.toMutableSet()
44    }
45
46    override fun getSupportedSourceVersion(): SourceVersion {
47        return SourceVersion.latest()
48    }
49
50    class DatabaseProcessingStep(context: Context) : ContextBoundProcessingStep(context) {
51        override fun process(
52                elementsByAnnotation: SetMultimap<Class<out Annotation>, Element>
53        ): MutableSet<Element> {
54            // TODO multi step support
55            val databases = elementsByAnnotation[Database::class.java]
56                    ?.map {
57                        DatabaseProcessor(context, MoreElements.asType(it)).process()
58                    }
59            val allDaoMethods = databases?.flatMap { it.daoMethods }
60            allDaoMethods?.let {
61                prepareDaosForWriting(databases, it)
62                it.forEach {
63                    DaoWriter(it.dao, context.processingEnv).write(context.processingEnv)
64                }
65            }
66
67            databases?.forEach { db ->
68                DatabaseWriter(db).write(context.processingEnv)
69                if (db.exportSchema) {
70                    val schemaOutFolder = context.schemaOutFolder
71                    if (schemaOutFolder == null) {
72                        context.logger.w(Warning.MISSING_SCHEMA_LOCATION, db.element,
73                                ProcessorErrors.MISSING_SCHEMA_EXPORT_DIRECTORY)
74                    } else {
75                        if (!schemaOutFolder.exists()) {
76                            schemaOutFolder.mkdirs()
77                        }
78                        val qName = db.element.qualifiedName.toString()
79                        val dbSchemaFolder = File(schemaOutFolder, qName)
80                        if (!dbSchemaFolder.exists()) {
81                            dbSchemaFolder.mkdirs()
82                        }
83                        db.exportSchema(File(dbSchemaFolder, "${db.version}.json"))
84                    }
85                }
86            }
87            return mutableSetOf()
88        }
89        override fun annotations(): MutableSet<out Class<out Annotation>> {
90            return mutableSetOf(Database::class.java, Dao::class.java, Entity::class.java)
91        }
92
93        /**
94         * Traverses all dao methods and assigns them suffix if they are used in multiple databases.
95         */
96        private fun prepareDaosForWriting(
97                databases: List<androidx.room.vo.Database>,
98                daoMethods: List<DaoMethod>) {
99            daoMethods.groupBy { it.dao.typeName }
100                    // if used only in 1 database, nothing to do.
101                    .filter { entry -> entry.value.size > 1 }
102                    .forEach { entry ->
103                        entry.value.groupBy { daoMethod ->
104                            // first suffix guess: Database's simple name
105                            val db = databases.first { db -> db.daoMethods.contains(daoMethod) }
106                            db.typeName.simpleName()
107                        }.forEach { (dbName, methods) ->
108                            if (methods.size == 1) {
109                                //good, db names do not clash, use db name as suffix
110                                methods.first().dao.setSuffix(dbName)
111                            } else {
112                                // ok looks like a dao is used in 2 different databases both of
113                                // which have the same name. enumerate.
114                                methods.forEachIndexed { index, method ->
115                                    method.dao.setSuffix("${dbName}_$index")
116                                }
117                            }
118                        }
119                    }
120        }
121    }
122
123    abstract class ContextBoundProcessingStep(val context: Context) : ProcessingStep
124}
125