From 378c2123f091a5827723b727f4423d6a31a33fc5 Mon Sep 17 00:00:00 2001
From: Dennis Huebner <dennis.huebner@gmail.com>
Date: Tue, 23 Jan 2024 11:21:46 +0100
Subject: [PATCH] Track augmented copies to apply deviation

---
 .../yang/processor/ProcessorUtility.java      | 22 ++++++++++++++++++-
 .../typefox/yang/processor/YangProcessor.java | 20 ++++++++++++-----
 .../yang/tests/processor/DeviationTest.xtend  |  3 ++-
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/ProcessorUtility.java b/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/ProcessorUtility.java
index f2f256bc..e4188e9d 100644
--- a/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/ProcessorUtility.java
+++ b/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/ProcessorUtility.java
@@ -4,7 +4,9 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
+import org.eclipse.emf.common.notify.impl.AdapterImpl;
 import org.eclipse.emf.ecore.EObject;
 import org.eclipse.emf.ecore.util.EcoreUtil.Copier;
 import org.eclipse.xtext.EcoreUtil2;
@@ -103,6 +105,7 @@ protected EObject createCopy(EObject eObject) {
 					// store text information. e.g. to serialize XPath
 					createCopy.eAdapters().add((CompositeNodeWithSemanticElement) node);
 				}
+				eObject.eAdapters().add(new CopiedObjectAdapter(createCopy));
 				return createCopy;
 			}
 		};
@@ -112,7 +115,7 @@ protected EObject createCopy(EObject eObject) {
 	}
 
 	public static String serializedXpath(XpathExpression reference) {
-		if(reference == null) {
+		if (reference == null) {
 			return null;
 		}
 		// TODO use serializer or implement a an own simple one
@@ -130,4 +133,21 @@ public static String serializedXpath(XpathExpression reference) {
 		return "leafref";
 	}
 
+	public static class CopiedObjectAdapter extends AdapterImpl {
+		final EObject copy;
+
+		public CopiedObjectAdapter(EObject copy) {
+			this.copy = copy;
+		}
+
+		public EObject getCopy() {
+			return copy;
+		}
+
+		public static Stream<CopiedObjectAdapter> findAll(EObject eObject) {
+			return eObject.eAdapters().stream().filter(a -> a instanceof CopiedObjectAdapter)
+					.map(a -> ((CopiedObjectAdapter) a));
+		}
+
+	}
 }
diff --git a/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/YangProcessor.java b/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/YangProcessor.java
index 23fb698f..de6c12f6 100644
--- a/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/YangProcessor.java
+++ b/yang-lsp/io.typefox.yang/src/main/java/io/typefox/yang/processor/YangProcessor.java
@@ -23,6 +23,7 @@
 import io.typefox.yang.processor.ProcessedDataModel.HasStatements;
 import io.typefox.yang.processor.ProcessedDataModel.ListData;
 import io.typefox.yang.processor.ProcessedDataModel.ModuleData;
+import io.typefox.yang.processor.ProcessorUtility.CopiedObjectAdapter;
 import io.typefox.yang.utils.YangNameUtils;
 import io.typefox.yang.yang.AbstractModule;
 import io.typefox.yang.yang.Action;
@@ -103,9 +104,9 @@ protected ProcessedDataModel processInternal(List<AbstractModule> modules, List<
 				augments.add((Augment) ele);
 			}
 		}));
-		
-		deviations.forEach(dev -> processDeviate(dev, processedModel));
+
 		augments.forEach(augm -> processAugment(augm, evalCtx));
+		deviations.forEach(dev -> processDeviate(dev, processedModel));
 
 		modules.forEach((module) -> {
 			String prefix = null;
@@ -207,12 +208,19 @@ protected void processDeviate(Deviate deviate, ProcessedDataModel processedModel
 				}
 				break;
 			}
-			Object eGet = targetNode.eContainer().eGet(targetNode.eContainingFeature(), true);
+			removeFromContainer(targetNode);
+			DeviationAdapter.add(targetNode, argument);
+			break;
+		}
+	}
+
+	private void removeFromContainer(EObject objToRemove) {
+		if (objToRemove.eContainer() != null) {
+			Object eGet = objToRemove.eContainer().eGet(objToRemove.eContainingFeature(), true);
 			if (eGet instanceof EList) {
-				DeviationAdapter.add(targetNode, argument);
-				((EList<?>) eGet).remove(targetNode);
+				((EList<?>) eGet).remove(objToRemove);
 			}
-			break;
+			CopiedObjectAdapter.findAll(objToRemove).forEach(adapter -> removeFromContainer(adapter.getCopy()));
 		}
 	}
 
diff --git a/yang-lsp/io.typefox.yang/src/test/java/io/typefox/yang/tests/processor/DeviationTest.xtend b/yang-lsp/io.typefox.yang/src/test/java/io/typefox/yang/tests/processor/DeviationTest.xtend
index 6d78c1ad..938c3209 100644
--- a/yang-lsp/io.typefox.yang/src/test/java/io/typefox/yang/tests/processor/DeviationTest.xtend
+++ b/yang-lsp/io.typefox.yang/src/test/java/io/typefox/yang/tests/processor/DeviationTest.xtend
@@ -310,7 +310,7 @@ class DeviationTest extends AbstractYangTest {
 	}
 
 	@Test
-	def void testDeviateAugmentedNode() {
+	def void testDeviateAugmentedNode_01() {
 
 		val mainModule = '''
 			module base-test-module {
@@ -345,4 +345,5 @@ class DeviationTest extends AbstractYangTest {
 		     +--rw simple-leaf?   string
 		'''.toString, new DataTreeSerializer().serialize(processedData.modules.get(0)).toString)
 	}
+
 }