Merge pull request #2997 from square/bquenaudon.2024-06-17.mapoptions

Don't crash with inlined map options
square · Jun 21, 2024 · 8d784ce · 8d784ce
2 parents 102ff34 + e7e19de
commit 8d784ce
Show file tree

Hide file tree

Showing 5 changed files with 375 additions and 39 deletions.
diff --git a/wire-schema/api/wire-schema.api b/wire-schema/api/wire-schema.api
@@ -370,21 +370,6 @@ public final class com/squareup/wire/schema/Field$Label : java/lang/Enum {
 	public static fun values ()[Lcom/squareup/wire/schema/Field$Label;
 }
 
-public final class com/squareup/wire/schema/LinkedOptionEntry {
-	public fun <init> (Lcom/squareup/wire/schema/internal/parser/OptionElement;Lcom/squareup/wire/schema/ProtoMember;Ljava/lang/Object;)V
-	public final fun component1 ()Lcom/squareup/wire/schema/internal/parser/OptionElement;
-	public final fun component2 ()Lcom/squareup/wire/schema/ProtoMember;
-	public final fun component3 ()Ljava/lang/Object;
-	public final fun copy (Lcom/squareup/wire/schema/internal/parser/OptionElement;Lcom/squareup/wire/schema/ProtoMember;Ljava/lang/Object;)Lcom/squareup/wire/schema/LinkedOptionEntry;
-	public static synthetic fun copy$default (Lcom/squareup/wire/schema/LinkedOptionEntry;Lcom/squareup/wire/schema/internal/parser/OptionElement;Lcom/squareup/wire/schema/ProtoMember;Ljava/lang/Object;ILjava/lang/Object;)Lcom/squareup/wire/schema/LinkedOptionEntry;
-	public fun equals (Ljava/lang/Object;)Z
-	public final fun getOptionElement ()Lcom/squareup/wire/schema/internal/parser/OptionElement;
-	public final fun getProtoMember ()Lcom/squareup/wire/schema/ProtoMember;
-	public final fun getValue ()Ljava/lang/Object;
-	public fun hashCode ()I
-	public fun toString ()Ljava/lang/String;
-}
-
 public final class com/squareup/wire/schema/Linker {
 	public fun <init> (Lcom/squareup/wire/schema/Loader;Lcom/squareup/wire/schema/ErrorCollector;ZZ)V
 	public fun <init> (Lcom/squareup/wire/schema/Loader;Lcom/squareup/wire/schema/ErrorCollector;ZZLjava/util/List;)V

diff --git a/wire-schema/api/wire-schema.klib.api b/wire-schema/api/wire-schema.klib.api
@@ -946,22 +946,6 @@ final class com.squareup.wire.schema/Field { // com.squareup.wire.schema/Field|n
     final var type // com.squareup.wire.schema/Field.type|{}type[0]
         final fun <get-type>(): com.squareup.wire.schema/ProtoType? // com.squareup.wire.schema/Field.type.<get-type>|<get-type>(){}[0]
 }
-final class com.squareup.wire.schema/LinkedOptionEntry { // com.squareup.wire.schema/LinkedOptionEntry|null[0]
-    constructor <init>(com.squareup.wire.schema.internal.parser/OptionElement, com.squareup.wire.schema/ProtoMember, kotlin/Any?) // com.squareup.wire.schema/LinkedOptionEntry.<init>|<init>(com.squareup.wire.schema.internal.parser.OptionElement;com.squareup.wire.schema.ProtoMember;kotlin.Any?){}[0]
-    final fun component1(): com.squareup.wire.schema.internal.parser/OptionElement // com.squareup.wire.schema/LinkedOptionEntry.component1|component1(){}[0]
-    final fun component2(): com.squareup.wire.schema/ProtoMember // com.squareup.wire.schema/LinkedOptionEntry.component2|component2(){}[0]
-    final fun component3(): kotlin/Any? // com.squareup.wire.schema/LinkedOptionEntry.component3|component3(){}[0]
-    final fun copy(com.squareup.wire.schema.internal.parser/OptionElement =..., com.squareup.wire.schema/ProtoMember =..., kotlin/Any? =...): com.squareup.wire.schema/LinkedOptionEntry // com.squareup.wire.schema/LinkedOptionEntry.copy|copy(com.squareup.wire.schema.internal.parser.OptionElement;com.squareup.wire.schema.ProtoMember;kotlin.Any?){}[0]
-    final fun equals(kotlin/Any?): kotlin/Boolean // com.squareup.wire.schema/LinkedOptionEntry.equals|equals(kotlin.Any?){}[0]
-    final fun hashCode(): kotlin/Int // com.squareup.wire.schema/LinkedOptionEntry.hashCode|hashCode(){}[0]
-    final fun toString(): kotlin/String // com.squareup.wire.schema/LinkedOptionEntry.toString|toString(){}[0]
-    final val optionElement // com.squareup.wire.schema/LinkedOptionEntry.optionElement|{}optionElement[0]
-        final fun <get-optionElement>(): com.squareup.wire.schema.internal.parser/OptionElement // com.squareup.wire.schema/LinkedOptionEntry.optionElement.<get-optionElement>|<get-optionElement>(){}[0]
-    final val protoMember // com.squareup.wire.schema/LinkedOptionEntry.protoMember|{}protoMember[0]
-        final fun <get-protoMember>(): com.squareup.wire.schema/ProtoMember // com.squareup.wire.schema/LinkedOptionEntry.protoMember.<get-protoMember>|<get-protoMember>(){}[0]
-    final val value // com.squareup.wire.schema/LinkedOptionEntry.value|{}value[0]
-        final fun <get-value>(): kotlin/Any? // com.squareup.wire.schema/LinkedOptionEntry.value.<get-value>|<get-value>(){}[0]
-}
 final class com.squareup.wire.schema/Linker { // com.squareup.wire.schema/Linker|null[0]
     constructor <init>(com.squareup.wire.schema/Loader, com.squareup.wire.schema/ErrorCollector, kotlin/Boolean, kotlin/Boolean) // com.squareup.wire.schema/Linker.<init>|<init>(com.squareup.wire.schema.Loader;com.squareup.wire.schema.ErrorCollector;kotlin.Boolean;kotlin.Boolean){}[0]
     constructor <init>(com.squareup.wire.schema/Loader, com.squareup.wire.schema/ErrorCollector, kotlin/Boolean, kotlin/Boolean, kotlin.collections/List<com.squareup.wire.schema/ProtoType>) // com.squareup.wire.schema/Linker.<init>|<init>(com.squareup.wire.schema.Loader;com.squareup.wire.schema.ErrorCollector;kotlin.Boolean;kotlin.Boolean;kotlin.collections.List<com.squareup.wire.schema.ProtoType>){}[0]

diff --git a/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/LinkedOptionEntry.kt b/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/LinkedOptionEntry.kt
@@ -17,7 +17,7 @@ package com.squareup.wire.schema
 
 import com.squareup.wire.schema.internal.parser.OptionElement
 
-data class LinkedOptionEntry(
+internal data class LinkedOptionEntry(
   val optionElement: OptionElement,
   val protoMember: ProtoMember,
   val value: Any?,

diff --git a/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/Options.kt b/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/Options.kt
@@ -34,7 +34,12 @@ class Options(
 
   val elements: List<OptionElement>
     get() {
-      return entries?.map { it.optionElement } ?: optionElements
+      return entries?.map {
+        // TODO(Benoit) this property is used to go from `Options` to `List<OptionElement>` but this
+        //  doesn't take into account what has been pruned. We should consume `it.value`  somehow and
+        //  select the option elements we are to fetch, or not.
+        it.optionElement
+      } ?: optionElements
     }
 
   val map: Map<ProtoMember, Any?>
@@ -322,7 +327,16 @@ class Options(
     when (o) {
       is Map<*, *> -> {
         for ((key, value) in o) {
-          val protoMember = key as ProtoMember
+          val protoMember = when (key) {
+            is ProtoMember -> key
+            else -> {
+              // When the key isn't a `ProtoMember`, this key/value pair is a inlined value of a map
+              // field. We don't need to track the key type in map fields for they are always of
+              // scalar types. We however have to check the value type.
+              gatherFields(sink, type, value!!, pruningRules)
+              continue
+            }
+          }
           if (pruningRules.prunes(protoMember)) continue
           sink.getOrPut(type, ::ArrayList).add(protoMember)
           gatherFields(sink, protoMember.type, value!!, pruningRules)
@@ -365,7 +379,21 @@ class Options(
       o is Map<*, *> -> {
         val map = mutableMapOf<ProtoMember, Any>()
         for ((key, value) in o) {
-          val protoMember = key as ProtoMember
+          val protoMember = when (key) {
+            is ProtoMember -> key
+            else -> {
+              // When the key isn't a `ProtoMember`, this key/value pair is a inlined value of a map
+              // field.
+              val retainedValue = retainAll(schema, markSet, type, value!!)
+              // if `retainedValue` is a map, its value represents an inline message, and we need to
+              // mark the proto member.
+              if (retainedValue is Map<*, *>) {
+                val (k, v) = retainedValue.entries.single()
+                map[k as ProtoMember] = v!!
+              }
+              continue
+            }
+          }
           val isCoreMemberOfGoogleProtobuf =
             protoMember.type in GOOGLE_PROTOBUF_OPTION_TYPES &&
               !schema.isExtensionField(protoMember)