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