fix: linked nodes hydration

jest/setup.js

@@ -2,3 +2,5 @@ const Enzyme = require('enzyme');
 const Adapter = require('enzyme-adapter-react-16');
 
 Enzyme.configure({ adapter: new Adapter() });
+
+jest.spyOn(console, 'error').mockImplementation(() => {});

packages/core/src/editor/actions.ts

@@ -65,19 +65,25 @@ export const Actions = (
       addNodeToParentAtIndex(node, parentId, index);
     }
 
-    if (!node.data.nodes) {
-      return;
+    if (node.data.nodes) {
+      const childToAdd = [...node.data.nodes];
+      node.data.nodes = [];
+      childToAdd.forEach((childId, index) =>
+        addTreeToParentAtIndex(
+          { rootNodeId: childId, nodes: tree.nodes },
+          node.id,
+          index
+        )
+      );
+    }
+
+    if (node.data.linkedNodes) {
+      Object.keys(node.data.linkedNodes).forEach((linkedId) => {
+        const nodeId = node.data.linkedNodes[linkedId];
+        state.nodes[nodeId] = tree.nodes[nodeId];
+        addTreeToParentAtIndex({ rootNodeId: nodeId, nodes: tree.nodes });
+      });
     }
-    // we need to deep clone here...
-    const childToAdd = [...node.data.nodes];
-    node.data.nodes = [];
-    childToAdd.forEach((childId, index) =>
-      addTreeToParentAtIndex(
-        { rootNodeId: childId, nodes: tree.nodes },
-        node.id,
-        index
-      )
-    );
   };
 
   const getParentAndValidate = (parentId: NodeId): Node => {
@@ -87,6 +93,32 @@ export const Actions = (
     return parent;
   };
 
+  const deleteNode = (id: NodeId, isLinkedNode: boolean = false) => {
+    const targetNode = state.nodes[id],
+      parentNode = state.nodes[targetNode.data.parent];
+
+    if (targetNode.data.nodes) {
+      // we deep clone here because otherwise immer will mutate the node
+      // object as we remove nodes
+      [...targetNode.data.nodes].forEach((childId) => deleteNode(childId));
+    }
+
+    if (isLinkedNode && parentNode.data.linkedNodes) {
+      const linkedId = Object.keys(parentNode.data.linkedNodes).filter(
+        (id) => parentNode.data.linkedNodes[id] === id
+      )[0];
+      if (linkedId) {
+        delete parentNode.data.linkedNodes[linkedId];
+      }
+    } else {
+      const parentChildren = parentNode.data.nodes;
+      parentChildren.splice(parentChildren.indexOf(id), 1);
+    }
+
+    updateEventsNode(state, id, true);
+    delete state.nodes[id];
+  };
+
   return {
     /**
      * @private
@@ -103,6 +135,11 @@ export const Actions = (
         parent.data.linkedNodes = {};
       }
 
+      const existingLinkedNode = parent.data.linkedNodes[id];
+      if (existingLinkedNode) {
+        deleteNode(existingLinkedNode, true);
+      }
+
       parent.data.linkedNodes[id] = tree.rootNodeId;
 
       tree.nodes[tree.rootNodeId].data.parent = parentId;
@@ -160,18 +197,7 @@ export const Actions = (
     delete(id: NodeId) {
       invariant(!query.node(id).isTopLevelNode(), ERROR_DELETE_TOP_LEVEL_NODE);
 
-      const targetNode = state.nodes[id];
-      if (targetNode.data.nodes) {
-        // we deep clone here because otherwise immer will mutate the node
-        // object as we remove nodes
-        [...targetNode.data.nodes].forEach((childId) => this.delete(childId));
-      }
-
-      const parentChildren = state.nodes[targetNode.data.parent].data.nodes;
-      parentChildren.splice(parentChildren.indexOf(id), 1);
-
-      updateEventsNode(state, id, true);
-      delete state.nodes[id];
+      deleteNode(id);
     },
 
     deserialize(input: SerializedNodes | string) {

packages/core/src/nodes/Element.tsx

@@ -29,11 +29,11 @@ export type Element<T extends React.ElementType> = {
 export function Element<T extends React.ElementType>({
   id,
   children,
-  ...otherProps
+  ...elementProps
 }: Element<T>) {
-  const props = {
+  const { is, custom, canvas, ...otherProps } = {
     ...defaultElementProps,
-    ...otherProps,
+    ...elementProps,
   };
 
   const { query, actions } = useInternalEditor();
@@ -44,55 +44,52 @@ export function Element<T extends React.ElementType>({
     },
   }));
 
-  const [internalId, setInternalId] = useState<NodeId | null>(null);
-  const [initialised, setInitialised] = useState(false);
+  const [linkedNodeId, setLinkedNodeId] = useState<NodeId | null>(null);
 
   useEffectOnce(() => {
-    invariant(id !== null, ERROR_TOP_LEVEL_ELEMENT_NO_ID);
+    invariant(!!id, ERROR_TOP_LEVEL_ELEMENT_NO_ID);
     const { id: nodeId, data } = node;
 
     if (inNodeContext) {
-      let internalId,
-        newProps = props;
+      let linkedNodeId;
 
       const existingNode =
         data.linkedNodes &&
         data.linkedNodes[id] &&
         query.node(data.linkedNodes[id]).get();
 
-      if (
-        existingNode &&
-        existingNode.data.type === props.is &&
-        typeof props.is !== 'string'
-      ) {
-        newProps = {
-          ...newProps,
+      // Render existing linked Node if it already exists (and is the same type as the JSX)
+      if (existingNode && existingNode.data.type === is) {
+        linkedNodeId = existingNode.id;
+
+        // Merge JSX and existing props
+        const mergedProps = {
           ...existingNode.data.props,
+          ...otherProps,
         };
-      }
-
-      const linkedElement = React.createElement(Element, newProps, children);
 
-      const tree = query
-        .parseReactElement(linkedElement)
-        .toNodeTree((node, jsx) => {
-          if (jsx === linkedElement) {
-            node.id = existingNode ? existingNode.id : node.id;
-            node.data = {
-              ...(existingNode ? existingNode.data.props : {}),
-              ...node.data,
-            };
-          }
-        });
-
-      internalId = tree.rootNodeId;
-      actions.addLinkedNodeFromTree(tree, nodeId, id);
+        actions.setProp(linkedNodeId, (props) =>
+          Object.keys(mergedProps).forEach(
+            (key) => (props[key] = mergedProps[key])
+          )
+        );
+      } else {
+        // otherwise, create and render a new linked Node
+        const linkedElement = React.createElement(
+          Element,
+          elementProps,
+          children
+        );
+
+        const tree = query.parseReactElement(linkedElement).toNodeTree();
+
+        linkedNodeId = tree.rootNodeId;
+        actions.addLinkedNodeFromTree(tree, nodeId, id);
+      }
 
-      setInternalId(internalId);
+      setLinkedNodeId(linkedNodeId);
     }
-
-    setInitialised(true);
   });
 
-  return initialised ? <NodeElement id={internalId} /> : null;
+  return linkedNodeId ? <NodeElement id={linkedNodeId} /> : null;
 }

packages/core/src/nodes/tests/Element.test.tsx

@@ -0,0 +1,127 @@
+import React from 'react';
+import { mount } from 'enzyme';
+import { Element } from '../Element';
+import { createTestNode } from '../../utils/createTestNode';
+
+let parentNode;
+let existingLinkedNode;
+let newLinkedNode = createTestNode('newLinkedNode');
+
+let toNodeTree = jest.fn().mockImplementation(() => ({
+  rootNodeId: newLinkedNode.id,
+}));
+
+let addLinkedNodeFromTree = jest.fn();
+let setProp = jest.fn();
+let parseReactElement = jest.fn().mockImplementation(() => ({
+  toNodeTree,
+}));
+
+jest.mock('../../editor/useInternalEditor', () => ({
+  useInternalEditor: () => ({
+    actions: {
+      addLinkedNodeFromTree,
+      setProp,
+    },
+    query: {
+      parseReactElement,
+      node: jest.fn().mockImplementation((id) => ({
+        get() {
+          if (id === 'parent-node') return parentNode;
+          else return existingLinkedNode;
+        },
+      })),
+    },
+  }),
+}));
+
+jest.mock('../useInternalNode', () => ({
+  useInternalNode: () => ({
+    node: parentNode,
+    inNodeContext: true,
+  }),
+}));
+
+const NodeElementTest = jest.fn().mockImplementation(() => null);
+
+jest.mock('../NodeElement', () => ({
+  NodeElement: jest.fn().mockImplementation((props) => NodeElementTest(props)),
+}));
+
+describe('<Element />', () => {
+  beforeEach(() => {
+    parentNode = createTestNode('test');
+  });
+
+  describe('when no id is passed', () => {
+    it('should throw error', () => {
+      expect(() => mount(<Element />)).toThrow();
+    });
+  });
+
+  describe('when there is no existing node', () => {
+    let elementProps, children;
+    beforeEach(() => {
+      elementProps = {
+        color: '#fff',
+      };
+
+      children = <h1>Child</h1>;
+      mount(
+        <Element id="test" {...elementProps}>
+          {children}
+        </Element>
+      );
+    });
+
+    it('should call query.parseReactElement()', () => {
+      expect(parseReactElement).toHaveBeenCalledWith(
+        <Element {...elementProps}>{children}</Element>
+      );
+    });
+    it('should call actions.addLinkedNodeFromTree()', () => {
+      expect(addLinkedNodeFromTree).toHaveBeenCalled();
+    });
+    it('should render a new linked Node', () => {
+      expect(NodeElementTest).toHaveBeenCalledWith({
+        id: newLinkedNode.id,
+      });
+    });
+  });
+
+  describe('when there is an existing node', () => {
+    beforeEach(() => {
+      existingLinkedNode = createTestNode('existing-linked-node', {
+        type: 'div',
+      });
+
+      parentNode = createTestNode('parent-node', {
+        linkedNodes: {
+          test: existingLinkedNode.id,
+        },
+      });
+    });
+    describe('when type is the same as JSX', () => {
+      beforeEach(() => {
+        mount(<Element id="test" />);
+      });
+
+      it('should render existing Node', () => {
+        expect(NodeElementTest).toHaveBeenCalledWith({
+          id: existingLinkedNode.id,
+        });
+      });
+    });
+    describe('when type is the different from JSX', () => {
+      beforeEach(() => {
+        mount(<Element id="test" is="h1" />);
+      });
+
+      it('should render a new linked Node', () => {
+        expect(NodeElementTest).toHaveBeenCalledWith({
+          id: newLinkedNode.id,
+        });
+      });
+    });
+  });
+});

packages/core/src/utils/createTestNode.ts

@@ -1,4 +1,4 @@
-export const createTestNode = (id, data, config: any = {}) => {
+export const createTestNode = (id, data = {}, config: any = {}) => {
   return {
     ...config,
     id,