From 0bbeb4c6eeb89675d7651090d3102c294af592e5 Mon Sep 17 00:00:00 2001
From: Shane Friedman <smoores-gpg@friedmans.us>
Date: Mon, 31 Jul 2023 23:32:30 -0400
Subject: [PATCH] Start plumbing through decorations

---
 demo/main.tsx                                 |  1 +
 src/components/DocNodeView.tsx                | 13 +++++++--
 src/components/EditorView.tsx                 |  4 ++-
 src/components/NodeView.tsx                   | 28 +++++++++++++++----
 .../DecorationInternal.ts                     |  9 +++++-
 5 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/demo/main.tsx b/demo/main.tsx
index fdee5415..ec7abaa5 100644
--- a/demo/main.tsx
+++ b/demo/main.tsx
@@ -240,6 +240,7 @@ root.render(<DemoEditor />);
 //           DecorationSet.create(state.doc, [
 //             Decoration.inline(5, 15, { class: "inline-deco" }),
 //             Decoration.node(35, 55, { class: "node-deco" }),
+//             Decoration.widget(40, () => document.createElement("div")),
 //           ])
 //         }
 //       >
diff --git a/src/components/DocNodeView.tsx b/src/components/DocNodeView.tsx
index 9a7b1ce3..b4baef07 100644
--- a/src/components/DocNodeView.tsx
+++ b/src/components/DocNodeView.tsx
@@ -14,16 +14,18 @@ import React, {
 import { ChildDescriptorsContext } from "../contexts/ChildDescriptorsContext.js";
 import { NodeViewContext } from "../contexts/NodeViewContext.js";
 import { NodeViewDesc, ViewDesc } from "../descriptors/ViewDesc.js";
+import { DecorationSourceInternal } from "../prosemirror-internal/DecorationInternal.js";
 
 import { NodeView } from "./NodeView.js";
 
 type Props = {
   node: Node;
   contentEditable: boolean;
+  decorations: DecorationSourceInternal;
 } & DetailedHTMLProps<HTMLAttributes<HTMLDivElement>, HTMLDivElement>;
 
 export const DocNodeView = forwardRef(function DocNodeView(
-  { node, contentEditable, ...props }: Props,
+  { node, contentEditable, decorations, ...props }: Props,
   ref: ForwardedRef<HTMLDivElement>
 ) {
   const { posToDesc, domToDesc } = useContext(NodeViewContext);
@@ -61,7 +63,14 @@ export const DocNodeView = forwardRef(function DocNodeView(
   const innerPos = 0;
   node.content.forEach((childNode, offset) => {
     const childPos = innerPos + offset;
-    children.push(<NodeView key={childPos} node={childNode} pos={childPos} />);
+    children.push(
+      <NodeView
+        key={childPos}
+        node={childNode}
+        pos={childPos}
+        decorations={decorations.forChild(offset, childNode)}
+      />
+    );
   });
 
   return (
diff --git a/src/components/EditorView.tsx b/src/components/EditorView.tsx
index 22d48198..829cfec1 100644
--- a/src/components/EditorView.tsx
+++ b/src/components/EditorView.tsx
@@ -16,6 +16,7 @@ import { NodeViewContext } from "../contexts/NodeViewContext.js";
 import { NodeViewDesc, ViewDesc } from "../descriptors/ViewDesc.js";
 import { useContentEditable } from "../hooks/useContentEditable.js";
 import { useSyncSelection } from "../hooks/useSyncSelection.js";
+import { DecorationSourceInternal } from "../prosemirror-internal/DecorationInternal.js";
 import { EditorViewInternal } from "../prosemirror-internal/EditorViewInternal.js";
 import { DOMNode, DOMSelection } from "../prosemirror-internal/dom.js";
 import {
@@ -61,7 +62,7 @@ export function EditorView(props: Props) {
     // keymap = {},
     nodeViews = {},
     dispatchTransaction: dispatchProp,
-    // decorations = DecorationSet.empty,
+    decorations = DecorationSet.empty,
     defaultState,
     state: stateProp,
     // ...mountProps
@@ -222,6 +223,7 @@ export function EditorView(props: Props) {
               ref={mountRef}
               node={state.doc}
               contentEditable={editable}
+              decorations={decorations as unknown as DecorationSourceInternal}
               // {...mountProps}
             />
             {children}
diff --git a/src/components/NodeView.tsx b/src/components/NodeView.tsx
index 356d4153..6f6e8e37 100644
--- a/src/components/NodeView.tsx
+++ b/src/components/NodeView.tsx
@@ -13,6 +13,7 @@ import React, {
 import { ChildDescriptorsContext } from "../contexts/ChildDescriptorsContext.js";
 import { NodeViewContext } from "../contexts/NodeViewContext.js";
 import { NodeViewDesc, ViewDesc } from "../descriptors/ViewDesc.js";
+import { DecorationSourceInternal } from "../prosemirror-internal/DecorationInternal.js";
 
 import { MarkView } from "./MarkView.js";
 import { NodeViewComponentProps } from "./NodeViewComponentProps.js";
@@ -23,14 +24,16 @@ import { TrailingHackView } from "./TrailingHackView.js";
 type Props = {
   node: Node;
   pos: number;
+  decorations: DecorationSourceInternal;
 };
 
-export function NodeView({ node, pos }: Props) {
+export function NodeView({ node, pos, decorations }: Props) {
   const { posToDesc, domToDesc, nodeViews, state } =
     useContext(NodeViewContext);
   const siblingDescriptors = useContext(ChildDescriptorsContext);
   const childDescriptors: ViewDesc[] = [];
   const domRef = useRef<HTMLElement | null>(null);
+  const nodeDomRef = useRef<HTMLElement | null>(null);
 
   useLayoutEffect(() => {
     if (!domRef.current) return;
@@ -44,7 +47,7 @@ export function NodeView({ node, pos }: Props) {
       DecorationSet.empty,
       domRef.current,
       firstChildDesc?.dom.parentElement ?? null,
-      domRef.current,
+      nodeDomRef.current ?? domRef.current,
       posToDesc,
       domToDesc
     );
@@ -75,7 +78,14 @@ export function NodeView({ node, pos }: Props) {
         </ChildDescriptorsContext.Consumer>
       );
     } else {
-      content.push(<NodeView key={childPos} node={childNode} pos={childPos} />);
+      content.push(
+        <NodeView
+          key={childPos}
+          node={childNode}
+          pos={childPos}
+          decorations={decorations.forChild(offset, childNode)}
+        />
+      );
     }
   });
 
@@ -97,7 +107,11 @@ export function NodeView({ node, pos }: Props) {
 
   if (Component) {
     return node.marks.reduce(
-      (element, mark) => <MarkView mark={mark}>{element}</MarkView>,
+      (element, mark) => (
+        <MarkView mark={mark} ref={nodeDomRef}>
+          {element}
+        </MarkView>
+      ),
       <Component
         ref={domRef}
         node={node}
@@ -117,7 +131,11 @@ export function NodeView({ node, pos }: Props) {
 
   if (outputSpec) {
     return node.marks.reduce(
-      (element, mark) => <MarkView mark={mark}>{element}</MarkView>,
+      (element, mark) => (
+        <MarkView ref={nodeDomRef} mark={mark}>
+          {element}
+        </MarkView>
+      ),
       <OutputSpec ref={domRef} outputSpec={outputSpec}>
         {children}
       </OutputSpec>
diff --git a/src/prosemirror-internal/DecorationInternal.ts b/src/prosemirror-internal/DecorationInternal.ts
index 10503ea8..d71cb1e8 100644
--- a/src/prosemirror-internal/DecorationInternal.ts
+++ b/src/prosemirror-internal/DecorationInternal.ts
@@ -1,4 +1,5 @@
-import { Decoration } from "prosemirror-view";
+import { Node } from "prosemirror-model";
+import { Decoration, DecorationSource } from "prosemirror-view";
 import { DecorationType } from "../decorations/DecorationType";
 import { ReactWidgetType } from "../decorations/ReactWidgetType";
 
@@ -11,3 +12,9 @@ export interface ReactWidgetDecoration extends Decoration {
   type: ReactWidgetType
   inline: false
 }
+
+export interface DecorationSourceInternal extends DecorationSource {
+  locals(node: Node): readonly Decoration[]
+  forChild(offset: number, child: Node): DecorationSourceInternal
+  eq(other: DecorationSource): boolean
+}