Read order of fields from the .class file

This CL fixes a bug in KSP where the order of fields returned in a
KSClassDeclaration does not necessarily match the order they were
declared in the class. It only impacts compiled kotlin classes because
Kotlin compiler orders field names when serializing the metadata.
https://github.com/google/ksp/issues/250

This particularly hurts Room as we generate the table structure based on
this order. Even though Room is clever enough to consider the schema the
same, the effects are still visible in both schema files and generated
sql database.

This CL is a port of a possible implementation in KSP:
https://github.com/google/ksp/pull/260
Unfortunately, merging that CL would mean KSP doing a lot more work for
each class. The plan there is to first fix it in the kotlin compiler if
possible and if not, add an API to get sorted fields (to avoid the cost
when not necessary). Meanwhile, this CL is a band aid solution in Room
so that developers won't have unexpected schema file changes when they
move to KSP from KAPT.

As this solution is best effort, it doesn't fail compilation if it fails
to read fields. I've added a strict mode to XProcessing so that we can
run it in tests with a requirement to always find the order but on
developer's machine, it won't require this. This strict mode can be
turned on via a system property, which we can use to collect more bug
information when it doesn't work.

Fixes: 177471948
Test: OrderOfFieldsTest
Change-Id: If0cf346a7f92dbbf2094e0076ea9ed02de4e32fc
diff --git a/room/compiler-processing/build.gradle b/room/compiler-processing/build.gradle
index d2365d4..856793a 100644
--- a/room/compiler-processing/build.gradle
+++ b/room/compiler-processing/build.gradle
@@ -53,6 +53,10 @@
     }
 }
 
