remove note views from context

wax-prosemirror-components/src/components/comments/ConnectedComment.js

@@ -30,11 +30,14 @@ export default ({ comment, top, commentId, recalculateTops }) => {
 
   const { state, dispatch } = activeView;
   const viewId = comment.attrs.viewid;
+  let allCommentsWithSameId = [];
 
-  const allCommentsWithSameId = DocumentHelpers.findAllMarksWithSameId(
-    view[viewId].state,
-    comment,
-  );
+  if (view[viewId]) {
+    allCommentsWithSameId = DocumentHelpers.findAllMarksWithSameId(
+      view[viewId].state,
+      comment,
+    );
+  }
 
   const commentMark = state.schema.marks.comment;
 

wax-prosemirror-components/src/components/rightArea/RightArea.js

@@ -66,11 +66,11 @@ export default ({ area }) => {
       if (markNodeEl) {
         boxEl = document.querySelector(`div[data-box="${id}"]`);
       }
-      if (boxEl) boxHeight = parseInt(boxEl.offsetHeight, 10);
-
-      // where the box should move to
-      top = boxEl ? annotationTop : -2000;
-
+      if (boxEl) {
+        boxHeight = parseInt(boxEl.offsetHeight, 10);
+        // where the box should move to
+        top = annotationTop;
+      }
       // if the above comment box has already taken up the height, move down
       if (pos > 0) {
         const previousBox = marksNodes[area][pos - 1];

wax-prosemirror-core/src/WaxContext.js

@@ -18,10 +18,8 @@ export default props => {
     view: props.view || {},
     activeView: props.activeView || {},
     activeViewId: props.activeViewId || {},
-    removeView: view => {
-      const newContext = { ...context, view };
-
-      setContext({ ...newContext });
+    removeView: deletedView => {
+      delete context.view[deletedView];
     },
     updateView: (newView, activeViewId) => {
       const view = Object.assign(context.view, newView);

wax-prosemirror-services/src/NoteService/NoteComponent.js

@@ -6,13 +6,13 @@ import { DocumentHelpers } from 'wax-prosemirror-utilities';
 import NoteEditor from './NoteEditor';
 
 export default () => {
+  const context = useContext(WaxContext);
   const {
     view,
     view: { main },
-  } = useContext(WaxContext);
+  } = context;
 
   const [notes, setNotes] = useState([]);
-
   const cleanUpNoteViews = () => {
     if (view) {
       const currentNotes = DocumentHelpers.findChildrenByType(
@@ -21,18 +21,24 @@ export default () => {
         true,
       );
       if (notes.length > currentNotes.length) {
-        // TODO remove from context views that no loger exist
         const difference = differenceBy(notes, currentNotes, 'node.attrs.id');
         difference.forEach((item, i) => {
-          // delete view[item.node.attrs.id];
+          context.removeView(item.node.attrs.id);
         });
+
+        // const newView = Object.keys(view).reduce((object, key) => {
+        //   if (key !== difference[0].node.attrs.id) {
+        //     object[key] = view[key];
+        //   }
+        //   return object;
+        // }, {});
       }
     }
   };
 
   useDeepCompareEffect(() => {
     setNotes(updateNotes(main));
-    // cleanUpNoteViews();
+    cleanUpNoteViews();
   }, [updateNotes(main)]);
 
   const noteComponent = useMemo(