From 8a62339ab9625d15d4ecd65836e641c5997b94b9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Ku=CC=88hn?= <kontakt@philipp-kuehn.com>
Date: Thu, 11 Feb 2021 23:14:17 +0100
Subject: [PATCH] refactoring

---
 .../Annotation/extension/AnnotationPlugin.ts  |  7 ++-
 .../Annotation/extension/AnnotationState.ts   | 47 ++++++++++---------
 2 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts b/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
index 29a382c4..406889f6 100644
--- a/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
+++ b/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
@@ -14,8 +14,11 @@ export interface AnnotationPluginOptions {
 export const AnnotationPlugin = (options: AnnotationPluginOptions) => new Plugin({
   key: AnnotationPluginKey,
   state: {
-    init(_, state) {
-      return AnnotationState.init(_, state)
+    init() {
+      return new AnnotationState({
+        HTMLAttributes: options.HTMLAttributes,
+        map: options.map,
+      })
     },
     apply(transaction, pluginState, oldState, newState) {
       return pluginState.apply(transaction, newState)
diff --git a/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts b/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
index a515826e..587d2260 100644
--- a/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
+++ b/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
@@ -1,15 +1,23 @@
 // @ts-nocheck
-import * as Y from 'yjs'
 import { EditorState } from 'prosemirror-state'
 import { Decoration, DecorationSet } from 'prosemirror-view'
 import { ySyncPluginKey, relativePositionToAbsolutePosition, absolutePositionToRelativePosition } from 'y-prosemirror'
 import { AnnotationPluginKey } from './AnnotationPlugin'
 
-export class AnnotationState {
-  private decorations: any
+export interface AnnotationStateOptions {
+  HTMLAttributes: {
+    [key: string]: any
+  },
+  map: any,
+}
 
-  constructor(decorations: any) {
-    this.decorations = decorations
+export class AnnotationState {
+  options: AnnotationStateOptions
+
+  decorations = DecorationSet.empty
+
+  constructor(options: AnnotationStateOptions) {
+    this.options = options
   }
 
   findAnnotation(id: number) {
@@ -23,18 +31,18 @@ export class AnnotationState {
   }
 
   annotationsAt(position: number) {
-    return this.decorations.find(position, position)
+    return this.decorations?.find(position, position)
   }
 
   apply(transaction: any, state: EditorState) {
     const ystate = ySyncPluginKey.getState(state)
-    const decs = ystate.doc.getMap('annotations')
+    const decs = this.options.map
 
     const action = transaction.getMeta(AnnotationPluginKey)
     const actionType = action && action.type
 
     if (action) {
-      let { decorations } = this
+      const { decorations } = this
 
       if (actionType === 'addAnnotation') {
         decs.set(action.data.id, {
@@ -51,18 +59,18 @@ export class AnnotationState {
           data: action.data,
         })
 
-        decorations = decorations.add(transaction.doc, [
-          Decoration.inline(action.from, action.to, { class: 'annotation' }, { data: action.data }),
+        this.decorations = decorations.add(transaction.doc, [
+          Decoration.inline(action.from, action.to, this.options.HTMLAttributes, { data: action.data }),
         ])
       } else if (actionType === 'deleteAnnotation') {
         decs.delete(action.id)
 
-        decorations = decorations.remove([
+        this.decorations = decorations.remove([
           this.findAnnotation(action.id),
         ])
       }
 
-      return new AnnotationState(decorations)
+      return this
     }
 
     if (ystate && ystate.isChangeOrigin) {
@@ -84,22 +92,19 @@ export class AnnotationState {
             dec.to,
             ystate.binding.mapping,
           ),
-          { class: 'annotation' },
+          this.options.HTMLAttributes,
           { data: dec.data },
         ))
       })
 
-      return new AnnotationState(DecorationSet.create(state.doc, decorations))
+      this.decorations = DecorationSet.create(state.doc, decorations)
+
+      return this
     }
 
     // Apply ProseMirror mapping
-    const decorations = this.decorations.map(transaction.mapping, transaction.doc)
-    return new AnnotationState(decorations)
-  }
+    this.decorations = this.decorations.map(transaction.mapping, transaction.doc)
 
-  static init(config: any, state: EditorState) {
-    const decorations = DecorationSet.create(state.doc, [])
-
-    return new AnnotationState(decorations)
+    return this
   }
 }