+tasks.withType(Test).configureEach {
+   it.systemProperty("androidx.room.compiler.processing.strict", "true")
+}
+
 androidx {
     name = "AndroidX Room XProcessor"
     type = LibraryType.ANNOTATION_PROCESSOR
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingConfig.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingConfig.kt
new file mode 100644
index 0000000..25fed28
--- /dev/null
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingConfig.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.room.compiler.processing
+
+/**
+ * Utility class to change some behavior in tests, like adding more strict tests.
+ */
+internal object XProcessingConfig {
+    /**
+     * When true, we do more strict checks and fail instead of workarounds or fallback
+     * behaviors. Set to true in room's own tests.
+     */
+    val STRICT_MODE by lazy {
+        System.getProperty("$PROP_PREFIX.strict").toBoolean()
+    }
+
+    private const val PROP_PREFIX = "androidx.room.compiler.processing"
+}
\ No newline at end of file
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFieldOrdering.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFieldOrdering.kt
new file mode 100644
index 0000000..aa3c1c2
--- /dev/null
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFieldOrdering.kt
@@ -0,0 +1,214 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.room.compiler.processing.ksp
+
+import androidx.room.compiler.processing.XFieldElement
+import androidx.room.compiler.processing.XProcessingConfig
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.Origin
+import java.lang.reflect.InvocationHandler
+import java.lang.reflect.Method
+import java.lang.reflect.Proxy
+
+/**
+ * When a compiled kotlin class is loaded from a `.class` file, its fields are not ordered in the
+ * same way they are declared in code.
+ * This particularly hurts Room where we generate the table structure in that order.
+ *
+ * This class implements a port of https://github.com/google/ksp/pull/260 via reflection until KSP
+ * (or kotlin compiler) fixes the problem. As this uses reflection, it is fail safe such that if it
+ * cannot find the correct order, it will just return in the order KSP returned instead of crashing.
+ */
+internal object KspFieldOrdering {
+    /**
+     * Sorts the given fields in the order they are declared in the backing class declaration.
+     */
+    fun orderFields(
+        owner: KSClassDeclaration,
+        fields: List<KspFieldElement>
+    ): List<KspFieldElement> {
+        // no reason to try to load .class if we don't have any fields to sort
+        if (fields.isEmpty()) return fields
+        val comparator = getFieldNamesComparator(owner)
+        return if (comparator == null) {
+            fields
+        } else {
+            fields.forEach {
+                // make sure each name gets registered so that if we didn't find it in .class for
+                // whatever reason, we keep the order given from KSP.
+                comparator.register(it.name)
+            }
+            fields.sortedWith(comparator)
+        }
+    }
+
+    /**
+     * Builds a field names comparator from the given class declaration if and only if its origin
+     * is CLASS.
+     * If it fails to find the order, returns null.
+     */
+    @Suppress("BanUncheckedReflection")
+    private fun getFieldNamesComparator(
+        ksClassDeclaration: KSClassDeclaration
+    ): FieldNameComparator? {
+        return try {
+            if (ksClassDeclaration.origin != Origin.CLASS) return null
+            val typeReferences = ReflectionReferences.getInstance(ksClassDeclaration) ?: return null
+            val descriptor = typeReferences.getDescriptorMethod.invoke(ksClassDeclaration)
+                ?: return null
+            if (!typeReferences.deserializedClassDescriptor.isInstance(descriptor)) {
+                return null
+            }
+            val descriptorSrc = typeReferences.descriptorSourceMethod.invoke(descriptor)
+                ?: return null
+            if (!typeReferences.kotlinJvmBinarySourceElement.isInstance(descriptorSrc)) {
+                return null
+            }
+            val binarySource = typeReferences.binaryClassMethod.invoke(descriptorSrc)
+                ?: return null
+
+            val fieldNameComparator = FieldNameComparator()
+            val invocationHandler = InvocationHandler { _, method, args ->
+                if (method.name == "visitField") {
+                    val nameAsString = typeReferences.asStringMethod.invoke(args[0])
+                    if (nameAsString is String) {
+                        fieldNameComparator.register(nameAsString)
+                    }
+                }
+                null
+            }
+
+            val proxy = Proxy.newProxyInstance(
+                ksClassDeclaration.javaClass.classLoader,
+                arrayOf(typeReferences.memberVisitor),
+                invocationHandler
+            )
+            typeReferences.visitMembersMethod.invoke(binarySource, proxy, null)
+            fieldNameComparator.seal()
+            fieldNameComparator
+        } catch (ignored: Throwable) {
+            // this is best effort, if it failed, just ignore
+            if (XProcessingConfig.STRICT_MODE) {
+                throw RuntimeException("failed to get fields", ignored)
+            }
+            null
+        }
+    }
+
+    /**
+     * Holder object to keep references to class/method instances.
+     */
+    private class ReflectionReferences private constructor(
+        classLoader: ClassLoader
+    ) {
+
+        val deserializedClassDescriptor: Class<*> = classLoader.loadClass(
+            "org.jetbrains.kotlin.serialization.deserialization.descriptors" +
+                ".DeserializedClassDescriptor"
+        )
+
+        val ksClassDeclarationDescriptorImpl: Class<*> = classLoader.loadClass(
+            "com.google.devtools.ksp.symbol.impl.binary.KSClassDeclarationDescriptorImpl"
+        )
+        val kotlinJvmBinarySourceElement: Class<*> = classLoader.loadClass(
+            "org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement"
+        )
+
+        val kotlinJvmBinaryClass: Class<*> = classLoader.loadClass(
+            "org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryClass"
+        )
+
+        val memberVisitor: Class<*> = classLoader.loadClass(
+            "org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryClass\$MemberVisitor"
+        )
+
+        val name: Class<*> = classLoader.loadClass(
+            "org.jetbrains.kotlin.name.Name"
+        )
+
+        val getDescriptorMethod: Method = ksClassDeclarationDescriptorImpl
+            .getDeclaredMethod("getDescriptor")
+
+        val descriptorSourceMethod: Method = deserializedClassDescriptor.getMethod("getSource")
+
+        val binaryClassMethod: Method = kotlinJvmBinarySourceElement.getMethod("getBinaryClass")
+
+        val visitMembersMethod: Method = kotlinJvmBinaryClass.getDeclaredMethod(
+            "visitMembers",
+            memberVisitor, ByteArray::class.java
+        )
+
+        val asStringMethod: Method = name.getDeclaredMethod("asString")
+
+        companion object {
+            private val FAILED = Any()
+            private var instance: Any? = null
+
+            /**
+             * Gets the cached instance or create a new one using the class loader of the given
+             * [ref] parameter.
+             */
+            fun getInstance(ref: Any): ReflectionReferences? {
+                if (instance == null) {
+                    instance = try {
+                        ReflectionReferences(ref::class.java.classLoader)
+                    } catch (ignored: Throwable) {
+                        FAILED
+                    }
+                }
+                return instance as? ReflectionReferences
+            }
+        }
+    }
+
+    private class FieldNameComparator : Comparator<XFieldElement> {
+        private var nextOrder: Int = 0
+        private var sealed: Boolean = false
+        private val orders = mutableMapOf<String, Int>()
+
+        /**
+         * Called when fields are read to lock the ordering.
+         * This is only relevant in tests as at runtime, we just do a best effort (add a new id
+         * for it) and continue.
+         */
+        fun seal() {
+            sealed = true
+        }
+
+        /**
+         * Registers the name with the next order id
+         */
+        fun register(name: String) {
+            getOrder(name)
+        }
+
+        /**
+         * Gets the order of the name. If it was not seen before, adds it to the list, giving it a
+         * new ID.
+         */
+        private fun getOrder(name: String) = orders.getOrPut(name) {
+            if (sealed && XProcessingConfig.STRICT_MODE) {
+                error("expected to find field $name but it is non-existent")
+            }
+            nextOrder++
+        }
+
+        override fun compare(field1: XFieldElement, field2: XFieldElement): Int {
+            return getOrder(field1.name).compareTo(getOrder(field2.name))
+        }
+    }
+}
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspTypeElement.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspTypeElement.kt
index b2ffecd..763a3fa 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspTypeElement.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspTypeElement.kt
@@ -108,13 +108,23 @@
      */
     private val _declaredProperties by lazy {
         val declaredProperties = declaration.getDeclaredProperties()
+            .map {
+                KspFieldElement(
+                    env = env,
+                    declaration = it,
+                    containing = this
+                )
+            }.let {
+                // only order instance fields, we don't care about the order of companion fields.
+                KspFieldOrdering.orderFields(declaration, it)
+            }
+
         val companionProperties = declaration
             .findCompanionObject()
             ?.getDeclaredProperties()
             ?.filter {
                 it.isStatic()
             }.orEmpty()
-        (declaredProperties + companionProperties)
             .map {
                 KspFieldElement(
                     env = env,
@@ -122,6 +132,7 @@
                     containing = this
                 )
             }
+        declaredProperties + companionProperties
     }
 
     private val _declaredFieldsIncludingSupers by lazy {
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OrderOfFieldsTest.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OrderOfFieldsTest.kt
new file mode 100644
index 0000000..24ce1ee
--- /dev/null
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OrderOfFieldsTest.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.room.compiler.processing
+
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.compileFiles
+import androidx.room.compiler.processing.util.getAllFieldNames
+import androidx.room.compiler.processing.util.runProcessorTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+
+/**
+ * see: https://github.com/google/ksp/issues/250
+ */
+class OrderOfFieldsTest {
+    @Test
+    fun outOfOrderKotlin() {
+        val libSource = Source.kotlin(
+            "lib.kt",
+            """
+            class KotlinClass {
+                val b: String = TODO()
+                val a: String = TODO()
+                val c: String = TODO()
+                val isB:String = TODO()
+                val isA:String = TODO()
+                val isC:String = TODO()
+            }
+            """.trimIndent()
+        )
+        val classpath = compileFiles(listOf(libSource))
+        runProcessorTest(
+            sources = emptyList(),
+            classpath = listOf(classpath)
+        ) { invocation ->
+            val element = invocation.processingEnv.requireTypeElement("KotlinClass")
+            assertThat(element.getAllFieldNames())
+                .containsExactly("b", "a", "c", "isB", "isA", "isC")
+                .inOrder()
+        }
+    }
+}
\ No newline at end of file
diff --git a/room/compiler/build.gradle b/room/compiler/build.gradle
index 86d8b57..ed7f3ea 100644
--- a/room/compiler/build.gradle
+++ b/room/compiler/build.gradle
@@ -250,6 +250,10 @@
     }
 }
 
+tasks.withType(Test).configureEach {
+    it.systemProperty("androidx.room.compiler.processing.strict", "true")
+}
+
 androidx {
     name = "Android Room Compiler"
     type = LibraryType.COMPILER_PLUGIN