From 7a9ee6d06e9cfdaa530310f5d9be9f5a52680d3b Mon Sep 17 00:00:00 2001
From: Matthew Pope <popematt@amazon.com>
Date: Mon, 19 Feb 2024 17:12:31 -0800
Subject: [PATCH] Fix plusType for schema that contains only ISL 2.0 version
 marker

---
 .../ionschema/internal/SchemaImpl_2_0.kt      |  7 +-
 .../ionschema/internal/SchemaImpl_2_0_Test.kt | 68 +++++++++++++++++++
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt b/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt
index 24663add..11567001 100644
--- a/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt
+++ b/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt
@@ -327,14 +327,15 @@ internal class SchemaImpl_2_0 internal constructor(
 
         val newIsl = mutableListOf<IonValue>()
         var newTypeAdded = false
-        isl.forEachIndexed { idx, value ->
+        isl.forEach { value ->
             if (!newTypeAdded) {
                 if (isType(value) && (value["name"] as? IonSymbol)?.stringValue() == type.name) {
                     // new type replaces existing type of the same name
                     newIsl.add(type.isl.clone())
                     newTypeAdded = true
-                    return@forEachIndexed
-                } else if (value.hasTypeAnnotation("schema_footer") || idx == isl.lastIndex) {
+                    return@forEach
+                } else if (value.hasTypeAnnotation("schema_footer")) {
+                    // Insert the new type right before the footer
                     newIsl.add(type.isl.clone())
                     newTypeAdded = true
                 }
diff --git a/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt b/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt
index 34c159c1..b4f269f6 100644
--- a/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt
+++ b/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt
@@ -2,6 +2,7 @@ package com.amazon.ionschema.internal
 
 import com.amazon.ion.IonStruct
 import com.amazon.ion.IonValue
+import com.amazon.ionschema.ION
 import com.amazon.ionschema.IonSchemaSystemBuilder
 import com.amazon.ionschema.IonSchemaTests
 import com.amazon.ionschema.IonSchemaVersion
@@ -65,6 +66,73 @@ class SchemaImpl_2_0_Test {
         assertEquals(expectedNewFooTypeIsl, newFooTypeIsl)
     }
 
+    @Test
+    fun `plusType(type) should correctly insert new type when a schema footer is present`() {
+        val schema = ISS.newSchema(
+            """
+            ${'$'}ion_schema_2_0
+            schema_header::{}
+            type::{ name: foo }
+            schema_footer::{}
+        """
+        )
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(setOf("foo"), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("foo", "bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+    }
+
+    @Test
+    fun `plusType(type) should correctly insert a new type into a schema that is empty aside from its version marker`() {
+        val schema = ISS.newSchema("\$ion_schema_2_0")
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(emptySet<Nothing>(), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+    }
+
+    @Test
+    fun `plusType(type) should preserve top-level open content in a schema`() {
+        val schema = ISS.newSchema(
+            """
+            ${'$'}ion_schema_2_0
+            abc
+            type::{ name: foo }
+            def
+        """
+        )
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(setOf("foo"), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("foo", "bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+
+        val expectedNewSchemaIsl = ION.loader.load(
+            """
+            ${'$'}ion_schema_2_0
+            abc
+            type::{ name: foo }
+            def
+            type::{ name: bar }
+        """
+        )
+        assertEquals(newSchema.isl, expectedNewSchemaIsl)
+    }
+
     @Test
     fun `getType(name) should return a declared type`() {
         val schema = ISS.newSchema("\$ion_schema_2_0 type::{ name: foo }")