Content Model: pass out segment nodes

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/context/defaultContentModelHandlers.ts

@@ -3,9 +3,9 @@ import { handleBlock } from '../handlers/handleBlock';
 import { handleBlockGroupChildren } from '../handlers/handleBlockGroupChildren';
 import { handleBr } from '../handlers/handleBr';
 import { handleDivider } from '../handlers/handleDivider';
-import { handleEntity } from '../handlers/handleEntity';
+import { handleEntityBlock, handleEntitySegment } from '../handlers/handleEntity';
 import { handleFormatContainer } from '../handlers/handleFormatContainer';
-import { handleGeneralModel } from '../handlers/handleGeneralModel';
+import { handleGeneralBlock, handleGeneralSegment } from '../handlers/handleGeneralModel';
 import { handleImage } from '../handlers/handleImage';
 import { handleList } from '../handlers/handleList';
 import { handleListItem } from '../handlers/handleListItem';
@@ -22,8 +22,10 @@ export const defaultContentModelHandlers: ContentModelHandlerMap = {
     block: handleBlock,
     blockGroupChildren: handleBlockGroupChildren,
     br: handleBr,
-    entity: handleEntity,
-    general: handleGeneralModel,
+    entityBlock: handleEntityBlock,
+    entitySegment: handleEntitySegment,
+    generalBlock: handleGeneralBlock,
+    generalSegment: handleGeneralSegment,
     divider: handleDivider,
     image: handleImage,
     list: handleList,

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleBlock.ts

@@ -24,15 +24,15 @@ export const handleBlock: ContentModelBlockHandler<ContentModelBlock> = (
             refNode = handlers.paragraph(doc, parent, block, context, refNode);
             break;
         case 'Entity':
-            refNode = handlers.entity(doc, parent, block, context, refNode);
+            refNode = handlers.entityBlock(doc, parent, block, context, refNode);
             break;
         case 'Divider':
             refNode = handlers.divider(doc, parent, block, context, refNode);
             break;
         case 'BlockGroup':
             switch (block.blockGroupType) {
                 case 'General':
-                    refNode = handlers.general(doc, parent, block, context, refNode);
+                    refNode = handlers.generalBlock(doc, parent, block, context, refNode);
                     break;
 
                 case 'FormatContainer':

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleBr.ts

@@ -1,23 +1,20 @@
+import { ContentModelBr, ContentModelSegmentHandler } from 'roosterjs-content-model-types';
 import { handleSegmentCommon } from '../utils/handleSegmentCommon';
-import {
-    ContentModelBr,
-    ContentModelHandler,
-    ModelToDomContext,
-} from 'roosterjs-content-model-types';
 
 /**
  * @internal
  */
-export const handleBr: ContentModelHandler<ContentModelBr> = (
-    doc: Document,
-    parent: Node,
-    segment: ContentModelBr,
-    context: ModelToDomContext
+export const handleBr: ContentModelSegmentHandler<ContentModelBr> = (
+    doc,
+    parent,
+    segment,
+    context,
+    segmentNodes
 ) => {
     const br = doc.createElement('br');
     const element = doc.createElement('span');
     element.appendChild(br);
     parent.appendChild(element);
 
-    handleSegmentCommon(doc, br, element, segment, context);
+    handleSegmentCommon(doc, br, element, segment, context, segmentNodes);
 };

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleEntity.ts

@@ -1,31 +1,68 @@
+import { addDelimiters, commitEntity, getObjectKeys, wrap } from 'roosterjs-editor-dom';
 import { applyFormat } from '../utils/applyFormat';
 import { Entity } from 'roosterjs-editor-types';
 import { reuseCachedElement } from '../utils/reuseCachedElement';
 import {
     ContentModelBlockHandler,
     ContentModelEntity,
+    ContentModelSegmentHandler,
     ModelToDomContext,
 } from 'roosterjs-content-model-types';
-import {
-    addDelimiters,
-    commitEntity,
-    getObjectKeys,
-    isBlockElement,
-    wrap,
-} from 'roosterjs-editor-dom';
 
 /**
  * @internal
  */
-export const handleEntity: ContentModelBlockHandler<ContentModelEntity> = (
-    doc: Document,
-    parent: Node,
-    entityModel: ContentModelEntity,
-    context: ModelToDomContext,
-    refNode: Node | null
+export const handleEntityBlock: ContentModelBlockHandler<ContentModelEntity> = (
+    _,
+    parent,
+    entityModel,
+    context,
+    refNode
+) => {
+    const wrapper = preprocessEntity(entityModel, context);
+
+    refNode = reuseCachedElement(parent, wrapper, refNode);
+    context.onNodeCreated?.(entityModel, wrapper);
+
+    return refNode;
+};
+
+/**
+ * @internal
+ */
+export const handleEntitySegment: ContentModelSegmentHandler<ContentModelEntity> = (
+    _,
+    parent,
+    entityModel,
+    context,
+    newSegments
 ) => {
-    const { id, type, isReadonly, format } = entityModel;
-    let wrapper = entityModel.wrapper;
+    const wrapper = preprocessEntity(entityModel, context);
+    const { format, isReadonly } = entityModel;
+
+    parent.appendChild(wrapper);
+    newSegments?.push(wrapper);
+
+    if (getObjectKeys(format).length > 0) {
+        const span = wrap(wrapper, 'span');
+
+        applyFormat(span, context.formatAppliers.segment, format, context);
+    }
+
+    if (context.addDelimiterForEntity && isReadonly) {
+        const [after, before] = addDelimiters(wrapper);
+
+        newSegments?.push(after, before);
+        context.regularSelection.current.segment = after;
+    } else {
+        context.regularSelection.current.segment = wrapper;
+    }
+
+    context.onNodeCreated?.(entityModel, wrapper);
+};
+
+function preprocessEntity(entityModel: ContentModelEntity, context: ModelToDomContext) {
+    let { id, type, isReadonly, wrapper } = entityModel;
 
     if (!context.allowCacheElement) {
         wrapper = wrapper.cloneNode(true /*deep*/) as HTMLElement;
@@ -42,30 +79,10 @@ export const handleEntity: ContentModelBlockHandler<ContentModelEntity> = (
                   isReadonly: !!isReadonly,
               }
             : null;
-    const isInlineEntity = !isBlockElement(wrapper);
 
     if (entity) {
         // Commit the entity attributes in case there is any change
         commitEntity(wrapper, entity.type, entity.isReadonly, entity.id);
     }
-
-    refNode = reuseCachedElement(parent, wrapper, refNode);
-
-    if (isInlineEntity && getObjectKeys(format).length > 0) {
-        const span = wrap(wrapper, 'span');
-
-        applyFormat(span, context.formatAppliers.segment, format, context);
-    }
-
-    if (context.addDelimiterForEntity && isInlineEntity && isReadonly) {
-        const [after] = addDelimiters(wrapper);
-
-        context.regularSelection.current.segment = after;
-    } else if (isInlineEntity) {
-        context.regularSelection.current.segment = wrapper;
-    }
-
-    context.onNodeCreated?.(entityModel, wrapper);
-
-    return refNode;
-};
+    return wrapper;
+}

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleGeneralModel.ts

@@ -1,24 +1,24 @@
 import { handleSegmentCommon } from '../utils/handleSegmentCommon';
-import { isGeneralSegment } from '../../modelApi/common/isGeneralSegment';
 import { isNodeOfType } from '../../domUtils/isNodeOfType';
 import { NodeType } from 'roosterjs-editor-types';
 import { reuseCachedElement } from '../utils/reuseCachedElement';
 import { wrap } from 'roosterjs-editor-dom';
 import {
     ContentModelBlockHandler,
     ContentModelGeneralBlock,
-    ModelToDomContext,
+    ContentModelGeneralSegment,
+    ContentModelSegmentHandler,
 } from 'roosterjs-content-model-types';
 
 /**
  * @internal
  */
-export const handleGeneralModel: ContentModelBlockHandler<ContentModelGeneralBlock> = (
-    doc: Document,
-    parent: Node,
-    group: ContentModelGeneralBlock,
-    context: ModelToDomContext,
-    refNode: Node | null
+export const handleGeneralBlock: ContentModelBlockHandler<ContentModelGeneralBlock> = (
+    doc,
+    parent,
+    group,
+    context,
+    refNode
 ) => {
     let node: Node = group.element;
 
@@ -31,15 +31,32 @@ export const handleGeneralModel: ContentModelBlockHandler<ContentModelGeneralBlo
         parent.insertBefore(node, refNode);
     }
 
-    if (isGeneralSegment(group) && isNodeOfType(node, NodeType.Element)) {
+    context.onNodeCreated?.(group, node);
+    context.modelHandlers.blockGroupChildren(doc, node, group, context);
+
+    return refNode;
+};
+
+/**
+ * @internal
+ */
+export const handleGeneralSegment: ContentModelSegmentHandler<ContentModelGeneralSegment> = (
+    doc,
+    parent,
+    group,
+    context,
+    segmentNodes
+) => {
+    const node = group.element.cloneNode() as HTMLElement;
+    group.element = node;
+    parent.appendChild(node);
+
+    if (isNodeOfType(node, NodeType.Element)) {
         const element = wrap(node, 'span');
 
-        handleSegmentCommon(doc, node, element, group, context);
-    } else {
+        handleSegmentCommon(doc, node, element, group, context, segmentNodes);
         context.onNodeCreated?.(group, node);
     }
 
     context.modelHandlers.blockGroupChildren(doc, node, group, context);
-
-    return refNode;
 };

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleImage.ts

@@ -1,20 +1,17 @@
 import { applyFormat } from '../utils/applyFormat';
+import { ContentModelImage, ContentModelSegmentHandler } from 'roosterjs-content-model-types';
 import { handleSegmentCommon } from '../utils/handleSegmentCommon';
 import { parseValueWithUnit } from '../../formatHandlers/utils/parseValueWithUnit';
-import {
-    ContentModelHandler,
-    ContentModelImage,
-    ModelToDomContext,
-} from 'roosterjs-content-model-types';
 
 /**
  * @internal
  */
-export const handleImage: ContentModelHandler<ContentModelImage> = (
-    doc: Document,
-    parent: Node,
-    imageModel: ContentModelImage,
-    context: ModelToDomContext
+export const handleImage: ContentModelSegmentHandler<ContentModelImage> = (
+    doc,
+    parent,
+    imageModel,
+    context,
+    segmentNodes
 ) => {
     const img = doc.createElement('img');
     const element = document.createElement('span');
@@ -53,5 +50,5 @@ export const handleImage: ContentModelHandler<ContentModelImage> = (
         };
     }
 
-    handleSegmentCommon(doc, img, element, imageModel, context);
+    handleSegmentCommon(doc, img, element, imageModel, context, segmentNodes);
 };

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleParagraph.ts

@@ -66,12 +66,14 @@ export const handleParagraph: ContentModelBlockHandler<ContentModelParagraph> =
                                 segmentType: 'Text',
                                 text: '',
                             },
-                            context
+                            context,
+                            []
                         );
                     }
 
                     paragraph.segments.forEach(segment => {
-                        context.modelHandlers.segment(doc, parent, segment, context);
+                        const newSegments: Node[] = [];
+                        context.modelHandlers.segment(doc, parent, segment, context, newSegments);
                     });
                 }
             };

packages-content-model/roosterjs-content-model-dom/lib/modelToDom/handlers/handleSegment.ts

@@ -1,17 +1,14 @@
-import {
-    ContentModelHandler,
-    ContentModelSegment,
-    ModelToDomContext,
-} from 'roosterjs-content-model-types';
+import { ContentModelSegment, ContentModelSegmentHandler } from 'roosterjs-content-model-types';
 
 /**
  * @internal
  */
-export const handleSegment: ContentModelHandler<ContentModelSegment> = (
-    doc: Document,
-    parent: Node,
-    segment: ContentModelSegment,
-    context: ModelToDomContext
+export const handleSegment: ContentModelSegmentHandler<ContentModelSegment> = (
+    doc,
+    parent,
+    segment,
+    context,
+    segmentNodes
 ) => {
     const regularSelection = context.regularSelection;
 
@@ -24,23 +21,23 @@ export const handleSegment: ContentModelHandler<ContentModelSegment> = (
 
     switch (segment.segmentType) {
         case 'Text':
-            context.modelHandlers.text(doc, parent, segment, context);
+            context.modelHandlers.text(doc, parent, segment, context, segmentNodes);
             break;
 
         case 'Br':
-            context.modelHandlers.br(doc, parent, segment, context);
+            context.modelHandlers.br(doc, parent, segment, context, segmentNodes);
             break;
 
         case 'Image':
-            context.modelHandlers.image(doc, parent, segment, context);
+            context.modelHandlers.image(doc, parent, segment, context, segmentNodes);
             break;
 
         case 'General':
-            context.modelHandlers.general(doc, parent, segment, context, null /*refNode*/);
+            context.modelHandlers.generalSegment(doc, parent, segment, context, segmentNodes);
             break;
 
         case 'Entity':
-            context.modelHandlers.entity(doc, parent, segment, context, null /*refNode*/);
+            context.modelHandlers.entitySegment(doc, parent, segment, context, segmentNodes);
             break;
     }