[form-builder] PTE: improve popover editing reference element logic

packages/@sanity/form-builder/src/inputs/PortableText/Input.tsx

@@ -1,7 +1,7 @@
 /* eslint-disable react/require-default-props */
 import classNames from 'classnames'
 import {Subject} from 'rxjs'
-import React, {useEffect, useState, useMemo, useCallback} from 'react'
+import React, {useEffect, useState, useMemo, useCallback, useRef} from 'react'
 import {FormFieldPresence} from '@sanity/base/presence'
 import {
   getPortableTextFeatures,
@@ -15,7 +15,7 @@ import {
   usePortableTextEditorSelection,
   HotkeyOptions
 } from '@sanity/portable-text-editor'
-import {Path, isKeySegment} from '@sanity/types'
+import {Path, isKeySegment, KeyedSegment} from '@sanity/types'
 import {uniqueId, isEqual} from 'lodash'
 import ActivateOnFocus from 'part:@sanity/components/utilities/activate-on-focus'
 import {Portal} from 'part:@sanity/components/portal'
@@ -42,7 +42,7 @@ type Props = {
   onBlur: () => void
   onChange: (event: PatchEvent) => void
   onCopy?: OnCopyFn
-  onFocus: (path) => void
+  onFocus: (path: Path) => void
   onPaste?: OnPasteFn
   onToggleFullscreen: () => void
   patche$: Subject<EditorPatch>
@@ -83,6 +83,14 @@ export default function PortableTextInput(props: Props) {
   const [isActive, setIsActive] = useState(false)
   const [objectEditData, setObjectEditData]: [ObjectEditData, any] = useState(null)
   const [initialSelection, setInitialSelection] = useState(undefined)
+  const popoverReferenceElement = useRef<HTMLElement | null>(null)
+
+  const setPopupReferenceElement = (key: string) => {
+    const refElement = document.querySelectorAll(`[data-pte-key="${key}"]`)[0]
+    if (refElement) {
+      popoverReferenceElement.current = refElement as HTMLElement
+    }
+  }
 
   // This will open the editing interfaces automatically according to the focusPath.
   // eslint-disable-next-line complexity
@@ -91,18 +99,19 @@ export default function PortableTextInput(props: Props) {
       const isChild = focusPath[1] === 'children'
       const isMarkdef = focusPath[1] === 'markDefs'
       const blockSegment = focusPath[0]
-
       // Annotation focus paths
       if (isMarkdef && isKeySegment(blockSegment)) {
-        const block = value && value.find(blk => blk._key === blockSegment._key)
+        // Get block from the editor value - as the props value may not be updated yet.
+        const block = (PortableTextEditor.getValue(editor) || []).find(
+          blk => blk._key === blockSegment._key
+        )
         const markDefSegment = focusPath[2]
-        // eslint-disable-next-line max-depth
         if (block && isKeySegment(markDefSegment)) {
           const span = block.children.find(
             child => Array.isArray(child.marks) && child.marks.includes(markDefSegment._key)
           )
-          // eslint-disable-next-line max-depth
           if (span) {
+            setPopupReferenceElement(span._key)
             const spanPath = [blockSegment, 'children', {_key: span._key}]
             setIsActive(true)
             PortableTextEditor.select(editor, {
@@ -128,6 +137,8 @@ export default function PortableTextInput(props: Props) {
           kind = 'inlineObject'
           path = path.concat(focusPath.slice(1, 3))
         }
+        const lastSemgent = path.slice(-1)[0] as KeyedSegment
+        setPopupReferenceElement(lastSemgent._key)
         setIsActive(true)
         PortableTextEditor.select(editor, {
           focus: {path, offset: 0},
@@ -242,7 +253,7 @@ export default function PortableTextInput(props: Props) {
   function renderChild(child, childType, attributes, defaultRender) {
     const isSpan = child._type === ptFeatures.types.span.name
     if (isSpan) {
-      return defaultRender(child)
+      return <span data-pte-key={child._key}>{defaultRender(child)}</span>
     }
     // eslint-disable-next-line react/prop-types
     const inlineMarkers = markers.filter(
@@ -309,6 +320,7 @@ export default function PortableTextInput(props: Props) {
         onClose={handleClose}
         onFocus={handleEditObjectFormBuilderFocus}
         readOnly={readOnly}
+        popoverReferenceElement={popoverReferenceElement}
         presence={presence}
         value={value}
       />

packages/@sanity/form-builder/src/inputs/PortableText/Objects/BlockObject.tsx

@@ -80,7 +80,7 @@ export const BlockObject: FunctionComponent<Props> = ({
     )
   }, [value, readOnly])
   return (
-    <div className={classnames} onDoubleClick={handleClickToOpen}>
+    <div className={classnames} onDoubleClick={handleClickToOpen} data-pte-key={value._key}>
       <div className={styles.previewContainer} style={readOnly ? {cursor: 'default'} : {}}>
         {blockPreview}
       </div>

packages/@sanity/form-builder/src/inputs/PortableText/Objects/EditObject.tsx

@@ -33,6 +33,7 @@ interface Props {
   onFocus: (path: Path) => void
   presence: FormFieldPresence[]
   readOnly: boolean
+  popoverReferenceElement: React.MutableRefObject<HTMLElement>
   value: PortableTextBlock[] | undefined
 }
 
@@ -47,6 +48,7 @@ export const EditObject = ({
   onFocus,
   presence,
   readOnly,
+  popoverReferenceElement,
   value
 }: Props) => {
   const editor = usePortableTextEditor()
@@ -173,11 +175,6 @@ export const EditObject = ({
     return null
   }
 
-  const editorElement: HTMLElement = PortableTextEditor.findDOMNode(
-    editor,
-    child ? child : block
-  ) as HTMLElement
-
   // Render the various editing interfaces
   if (editModalLayout === 'fullscreen') {
     return (
@@ -196,7 +193,6 @@ export const EditObject = ({
       />
     )
   }
-
   if (editModalLayout === 'popover' || kind === 'annotation') {
     return (
       <PopoverObjectEditing
@@ -209,8 +205,8 @@ export const EditObject = ({
         onFocus={onFocus}
         path={formBuilderPath}
         presence={presence}
+        popoverReferenceElement={popoverReferenceElement}
         readOnly={readOnly}
-        referenceElement={editorElement}
         type={type}
       />
     )

packages/@sanity/form-builder/src/inputs/PortableText/Objects/InlineObject.tsx

@@ -48,7 +48,7 @@ export const InlineObject: FunctionComponent<Props> = ({
   const isEmpty = !value || isEqual(Object.keys(value), ['_key', '_type'])
 
   return (
-    <span className={classnames} onClick={handleOpen}>
+    <span className={classnames} onClick={handleOpen} data-pte-key={value._key}>
       <span
         className={styles.previewContainer}
         style={readOnly ? {cursor: 'default'} : {}} // TODO: Probably move to styles aka. className?

packages/@sanity/form-builder/src/inputs/PortableText/Objects/renderers/PopoverObjectEditing.tsx

@@ -1,5 +1,5 @@
 /* eslint-disable react/prop-types */
-import React, {FunctionComponent, useMemo} from 'react'
+import React, {FunctionComponent} from 'react'
 
 import Popover from 'part:@sanity/components/dialogs/popover'
 import Stacked from 'part:@sanity/components/utilities/stacked'
@@ -20,9 +20,9 @@ interface Props {
   onClose: (event: React.SyntheticEvent) => void
   onFocus: (arg0: Path) => void
   path: Path
+  popoverReferenceElement: React.MutableRefObject<HTMLElement>
   presence: FormFieldPresence[]
   readOnly: boolean
-  referenceElement: HTMLElement
   type: Type
 }
 
@@ -37,17 +37,16 @@ export const PopoverObjectEditing: FunctionComponent<Props> = ({
   path,
   presence,
   readOnly,
-  referenceElement,
+  popoverReferenceElement,
   type
 }) => {
   const handleChange = (patchEvent: PatchEvent): void => onChange(patchEvent, path)
-  const element = useMemo(() => referenceElement, [])
   return (
     <Stacked>
       {() => (
         <Popover
           placement="bottom"
-          referenceElement={element}
+          referenceElement={popoverReferenceElement.current}
           onClickOutside={onClose}
           onEscape={onClose}
           onClose={onClose}

packages/@sanity/form-builder/src/inputs/PortableText/Text/Annotation.tsx

@@ -4,9 +4,9 @@ import classNames from 'classnames'
 import {PortableTextChild, Type, RenderAttributes} from '@sanity/portable-text-editor'
 
 import {FOCUS_TERMINATOR} from '@sanity/util/paths'
+import {Path} from '@sanity/types'
 import {PatchEvent} from '../../../PatchEvent'
 import {Marker} from '../../../typedefs'
-import {Path} from '@sanity/types'
 
 import styles from './Annotation.css'
 

packages/@sanity/form-builder/src/inputs/PortableText/Toolbar/helpers.tsx

@@ -232,7 +232,10 @@ export function getBlockStyleSelectProps(
   }
 }
 
-function getInsertMenuIcon(type: Type, fallbackIcon: Function): React.ComponentType {
+function getInsertMenuIcon(
+  type: Type,
+  fallbackIcon: () => React.ReactElement
+): React.ComponentType {
   const referenceIcon = get(type, 'to[0].icon')
 
   return type.icon || (type.type && type.type.icon) || referenceIcon || fallbackIcon