DRY

Author: trangdoan982

apps/roam/src/components/canvas/Tldraw.tsx

@@ -56,6 +56,7 @@ import {
 import "tldraw/tldraw.css";
 import tldrawStyles from "./tldrawStyles";
 import { DragHandleOverlay } from "./overlays/DragHandleOverlay";
+import { isDiscourseNodeShape } from "./canvasUtils";
 import getDiscourseNodes, { DiscourseNode } from "~/utils/getDiscourseNodes";
 import getDiscourseRelations, {
   DiscourseRelation,
@@ -507,12 +508,6 @@ const TldrawCanvasShared = ({
     );
   };
 
-  const isDiscourseNodeShape = (
-    shape: TLShape,
-  ): shape is DiscourseNodeShape => {
-    return allNodes.some((node) => node.type === shape.type);
-  };
-
   // Add state for tracking relation creation
   const relationCreationRef = useRef<{
     isCreating: boolean;
@@ -537,7 +532,7 @@ const TldrawCanvasShared = ({
         relationCreationRef.current.toolType = currentToolId;
 
         // If we clicked on a discourse node, record it as the source
-        if (shapeAtPoint && isDiscourseNodeShape(shapeAtPoint)) {
+        if (shapeAtPoint && isDiscourseNodeShape(app, shapeAtPoint)) {
           relationCreationRef.current.sourceShapeId = shapeAtPoint.id;
         }
       }
@@ -562,7 +557,7 @@ const TldrawCanvasShared = ({
           relationCreationRef.current.relationShapeId = relationShape.id;
 
           // Check if we have a target shape
-          if (shapeAtPoint && isDiscourseNodeShape(shapeAtPoint)) {
+          if (shapeAtPoint && isDiscourseNodeShape(app, shapeAtPoint)) {
             posthog.capture("Canvas: Relation Created", {
               relationType: relationShape.type,
               toolType: relationCreationRef.current.toolType || "",

apps/roam/src/components/canvas/canvasUtils.ts

@@ -0,0 +1,30 @@
+import { Editor, TLShape } from "tldraw";
+import {
+  BaseDiscourseNodeUtil,
+  DiscourseNodeShape,
+} from "~/components/canvas/DiscourseNodeUtil";
+import { discourseContext } from "~/components/canvas/Tldraw";
+
+export const isDiscourseNodeShape = (
+  editor: Editor,
+  shape: TLShape,
+): shape is DiscourseNodeShape => {
+  try {
+    return editor.getShapeUtil(shape) instanceof BaseDiscourseNodeUtil;
+  } catch {
+    return false;
+  }
+};
+
+export const getAllRelations = () =>
+  Object.values(discourseContext.relations).flat();
+
+export const hasValidRelationTypes = (
+  sourceNodeType: string,
+  targetNodeType: string,
+): boolean =>
+  getAllRelations().some(
+    (r) =>
+      (r.source === sourceNodeType && r.destination === targetNodeType) ||
+      (r.source === targetNodeType && r.destination === sourceNodeType),
+  );

apps/roam/src/components/canvas/overlays/DragHandleOverlay.tsx

@@ -1,23 +1,17 @@
 import React, { useCallback, useEffect, useRef, useState } from "react";
-import {
-  TLShapeId,
-  createShapeId,
-  useEditor,
-  useValue,
-} from "tldraw";
-import {
-  BaseDiscourseNodeUtil,
-  DiscourseNodeShape,
-} from "~/components/canvas/DiscourseNodeUtil";
+import { TLShapeId, createShapeId, useEditor, useValue } from "tldraw";
+import { DiscourseNodeShape } from "~/components/canvas/DiscourseNodeUtil";
 import {
   BaseDiscourseRelationUtil,
   DiscourseRelationShape,
   getRelationColor,
 } from "~/components/canvas/DiscourseRelationShape/DiscourseRelationUtil";
+import { createOrUpdateArrowBinding } from "~/components/canvas/DiscourseRelationShape/helpers";
 import {
-  createOrUpdateArrowBinding,
-} from "~/components/canvas/DiscourseRelationShape/helpers";
-import { discourseContext } from "~/components/canvas/Tldraw";
+  getAllRelations,
+  hasValidRelationTypes,
+  isDiscourseNodeShape,
+} from "~/components/canvas/canvasUtils";
 import { dispatchToastEvent } from "~/components/canvas/ToastListener";
 import { RelationTypeDropdown } from "./RelationTypeDropdown";
 
@@ -67,30 +61,6 @@ const getEdgeMidpoints = (bounds: {
   },
 ];
 
-const isDiscourseNode = (
-  editor: ReturnType<typeof useEditor>,
-  shapeType: string,
-): boolean => {
-  try {
-    const util = editor.getShapeUtil(shapeType);
-    return util instanceof BaseDiscourseNodeUtil;
-  } catch {
-    return false;
-  }
-};
-
-const hasValidRelationTypes = (
-  sourceNodeType: string,
-  targetNodeType: string,
-): boolean => {
-  const allRelations = Object.values(discourseContext.relations).flat();
-  return allRelations.some(
-    (r) =>
-      (r.source === sourceNodeType && r.destination === targetNodeType) ||
-      (r.source === targetNodeType && r.destination === sourceNodeType),
-  );
-};
-
 export const DragHandleOverlay = () => {
   const editor = useEditor();
 
@@ -120,7 +90,7 @@ export const DragHandleOverlay = () => {
     () => {
       if (isDragging || pending) return sourceNodeRef.current;
       const shape = editor.getOnlySelectedShape();
-      if (shape && isDiscourseNode(editor, shape.type)) {
+      if (shape && isDiscourseNodeShape(editor, shape)) {
         return shape as DiscourseNodeShape;
       }
       return null;
@@ -187,7 +157,7 @@ export const DragHandleOverlay = () => {
           hitFrameInside: true,
           margin: 0,
           filter: (s) =>
-            isDiscourseNode(editor, s.type) &&
+            isDiscourseNodeShape(editor, s) &&
             s.id !== selectedNode.id &&
             !s.isLocked,
         });
@@ -215,7 +185,7 @@ export const DragHandleOverlay = () => {
           hitFrameInside: true,
           margin: 0,
           filter: (s) =>
-            isDiscourseNode(editor, s.type) &&
+            isDiscourseNodeShape(editor, s) &&
             s.id !== selectedNode.id &&
             !s.isLocked,
         });
@@ -276,8 +246,9 @@ export const DragHandleOverlay = () => {
     (relationId: string) => {
       if (!pending) return;
 
-      const allRelations = Object.values(discourseContext.relations).flat();
-      const selectedRelation = allRelations.find((r) => r.id === relationId);
+      const selectedRelation = getAllRelations().find(
+        (r) => r.id === relationId,
+      );
       if (!selectedRelation) {
         setPending(null);
         sourceNodeRef.current = null;

apps/roam/src/components/canvas/overlays/RelationTypeDropdown.tsx

@@ -1,12 +1,10 @@
 import React, { useCallback, useEffect, useMemo, useRef } from "react";
-import {
-  TLShapeId,
-  useEditor,
-  DefaultColorThemePalette,
-} from "tldraw";
-import { discourseContext } from "~/components/canvas/Tldraw";
-import { BaseDiscourseNodeUtil } from "~/components/canvas/DiscourseNodeUtil";
+import { TLShapeId, useEditor, DefaultColorThemePalette } from "tldraw";
 import { getRelationColor } from "~/components/canvas/DiscourseRelationShape/DiscourseRelationUtil";
+import {
+  getAllRelations,
+  isDiscourseNodeShape,
+} from "~/components/canvas/canvasUtils";
 
 type RelationTypeDropdownProps = {
   sourceId: TLShapeId;
@@ -36,21 +34,15 @@ export const RelationTypeDropdown = ({
     const endNodeType = endNode.type;
 
     // Verify both are discourse nodes
-    try {
-      const startUtil = editor.getShapeUtil(startNode);
-      const endUtil = editor.getShapeUtil(endNode);
-      if (
-        !(startUtil instanceof BaseDiscourseNodeUtil) ||
-        !(endUtil instanceof BaseDiscourseNodeUtil)
-      )
-        return [];
-    } catch {
+    if (
+      !isDiscourseNodeShape(editor, startNode) ||
+      !isDiscourseNodeShape(editor, endNode)
+    )
       return [];
-    }
 
     const colorPalette = DefaultColorThemePalette.lightMode;
     const validTypes: { id: string; label: string; color: string }[] = [];
-    const allRelations = Object.values(discourseContext.relations).flat();
+    const allRelations = getAllRelations();
     const seenLabels = new Set<string>();
 
     for (const relation of allRelations) {