From f00e8ed93c88fda00c1bbccf3a0bd47b8c073c82 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@vmware.com>
Date: Tue, 11 Jan 2022 13:06:02 +0100
Subject: [PATCH] Use index instead of iterator to map position and map keys
 for updates.

This commit removes usage of the iterator and replaces map key and positional parameter mappings with an index based token lookup.

Closes #3921
Original pull request: #3930.
---
 .../mongodb/core/convert/QueryMapper.java     | 46 ++++++++++---------
 .../core/convert/UpdateMapperUnitTests.java   | 46 +++++++++++++++++++
 2 files changed, 70 insertions(+), 22 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java
index bc01ce6ff..49a7484d9 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java
@@ -1388,6 +1388,14 @@ public class QueryMapper {
 				this.currentIndex = 0;
 			}
 
+			String nextToken() {
+				return pathParts.get(currentIndex+1);
+			}
+
+			boolean hasNexToken() {
+				return pathParts.size() > currentIndex+1;
+			}
+
 			/**
 			 * Maps the property name while retaining potential positional operator {@literal $}.
 			 *
@@ -1397,31 +1405,25 @@ public class QueryMapper {
 			protected String mapPropertyName(MongoPersistentProperty property) {
 
 				StringBuilder mappedName = new StringBuilder(PropertyToFieldNameConverter.INSTANCE.convert(property));
-
-				boolean inspect = iterator.hasNext();
-
-				while (inspect) {
-
-					String partial = iterator.next();
-					currentIndex++;
-
-					boolean isPositional = isPositionalParameter(partial) && property.isCollectionLike() ;
-					if(property.isMap() && currentPropertyRoot.equals(partial) && iterator.hasNext()){
-						partial = iterator.next();
-						currentIndex++;
-					}
-
-					if (isPositional || property.isMap() && !currentPropertyRoot.equals(partial)) {
-						mappedName.append(".").append(partial);
-					}
-
-					inspect = isPositional && iterator.hasNext();
+				if(!hasNexToken()) {
+					return mappedName.toString();
 				}
 
-				if(currentIndex + 1 < pathParts.size()) {
-					currentIndex++;
-					currentPropertyRoot = pathParts.get(currentIndex);
+				String nextToken = nextToken();
+				if(isPositionalParameter(nextToken)) {
+
+					mappedName.append(".").append(nextToken);
+					currentIndex+=2;
+					return mappedName.toString();
 				}
+
+				if(property.isMap()) {
+
+					mappedName.append(".").append(nextToken);
+					currentIndex+=2;
+					return mappedName.toString();
+				}
+				currentIndex++;
 				return mappedName.toString();
 			}
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
index 111567bea..6e3e492f4 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
@@ -1251,6 +1251,36 @@ class UpdateMapperUnitTests {
 		assertThat(mappedUpdate).isEqualTo("{\"$set\": {\"intKeyedMap.1a.map.0b\": \"testing\"}}");
 	}
 
+	@Test // GH-3921
+	void mapNumericKeyInPathHavingComplexMapValyeTypes() {
+
+		Update update = new Update().set("testInnerData.testMap.1.intValue", "4");
+		Document mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(TestData.class));
+
+		assertThat(mappedUpdate).isEqualTo(new org.bson.Document("$set",new org.bson.Document("testInnerData.testMap.1.intValue","4")));
+	}
+
+	@Test // GH-3921
+	void mapNumericKeyInPathNotMatchingExistingProperties() {
+
+		Update update = new Update().set("testInnerData.imaginaryMap.1.nonExistingProperty", "4");
+		Document mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(TestData.class));
+
+		assertThat(mappedUpdate).isEqualTo(new org.bson.Document("$set",new org.bson.Document("testInnerData.imaginaryMap.1.noExistingProperty","4")));
+	}
+
+	@Test // GH-3921
+	void mapNumericKeyInPathPartiallyMatchingExistingProperties() {
+
+		Update update = new Update().set("testInnerData.testMap.1.nonExistingProperty.2.someValue", "4");
+		Document mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(TestData.class));
+
+		assertThat(mappedUpdate).isEqualTo(new org.bson.Document("$set",new org.bson.Document("testInnerData.testMap.1.nonExistingProperty.2.someValue","4")));
+	}
+
 	static class DomainTypeWrappingConcreteyTypeHavingListOfInterfaceTypeAttributes {
 		ListModelWrapper concreteTypeWithListAttributeOfInterfaceType;
 	}
@@ -1621,4 +1651,20 @@ class UpdateMapperUnitTests {
 		Map<String, Map<String, Map<String, Object>>> levelOne;
 	}
 
+	@Data
+	private static class TestData {
+		@Id
+		private String id;
+		private TestInnerData testInnerData;
+	}
+
+	@Data
+	private static class TestInnerData {
+		private Map<Integer, TestValue> testMap;
+	}
+
+	@Data
+	private static class TestValue {
+		private int intValue;
+	}
 }