From 50a12121f2fd14679e79841c4fdfc8e8a68569b0 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 a360b8be9..1013628de 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
@@ -1411,6 +1411,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 $}.
 			 *
@@ -1420,31 +1428,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 15704af03..28d522fba 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
@@ -1309,6 +1309,36 @@ class UpdateMapperUnitTests {
 		assertThat(mappedUpdate).isEqualTo(new org.bson.Document("$set",new org.bson.Document("customers", Arrays.asList("c-name"))));
 	}
 
+	@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;
 	}
@@ -1710,4 +1740,20 @@ class UpdateMapperUnitTests {
 		private List<Sample> samples;
 	}
 
+	@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;
+	}
 }