From b75db97c26b82a343e17e7614fd84ef40db94683 Mon Sep 17 00:00:00 2001
From: Mitchell McCaffrey <mitchemmc@gmail.com>
Date: Thu, 10 Jun 2021 18:04:32 +1000
Subject: [PATCH] Add custom drag context wrapper to inject the overlay rect on
 drag end

---
 src/components/token/TokenBar.js | 34 ++++-----------
 src/contexts/DragContext.js      | 75 ++++++++++++++++++++++++++++++++
 src/contexts/TileDragContext.js  | 13 +++---
 src/modals/SelectTokensModal.js  |  4 +-
 4 files changed, 94 insertions(+), 32 deletions(-)
 create mode 100644 src/contexts/DragContext.js

diff --git a/src/components/token/TokenBar.js b/src/components/token/TokenBar.js
index 7f42544..ea4ef85 100644
--- a/src/components/token/TokenBar.js
+++ b/src/components/token/TokenBar.js
@@ -1,10 +1,9 @@
-import React, { useState, useRef } from "react";
+import React, { useState } from "react";
 import { createPortal } from "react-dom";
 import { Box, Flex } from "theme-ui";
 import SimpleBar from "simplebar-react";
 import {
   DragOverlay,
-  DndContext,
   MouseSensor,
   TouchSensor,
   KeyboardSensor,
@@ -24,6 +23,7 @@ import usePreventSelect from "../../hooks/usePreventSelect";
 import { useTokenData } from "../../contexts/TokenDataContext";
 import { useAuth } from "../../contexts/AuthContext";
 import { useMapStage } from "../../contexts/MapStageContext";
+import DragContext from "../../contexts/DragContext";
 
 import {
   createTokenState,
@@ -40,10 +40,6 @@ function TokenBar({ onMapTokensStateCreate }) {
   const [dragId, setDragId] = useState();
 
   const mapStageRef = useMapStage();
-  // Use a ref to the drag overlay to get it's position on dragEnd
-  // TODO: use active.rect when dnd-kit bug is fixed
-  // https://github.com/clauderic/dnd-kit/issues/238
-  const dragOverlayRef = useRef();
 
   const mouseSensor = useSensor(MouseSensor, {
     activationConstraint: { distance: 5 },
@@ -61,13 +57,12 @@ function TokenBar({ onMapTokensStateCreate }) {
     preventSelect();
   }
 
-  function handleDragEnd({ active }) {
+  function handleDragEnd({ active, overlayNodeClientRect }) {
     setDragId(null);
 
     const mapStage = mapStageRef.current;
-    const dragOverlay = dragOverlayRef.current;
-    if (mapStage && dragOverlay) {
-      const dragRect = dragOverlay.getBoundingClientRect();
+    if (mapStage) {
+      const dragRect = overlayNodeClientRect;
       const dragPosition = {
         x: dragRect.left + dragRect.width / 2,
         y: dragRect.top + dragRect.height / 2,
@@ -146,7 +141,7 @@ function TokenBar({ onMapTokensStateCreate }) {
   }
 
   return (
-    <DndContext
+    <DragContext
       onDragStart={handleDragStart}
       onDragEnd={handleDragEnd}
       onDragCancel={handleDragCancel}
@@ -185,24 +180,13 @@ function TokenBar({ onMapTokensStateCreate }) {
           <SelectTokensButton onMapTokensStateCreate={onMapTokensStateCreate} />
         </Flex>
         {createPortal(
-          <DragOverlay
-            // Ensure a drop animation plays to allow us to get the position of the drag overlay in drag end
-            dropAnimation={{
-              dragSourceOpacity: 0,
-              duration: 1,
-              easing: "ease",
-            }}
-          >
-            {dragId && (
-              <div ref={dragOverlayRef}>
-                {renderToken(findGroup(tokenGroups, dragId), false)}
-              </div>
-            )}
+          <DragOverlay dropAnimation={null}>
+            {dragId && renderToken(findGroup(tokenGroups, dragId), false)}
           </DragOverlay>,
           document.body
         )}
       </Box>
-    </DndContext>
+    </DragContext>
   );
 }
 
diff --git a/src/contexts/DragContext.js b/src/contexts/DragContext.js
new file mode 100644
index 0000000..2614b0f
--- /dev/null
+++ b/src/contexts/DragContext.js
@@ -0,0 +1,75 @@
+// eslint-disable-next-line no-unused-vars
+import React, { useRef, ReactNode } from "react";
+import {
+  DndContext,
+  useDndContext,
+  useDndMonitor,
+  // eslint-disable-next-line no-unused-vars
+  DragEndEvent,
+} from "@dnd-kit/core";
+
+/**
+ * Wrap a dnd-kit DndContext with a position monitor to get the
+ * active drag element on drag end
+ * TODO: use look into fixing this upstream
+ * Related: https://github.com/clauderic/dnd-kit/issues/238
+ */
+
+/**
+ * @typedef DragEndOverlayEvent
+ * @property {DOMRect} overlayNodeClientRect
+ *
+ * @typedef {DragEndEvent & DragEndOverlayEvent} DragEndWithOverlayProps
+ */
+
+/**
+ * @callback DragEndWithOverlayEvent
+ * @param {DragEndWithOverlayProps} props
+ */
+
+/**
+ * @typedef CustomDragProps
+ * @property {DragEndWithOverlayEvent=} onDragEnd
+ * @property {ReactNode} children
+ */
+
+/**
+ * @param {CustomDragProps} props
+ */
+function DragPositionMonitor({ children, onDragEnd }) {
+  const { overlayNode } = useDndContext();
+
+  const overlayNodeClientRectRef = useRef();
+  function handleDragMove() {
+    if (overlayNode?.nodeRef?.current) {
+      overlayNodeClientRectRef.current = overlayNode.nodeRef.current.getBoundingClientRect();
+    }
+  }
+
+  function handleDragEnd(props) {
+    onDragEnd &&
+      onDragEnd({
+        ...props,
+        overlayNodeClientRect: overlayNodeClientRectRef.current,
+      });
+  }
+  useDndMonitor({ onDragEnd: handleDragEnd, onDragMove: handleDragMove });
+
+  return children;
+}
+
+/**
+ * TODO: Import Props interface from dnd-kit with conversion to Typescript
+ * @param {CustomDragProps} props
+ */
+function DragContext({ children, onDragEnd, ...props }) {
+  return (
+    <DndContext {...props}>
+      <DragPositionMonitor onDragEnd={onDragEnd}>
+        {children}
+      </DragPositionMonitor>
+    </DndContext>
+  );
+}
+
+export default DragContext;
diff --git a/src/contexts/TileDragContext.js b/src/contexts/TileDragContext.js
index 2e8c229..6f6644e 100644
--- a/src/contexts/TileDragContext.js
+++ b/src/contexts/TileDragContext.js
@@ -1,6 +1,5 @@
 import React, { useState, useContext } from "react";
 import {
-  DndContext,
   MouseSensor,
   TouchSensor,
   KeyboardSensor,
@@ -9,6 +8,8 @@ import {
   closestCenter,
 } from "@dnd-kit/core";
 
+import DragContext from "./DragContext";
+
 import { useGroup } from "./GroupContext";
 
 import { moveGroupsInto, moveGroups, ungroup } from "../helpers/group";
@@ -108,7 +109,7 @@ export function TileDragProvider({
   }
 
   function handleDragEnd(event) {
-    const { active, over } = event;
+    const { active, over, overlayNodeClientRect } = event;
 
     setDragId();
     setOverId();
@@ -143,7 +144,9 @@ export function TileDragProvider({
         }
         onGroupsChange(newGroups);
       } else if (over.id === ADD_TO_MAP_ID) {
-        onDragAdd && onDragAdd(selectedGroupIds, over.rect);
+        onDragAdd &&
+          overlayNodeClientRect &&
+          onDragAdd(selectedGroupIds, overlayNodeClientRect);
       } else if (!filter) {
         // Hanlde tile move only if we have no filter
         const overGroupIndex = activeGroups.findIndex(
@@ -210,7 +213,7 @@ export function TileDragProvider({
   const value = { dragId, overId, dragCursor };
 
   return (
-    <DndContext
+    <DragContext
       onDragStart={handleDragStart}
       onDragEnd={handleDragEnd}
       onDragOver={handleDragOver}
@@ -221,7 +224,7 @@ export function TileDragProvider({
       <TileDragContext.Provider value={value}>
         {children}
       </TileDragContext.Provider>
-    </DndContext>
+    </DragContext>
   );
 }
 
diff --git a/src/modals/SelectTokensModal.js b/src/modals/SelectTokensModal.js
index 42606b4..421e0e9 100644
--- a/src/modals/SelectTokensModal.js
+++ b/src/modals/SelectTokensModal.js
@@ -144,8 +144,8 @@ function SelectTokensModal({ isOpen, onRequestClose, onMapTokensStateCreate }) {
   const mapStageRef = useMapStage();
   function handleTokensAddToMap(groupIds, rect) {
     let clientPosition = new Vector2(
-      rect.width / 2 + rect.offsetLeft,
-      rect.height / 2 + rect.offsetTop
+      rect.width / 2 + rect.left,
+      rect.height / 2 + rect.top
     );
     const mapStage = mapStageRef.current;
     if (!mapStage) {