Change MutableStructFields to Implement MutableCollection

Implementing MutableCollection instead of Iterator will allow extra
functionality built into MutableCollection. To simplify implementing
`iterator()` on MutableStructFields the internal backing field was also changed
to MutableMap<String, MutableList<StructField>> at the expense of slightly more
overhead because the underlying struct's field name is stored in the map key as
well as value. This can be optimized in the future if desired.

api/IonElement.api

@@ -577,19 +577,19 @@ public abstract interface class com/amazon/ionelement/api/LobElement : com/amazo
 	public abstract fun withoutMetas ()Lcom/amazon/ionelement/api/LobElement;
 }
 
-public abstract interface class com/amazon/ionelement/api/MutableStructFields : java/lang/Iterable, kotlin/jvm/internal/markers/KMappedMarker {
-	public abstract fun add (Lcom/amazon/ionelement/api/StructField;)Lcom/amazon/ionelement/api/MutableStructFields;
-	public abstract fun add (Ljava/lang/String;Lcom/amazon/ionelement/api/IonElement;)Lcom/amazon/ionelement/api/MutableStructFields;
+public abstract interface class com/amazon/ionelement/api/MutableStructFields : java/util/Collection, kotlin/jvm/internal/markers/KMutableCollection {
+	public abstract fun add (Lcom/amazon/ionelement/api/StructField;)Z
+	public abstract fun add (Ljava/lang/String;Lcom/amazon/ionelement/api/IonElement;)Z
+	public abstract fun clearField (Ljava/lang/String;)Z
 	public abstract fun containsField (Ljava/lang/String;)Z
 	public abstract fun get (Ljava/lang/String;)Lcom/amazon/ionelement/api/AnyElement;
 	public abstract fun getAll (Ljava/lang/String;)Ljava/util/Collection;
 	public abstract fun getOptional (Ljava/lang/String;)Lcom/amazon/ionelement/api/AnyElement;
-	public abstract fun plusAssign (Lcom/amazon/ionelement/api/StructField;)V
-	public abstract fun plusAssign (Ljava/lang/Iterable;)V
-	public abstract fun remove (Lcom/amazon/ionelement/api/StructField;)Lcom/amazon/ionelement/api/MutableStructFields;
-	public abstract fun removeAll (Ljava/lang/String;)Lcom/amazon/ionelement/api/MutableStructFields;
-	public abstract fun set (Ljava/lang/String;Lcom/amazon/ionelement/api/IonElement;)Lcom/amazon/ionelement/api/MutableStructFields;
-	public abstract fun setAll (Ljava/lang/Iterable;)Lcom/amazon/ionelement/api/MutableStructFields;
+	public abstract fun plusAssign (Ljava/util/Collection;)V
+	public abstract fun remove (Lcom/amazon/ionelement/api/StructField;)Z
+	public abstract fun removeAll (Ljava/util/Collection;)Z
+	public abstract fun set (Ljava/lang/String;Lcom/amazon/ionelement/api/IonElement;)V
+	public abstract fun setAll (Ljava/lang/Iterable;)V
 }
 
 public abstract interface class com/amazon/ionelement/api/SeqElement : com/amazon/ionelement/api/ContainerElement {

src/com/amazon/ionelement/api/MutableStructFields.kt

@@ -3,7 +3,7 @@ package com.amazon.ionelement.api
 /**
  * Represents a mutable view of a collection of struct fields.
  */
-public interface MutableStructFields : Iterable<StructField> {
+public interface MutableStructFields : MutableCollection<StructField> {
 
     /**
      * Retrieves the value of the first field found with the specified name.
@@ -26,34 +26,44 @@ public interface MutableStructFields : Iterable<StructField> {
     /**
      * If one or more fields with the specified name already exists, this replaces all of them with the value provided.
      *
-     * Otherwise, a new field with the given name and value is added to the collection
+     * Otherwise, a new field with the given name and value is added to the collection.
+     *
+     * Returns true if a field was replaced.
      */
-    public operator fun set(fieldName: String, value: IonElement): MutableStructFields
+    public operator fun set(fieldName: String, value: IonElement)
 
     /**
      * Adds all the given fields to the collection. For existing fields with the same names as the fields provided, all
      * instances of those fields will be removed.
      */
-    public fun setAll(fields: Iterable<StructField>): MutableStructFields
+    public fun setAll(fields: Iterable<StructField>)
 
     /**
      * Adds a new field to the collection with the given name and value. The collection may have multiple fields with
      * the same name.
      */
-    public fun add(fieldName: String, value: IonElement): MutableStructFields
+    public fun add(fieldName: String, value: IonElement): Boolean
 
-    /** Adds the given field to the collection. The collection may have multiple fields with the same name. */
-    public fun add(field: StructField): MutableStructFields
+    /** Adds the given field to the collection. The collection may have multiple fields with the same name.
+     *
+     * Repeated fields are allowed, so this will always return true. */
+    public override fun add(element: StructField): Boolean
 
-    /** Removes a random occurrence of a field the matches the given field, or does nothing if no field exists */
-    public fun remove(field: StructField): MutableStructFields
+    /** Removes a random occurrence of a field the matches the given field, or does nothing if no field exists.
+     *
+     * Returns true is a field was removed. */
+    public override fun remove(element: StructField): Boolean
 
-    /** Removes all fields found with the given name or does nothing if no fields exist */
-    public fun removeAll(fieldName: String): MutableStructFields
+    /** Removes all occurrence of a field the matches the given field name, or does nothing if no field exists.
+     *
+     * Returns true if a field was removed. */
+    public fun clearField(fieldName: String): Boolean
 
-    /** Adds the given field to the collection */
-    public operator fun plusAssign(field: StructField)
+    /** Removes all of this collection's elements that are also contained in the specified collection.
+     *
+     *  At most one field per element in the give collection is removed. */
+    public override fun removeAll(elements: Collection<StructField>): Boolean
 
     /** Adds all the given fields to the collection */
-    public operator fun plusAssign(fields: Iterable<StructField>)
+    public operator fun plusAssign(fields: Collection<StructField>)
 }

src/com/amazon/ionelement/impl/MutableStructFieldsImpl.kt

@@ -5,82 +5,173 @@ import com.amazon.ionelement.api.IonElement
 import com.amazon.ionelement.api.MutableStructFields
 import com.amazon.ionelement.api.StructField
 import com.amazon.ionelement.api.field
+import java.util.NoSuchElementException
 
-internal class MutableStructFieldsImpl(private val fields: MutableMap<String, MutableList<AnyElement>>) :
+internal class MutableStructFieldsImpl(private val fields: MutableMap<String, MutableList<StructField>>) :
     MutableStructFields {
+    override val size: Int
+        get() = fields.values.sumBy { it.size }
 
     override fun get(fieldName: String): AnyElement {
-        return requireNotNull(fields[fieldName]?.firstOrNull()) {
+        return requireNotNull(fields[fieldName]?.firstOrNull()?.value) {
             "Required struct field '$fieldName' missing"
         }
     }
 
     override fun getOptional(fieldName: String): AnyElement? {
-        return fields[fieldName]?.firstOrNull()
+        return fields[fieldName]?.firstOrNull()?.value
     }
 
     override fun getAll(fieldName: String): Collection<AnyElement> {
-        return fields[fieldName] ?: mutableListOf()
+        return fields[fieldName]?.map { it.value } ?: mutableListOf()
     }
 
     override fun containsField(fieldName: String): Boolean {
         return fieldName in fields
     }
 
-    override fun set(fieldName: String, value: IonElement): MutableStructFields {
+    override fun set(fieldName: String, value: IonElement) {
         val values = fields.getOrPut(fieldName, ::mutableListOf)
         values.clear()
-        values.add(value as AnyElement)
-        return this
+        values.add(field(fieldName, value))
     }
 
-    override fun setAll(fields: Iterable<StructField>): MutableStructFields {
+    override fun setAll(fields: Iterable<StructField>) {
         fields.groupByTo(mutableMapOf(), { it.name }, { it.value }).forEach { (name, values) ->
-            this.fields[name] = values
+            this.fields[name] = values.map { field(name, it) }.toMutableList()
         }
-        return this
     }
 
-    override fun add(fieldName: String, value: IonElement): MutableStructFields {
+    override fun add(fieldName: String, value: IonElement): Boolean {
         value as AnyElement
         val values = fields[fieldName]
         if (values == null) {
-            fields[fieldName] = mutableListOf(value)
+            fields[fieldName] = mutableListOf(field(fieldName, value))
         } else {
-            values.add(value)
+            values.add(field(fieldName, value))
         }
+        return true
+    }
+
+    override fun add(element: StructField): Boolean {
+        return add(element.name, element.value)
+    }
+
+    override fun remove(element: StructField): Boolean {
+        return fields[element.name]?.remove(element) ?: false
+    }
+
+    override fun clearField(fieldName: String): Boolean {
+        fields[fieldName]?.clear() ?: return false
+        return true
+    }
+
+    override fun removeAll(elements: Collection<StructField>): Boolean {
+        var modified = false
+        for (element in elements) {
+            modified = remove(element) || modified
+        }
+        return modified
+    }
+
+    override fun plusAssign(fields: Collection<StructField>) {
+        addAll(fields)
+    }
+
+    override fun addAll(elements: Collection<StructField>): Boolean {
+        elements.forEach { add(it) }
+        return true
+    }
+
+    override fun clear() {
+        fields.clear()
+    }
+
+    override fun contains(element: StructField): Boolean {
+        return fields[element.name]?.contains(element) ?: false
+    }
+
+    override fun containsAll(elements: Collection<StructField>): Boolean {
+        return elements.all { contains(it) }
+    }
 
-        return this
+    override fun isEmpty(): Boolean {
+        return fields.all { it.value.isEmpty() }
     }
 
-    override fun add(field: StructField): MutableStructFields {
-        add(field.name, field.value)
-        return this
+    override fun iterator(): MutableIterator<StructField> {
+        return MutableStructFieldsIterator(fields)
     }
 
-    override fun remove(field: StructField): MutableStructFields {
-        fields[field.name]?.remove(field.value)
-        return this
+    override fun retainAll(elements: Collection<StructField>): Boolean {
+        var modified = false
+        val it = iterator()
+        while (it.hasNext()) {
+            val field = it.next()
+            if (field !in elements) {
+                it.remove()
+                modified = true
+            }
+        }
+        return modified
+    }
+}
+
+internal class MutableStructFieldsIterator(fieldsMap: MutableMap<String, MutableList<StructField>>) :
+    MutableIterator<StructField> {
+
+    private val mapIterator = fieldsMap.iterator()
+    private var listIterator: MutableIterator<StructField> =
+        if (mapIterator.hasNext()) {
+            mapIterator.next().value.iterator()
+        } else {
+            EMPTY_ITERATOR
+        }
+
+    override fun hasNext(): Boolean {
+        if (listIterator.hasNext()) {
+            return true
+        }
+
+        nextFieldList()
+
+        return listIterator.hasNext()
     }
 
-    override fun removeAll(fieldName: String): MutableStructFields {
-        fields[fieldName]?.clear()
-        return this
+    override fun next(): StructField {
+        nextFieldList()
+
+        return listIterator.next()
     }
 
-    override fun plusAssign(field: StructField) {
-        add(field.name, field.value)
+    override fun remove() {
+        listIterator.remove()
     }
 
-    override fun plusAssign(fields: Iterable<StructField>) {
-        fields.forEach { add(it) }
+    private fun nextFieldList() {
+        while (!listIterator.hasNext()) {
+            if (mapIterator.hasNext()) {
+                listIterator = mapIterator.next().value.iterator()
+            } else {
+                listIterator = EMPTY_ITERATOR
+                break
+            }
+        }
     }
 
-    override fun iterator(): Iterator<StructField> {
-        return fields.flatMap { (fieldName, values) ->
-            values.map { value ->
-                field(fieldName, value)
+    companion object {
+        private val EMPTY_ITERATOR = object : MutableIterator<StructField> {
+            override fun hasNext(): Boolean {
+                return false
             }
-        }.iterator()
+
+            override fun next(): StructField {
+                throw NoSuchElementException()
+            }
+
+            override fun remove() {
+                throw IllegalStateException()
+            }
+        }
     }
 }

src/com/amazon/ionelement/impl/StructElementImpl.kt

@@ -69,8 +69,12 @@ internal class StructElementImpl(
         }
 
     override fun mutableFields(): MutableStructFields {
-        val internalMap = mutableMapOf<String, MutableList<AnyElement>>()
-        return MutableStructFieldsImpl(fieldsByName.mapValuesTo(internalMap) { it.value.toMutableList() })
+        val internalMap = mutableMapOf<String, MutableList<StructField>>()
+        return MutableStructFieldsImpl(
+            fieldsByName.mapValuesTo(internalMap) { (name, values) ->
+                values.map { field(name, it) }.toMutableList()
+            }
+        )
     }
 
     override fun update(body: MutableStructFields.() -> Unit): StructElement {
@@ -92,9 +96,11 @@ internal class StructElementImpl(
     override fun copy(annotations: List<String>, metas: MetaContainer): StructElementImpl =
         StructElementImpl(allFields, annotations.toPersistentList(), metas.toPersistentMap())
 
-    override fun withAnnotations(vararg additionalAnnotations: String): StructElementImpl = _withAnnotations(*additionalAnnotations)
+    override fun withAnnotations(vararg additionalAnnotations: String): StructElementImpl =
+        _withAnnotations(*additionalAnnotations)
 
-    override fun withAnnotations(additionalAnnotations: Iterable<String>): StructElementImpl = _withAnnotations(additionalAnnotations)
+    override fun withAnnotations(additionalAnnotations: Iterable<String>): StructElementImpl =
+        _withAnnotations(additionalAnnotations)
 
     override fun withoutAnnotations(): StructElementImpl = _withoutAnnotations()
     override fun withMetas(additionalMetas: MetaContainer): StructElementImpl = _withMetas(additionalMetas)