diff --git a/src/editor-mode/decorator/widgets/ColorButton.ts b/src/editor-mode/decorator/widgets/ColorButton.ts
index 3c4ddb6..3ecc425 100644
--- a/src/editor-mode/decorator/widgets/ColorButton.ts
+++ b/src/editor-mode/decorator/widgets/ColorButton.ts
@@ -1,30 +1,24 @@
 import { WidgetType, EditorView, Decoration } from "@codemirror/view"
 import { Menu } from "obsidian";
-import { ColorMenu } from "src/editor-mode/ui-components";
-import { PlainRange, Token } from "src/types";
+import { TagMenu } from "src/editor-mode/ui-components";
+import { Format } from "src/enums";
+import { Token } from "src/types";
 
 /**
  * These code snippets are taken from 
  * https://github.com/Superschnizel/obisdian-fast-text-color/blob/master/src/widgets/ColorWidget.ts
  * with some modifications.
-*/
+ */
 export class ColorButton extends WidgetType {
     menu: Menu;
-    colorMenu: ColorMenu;
-    tagRange: PlainRange;
-    openRange: PlainRange;
-    closeRange: PlainRange;
+    btnPos: number;
+    colorMenu: TagMenu;
     constructor(token: Token) {
         super();
-        this.openRange = { from: token.from, to: token.from + token.openLen };
-        this.tagRange = { from: this.openRange.to, to: this.openRange.to + token.tagLen };
-        this.closeRange = { from: token.to - token.closeLen, to: token.to };
+        this.btnPos = token.from + token.openLen;
     }
     eq(other: ColorButton) {
-        return (
-            other.tagRange.from == this.tagRange.from &&
-            other.tagRange.to == this.tagRange.to
-        );
+        return this.btnPos == other.btnPos;
     }
     toDOM(view: EditorView): HTMLElement {
         let btn = document.createElement("span");
@@ -32,12 +26,9 @@ export class ColorButton extends WidgetType {
         btn.className = "cm-highlight-color-btn";
         btn.onclick = () => {
             view.dispatch({
-                selection: {
-                    anchor: this.tagRange.from,
-                    head: this.tagRange.to
-                }
+                selection: { anchor: this.btnPos }
             });
-            this.colorMenu = ColorMenu.create(view, this.openRange, this.tagRange, this.closeRange, true);
+            this.colorMenu ||= TagMenu.create(view, Format.HIGHLIGHT);
             this.colorMenu.showMenu();
         }
         return btn;