Improve performance of getSelectedBlocks (#17582)

Other (engine): Improve performance of `Selection#getSelectedBlocks` when selection contains block elements with many blocks inside (e.g. table). Closes #17629.

packages/ckeditor5-engine/src/model/selection.ts

@@ -667,12 +667,22 @@ export default class Selection extends /* #__PURE__ */ EmitterMixin( TypeCheckab
 				yield startBlock as any;
 			}
 
-			for ( const value of range.getWalker() ) {
+			const treewalker = range.getWalker();
+
+			for ( const value of treewalker ) {
 				const block = value.item;
 
 				if ( value.type == 'elementEnd' && isUnvisitedTopBlock( block as any, visited, range ) ) {
 					yield block as Element;
 				}
+				// If element is block, we can skip its children and jump to the end of it.
+				else if (
+					value.type == 'elementStart' &&
+					block.is( 'model:element' ) &&
+					block.root.document!.model.schema.isBlock( block )
+				) {
+					treewalker.jumpTo( Position._createAt( block, 'end' ) );
+				}
 			}
 
 			const endBlock = getParentBlock( range.end, visited );

packages/ckeditor5-engine/src/model/treewalker.ts

@@ -183,6 +183,32 @@ export default class TreeWalker implements Iterable<TreeWalkerValue> {
 		}
 	}
 
+	/**
+	 * Moves tree walker {@link #position} to provided `position`. Tree walker will
+	 * continue traversing from that position.
+	 *
+	 * Note: in contrary to {@link ~TreeWalker#skip}, this method does not iterate over the nodes along the way.
+	 * It simply sets the current tree walker position to a new one.
+	 * From the performance standpoint, it is better to use {@link ~TreeWalker#jumpTo} rather than {@link ~TreeWalker#skip}.
+	 *
+	 * If the provided position is before the start boundary, the position will be
+	 * set to the start boundary. If the provided position is after the end boundary,
+	 * the position will be set to the end boundary.
+	 * This is done to prevent the treewalker from traversing outside the boundaries.
+	 *
+	 * @param position Position to jump to.
+	 */
+	public jumpTo( position: Position ): void {
+		if ( this._boundaryStartParent && position.isBefore( this.boundaries!.start ) ) {
+			position = this.boundaries!.start;
+		} else if ( this._boundaryEndParent && position.isAfter( this.boundaries!.end ) ) {
+			position = this.boundaries!.end;
+		}
+
+		this._position = position.clone();
+		this._visitedParent = position.parent;
+	}
+
 	/**
 	 * Gets the next tree walker's value.
 	 */

packages/ckeditor5-engine/src/view/treewalker.ts

@@ -159,6 +159,31 @@ export default class TreeWalker implements IterableIterator<TreeWalkerValue> {
 		}
 	}
 
+	/**
+	 * Moves tree walker {@link #position} to provided `position`. Tree walker will
+	 * continue traversing from that position.
+	 *
+	 * Note: in contrary to {@link ~TreeWalker#skip}, this method does not iterate over the nodes along the way.
+	 * It simply sets the current tree walker position to a new one.
+	 * From the performance standpoint, it is better to use {@link ~TreeWalker#jumpTo} rather than {@link ~TreeWalker#skip}.
+	 *
+	 * If the provided position is before the start boundary, the position will be
+	 * set to the start boundary. If the provided position is after the end boundary,
+	 * the position will be set to the end boundary.
+	 * This is done to prevent the treewalker from traversing outside the boundaries.
+	 *
+	 * @param position Position to jump to.
+	 */
+	public jumpTo( position: Position ): void {
+		if ( this._boundaryStartParent && position.isBefore( this.boundaries!.start ) ) {
+			position = this.boundaries!.start;
+		} else if ( this._boundaryEndParent && position.isAfter( this.boundaries!.end ) ) {
+			position = this.boundaries!.end;
+		}
+
+		this._position = position.clone();
+	}
+
 	/**
 	 * Gets the next tree walker's value.
 	 *

packages/ckeditor5-engine/tests/model/treewalker.js

@@ -720,6 +720,70 @@ describe( 'TreeWalker', () => {
 			} );
 		} );
 	} );
+
+	describe( 'jumpTo', () => {
+		it( 'should jump to the given position', () => {
+			const walker = new TreeWalker( {
+				startPosition: Position._createAt( paragraph, 0 )
+			} );
+
+			walker.jumpTo( new Position( paragraph, [ 2 ] ) );
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 3 );
+		} );
+
+		it( 'cannot move position before the #_boundaryStartParent', () => {
+			const range = new Range(
+				new Position( paragraph, [ 2 ] ),
+				new Position( paragraph, [ 4 ] )
+			);
+			const walker = new TreeWalker( {
+				boundaries: range
+			} );
+
+			const positionBeforeAllowedRange = new Position( paragraph, [ 0 ] );
+
+			walker.jumpTo( positionBeforeAllowedRange );
+
+			// `jumpTo()` autocorrected the position to the first allowed position.
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 3 );
+		} );
+
+		it( 'cannot move position after the #_boundaryEndParent', () => {
+			const range = new Range(
+				new Position( paragraph, [ 0 ] ),
+				new Position( paragraph, [ 2 ] )
+			);
+			const walker = new TreeWalker( {
+				boundaries: range
+			} );
+
+			const positionAfterAllowedRange = new Position( paragraph, [ 4 ] );
+
+			// `jumpTo()` autocorrected the position to the last allowed position.
+			walker.jumpTo( positionAfterAllowedRange );
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+		} );
+	} );
 } );
 
 function expectValue( value, expected, options ) {

packages/ckeditor5-engine/tests/view/treewalker.js

@@ -1216,6 +1216,70 @@ describe( 'TreeWalker', () => {
 			} );
 		} );
 	} );
+
+	describe( 'jumpTo', () => {
+		it( 'should jump to the given position', () => {
+			const walker = new TreeWalker( {
+				startPosition: Position._createAt( paragraph, 0 )
+			} );
+
+			walker.jumpTo( new Position( paragraph, 2 ) );
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( img2 );
+			expect( walker.position.offset ).to.equal( 0 );
+		} );
+
+		it( 'cannot move position before the #_boundaryStartParent', () => {
+			const range = new Range(
+				new Position( paragraph, 2 ),
+				new Position( paragraph, 4 )
+			);
+			const walker = new TreeWalker( {
+				boundaries: range
+			} );
+
+			const positionBeforeAllowedRange = new Position( paragraph, 0 );
+
+			walker.jumpTo( positionBeforeAllowedRange );
+
+			// `jumpTo()` autocorrected the position to the first allowed position.
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( img2 );
+			expect( walker.position.offset ).to.equal( 0 );
+		} );
+
+		it( 'cannot move position after the #_boundaryEndParent', () => {
+			const range = new Range(
+				new Position( paragraph, 0 ),
+				new Position( paragraph, 2 )
+			);
+			const walker = new TreeWalker( {
+				boundaries: range
+			} );
+
+			const positionAfterAllowedRange = new Position( paragraph, 4 );
+
+			// `jumpTo()` autocorrected the position to the last allowed position.
+			walker.jumpTo( positionAfterAllowedRange );
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+
+			walker.next();
+
+			expect( walker.position.parent ).to.equal( paragraph );
+			expect( walker.position.offset ).to.equal( 2 );
+		} );
+	} );
 } );
 
 function expectValue( value, expected, options = {} ) {