Add special cases

packages/playground/data-liberation/src/import/WP_Topological_Sorter.php

@@ -50,8 +50,8 @@ public function map_category( $byte_offset, $data ) {
 		}
 
 		$this->categories[ $data['slug'] ] = array(
-			'byte_offset' => $byte_offset,
 			'parent'      => $data['parent'],
+			'byte_offset' => $byte_offset,
 			'visited'     => false,
 		);
 	}
@@ -84,20 +84,34 @@ public function map_post( $byte_offset, $data ) {
 	}
 
 	/**
-	 * Get the byte offset of an element.
+	 * Get the byte offset of an element, and remove it from the list.
 	 */
 	public function get_byte_offset( $id ) {
 		if ( ! $this->sorted ) {
 			return false;
 		}
 
 		if ( isset( $this->posts[ $id ] ) ) {
-			return $this->posts[ $id ];
+			$ret = $this->posts[ $id ];
+
+			// Remove the element from the array.
+			unset( $this->posts[ $id ] );
+
+			if ( 0 === count( $this->posts ) ) {
+				// All posts have been processed.
+				$this->reset();
+			}
+
+			return $ret;
 		}
 
 		return false;
 	}
 
+	public function is_sorted() {
+		return $this->sorted;
+	}
+
 	/**
 	 * Sort posts topologically.
 	 *
@@ -106,21 +120,23 @@ public function get_byte_offset( $id ) {
 	 *
 	 * Sorted posts will be stored as attachments and posts/pages separately.
 	 */
-	public function sort_topologically() {
+	public function sort_topologically( $empty_memory = true ) {
 		foreach ( $this->categories as $slug => $category ) {
 			$this->topological_category_sort( $slug, $category );
 		}
 
 		$this->sort_parent_child( $this->posts );
 
 		// Empty some memory.
-		foreach ( $this->posts as $id => $element ) {
-			if ( ! $element[2] ) {
-				// The element have not been moved, unset it.
-				unset( $this->posts[ $id ] );
-			} else {
-				// Save only the byte offset.
-				$this->posts[ $id ] = $element[1];
+		if ( $empty_memory ) {
+			foreach ( $this->posts as $id => $element ) {
+				if ( ! $element[2] ) {
+					// The element have not been moved, unset it.
+					unset( $this->posts[ $id ] );
+				} else {
+					// Save only the byte offset.
+					$this->posts[ $id ] = $element[1];
+				}
 			}
 		}
 
@@ -137,8 +153,29 @@ public function sort_topologically() {
 	 */
 	private function sort_parent_child( &$elements ) {
 		// Sort the array in-place.
-		reset( $elements );
-		$position = key( $elements );
+		// reset( $elements );
+		$position = 0; // key( $elements );
+		$length   = count( $elements );
+
+		if ( $length < 2 ) {
+			// No need to sort.
+			return;
+		}
+
+		if ( 2 === $length ) {
+			$keys = array_keys( $elements );
+
+			// First element has a parent and is the second.
+			if ( $elements[ $keys[0] ][0] && $keys[1] === $elements[ $keys[0] ][0] ) {
+				// Swap.
+				$elements = array_reverse( $elements, true );
+
+				// Set the second as 'moved'.
+				$elements[ $keys[1] ][2] = true;
+			}
+
+			return;
+		}
 
 		foreach ( $elements as $id => $element ) {
 			if ( empty( $element[0] ) ) {
@@ -163,7 +200,7 @@ private function move_element( &$elements, $id, &$position ) {
 
 		$element = $elements[ $id ];
 
-		if ( $id <= $position ) {
+		if ( $id < $position ) {
 			// Already in the correct position.
 			return;
 		}

packages/playground/data-liberation/tests/WPTopologicalSorterTests.php

@@ -18,13 +18,13 @@ public function test_import_one_post() {
 	public function test_parent_after_child() {
 		$sorter = new WP_Topological_Sorter();
 
-		$sorter->map_post( 0, $this->generate_post( 1, 2 ) );
-		$sorter->map_post( 1, $this->generate_post( 2, 0 ) );
+		$sorter->map_post( 10, $this->generate_post( 1, 2 ) );
+		$sorter->map_post( 20, $this->generate_post( 2, 0 ) );
 		$sorter->sort_topologically();
 
-		$this->assertEquals( array( 2 => 1 ), $sorter->posts );
+		$this->assertEquals( array( 2 => 20 ), $sorter->posts );
 		$this->assertFalse( $sorter->get_byte_offset( 1 ) );
-		$this->assertEquals( 1, $sorter->get_byte_offset( 2 ) );
+		$this->assertEquals( 20, $sorter->get_byte_offset( 2 ) );
 	}
 
 	public function test_child_after_parent() {
@@ -58,7 +58,7 @@ public function test_chain_parent_child_after() {
 		$sorter->map_post( 30, $this->generate_post( 3, 0 ) );
 		$sorter->sort_topologically();
 
-		$this->assertEquals( array( 3 => 30 ), $sorter->posts );
+		$this->assertEquals( array( 3 => 30, 2 => 20 ), $sorter->posts );
 	}
 
 	public function test_reverse_order() {
@@ -70,6 +70,22 @@ public function test_reverse_order() {
 		$this->assertEquals( array(), $sorter->posts );
 	}
 
+	public function test_get_byte_offsets_consume_array() {
+		$sorter = new WP_Topological_Sorter();
+
+		$this->multiple_map_posts( $sorter, array( 3, 1, 2 ) );
+		$sorter->sort_topologically();
+
+		$this->assertEquals( array( 3 => 10 ), $sorter->posts );
+
+		// $this->assertEquals( 10, $sorter->get_byte_offset( 1 ) );
+		// $this->assertEquals( 20, $sorter->get_byte_offset( 2 ) );
+		// $this->assertEquals( 30, $sorter->get_byte_offset( 3 ) );
+
+		$this->assertFalse( $sorter->get_byte_offset( 1 ) );
+		$this->assertFalse( $sorter->is_sorted() );
+	}
+
 	/**
 	 * This map a list of posts [3, 2, 1] of the form:
 	 *   post_id: 1, 2, 3
@@ -79,7 +95,7 @@ public function test_reverse_order() {
 	private function multiple_map_posts( $sorter, $parents ) {
 		foreach ( $parents as $i => $parent ) {
 			$post = $this->generate_post( $i + 1, $parent );
-			$sorter->map_post( 10 * $parent + 10, $post );
+			$sorter->map_post( 10 * $i + 10, $post );
 		}
 	}