Restore annotations when data trace is restored for undo of command

app/databrowser/src/main/java/org/csstudio/trends/databrowser3/ui/properties/DeleteItemsCommand.java

@@ -9,7 +9,9 @@
 
 import java.text.MessageFormat;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.csstudio.trends.databrowser3.Messages;
 import org.csstudio.trends.databrowser3.model.AnnotationInfo;
@@ -26,6 +28,7 @@ public class DeleteItemsCommand extends UndoableAction
 {
     final private Model model;
     final private List<ModelItem> items;
+    final private Map<ModelItem, List<AnnotationInfo>> mapModelItemAnnotations;
 
     /** Register and perform the command
      *  @param operations_manager OperationsManager where command will be registered
@@ -40,6 +43,7 @@ public DeleteItemsCommand(final UndoableActionManager operations_manager,
         // This list could be a reference to the model's list.
         // Since we will loop over this list, assert that there are no co-modification problems by creating a copy.
         this.items = new ArrayList<>(items);
+        this.mapModelItemAnnotations = new HashMap<>();
         operations_manager.execute(this);
     }
 
@@ -60,12 +64,20 @@ public void run()
     }
 
     private List<AnnotationInfo> deleteAnnotations(ModelItem item) {
-        final List<AnnotationInfo> annotations = new ArrayList<>(model.getAnnotations());
+        // check for annotations to remove because their item is deleted
         final int item_index = model.getItems().indexOf(item);
+        final List<AnnotationInfo> deletedAnnotations = new ArrayList<>();
+        final List<AnnotationInfo> annotations = new ArrayList<>();
+        for (AnnotationInfo annotation : model.getAnnotations())
+        {
+            if (annotation.getItemIndex() == item_index)
+                deletedAnnotations.add(annotation);
+            else
+                annotations.add(annotation);
+        }
 
-        // Check for annotations to remove because their item is deleted
-        //     possibly recalculate item index for remaining annotations
-        if (annotations.removeIf(anno -> anno.getItemIndex() == item_index)
+        // possibly recalculate item index for remaining annotations
+        if (!deletedAnnotations.isEmpty()
                 && !annotations.isEmpty()
                 && item_index < (model.getItems().size() - 1))
         {
@@ -81,6 +93,9 @@ private List<AnnotationInfo> deleteAnnotations(ModelItem item) {
             }
         }
 
+        // keep track of deleted annotations
+        mapModelItemAnnotations.put(item, deletedAnnotations);
+
         // annotations after item has been deleted
         return annotations;
     }
@@ -94,6 +109,12 @@ public void undo()
             try
             {
                 model.addItem(item);
+
+                final List<AnnotationInfo> annotations = restoreAnnotations(item);
+
+                // Any changes?
+                if (!annotations.equals(model.getAnnotations()))
+                    model.setAnnotations(annotations);
             }
             catch (Exception ex)
             {
@@ -103,6 +124,19 @@ public void undo()
                         ex);
             }
         }
-        // Note that we do not restore removed annotations at this time...
     }
+
+    private List<AnnotationInfo> restoreAnnotations(ModelItem item)
+    {
+        final List<AnnotationInfo> annotations = new ArrayList<>(model.getAnnotations());
+        final int item_index = model.getItems().indexOf(item);
+
+        // check for annotations to restore because their item was restored
+        for (AnnotationInfo ai : mapModelItemAnnotations.get(item))
+            annotations.add(new AnnotationInfo(ai.isInternal(), item_index, ai.getTime(), ai.getValue(), ai.getOffset(), ai.getText()));
+
+        // annotations after item has been restored
+        return annotations;
+    }
+
 }