Preserve selection when collapsing items in an outline view.

This fixes the error where if a row was selected below a certain item and that item was collapsed, the row selection would remain fixed in place even while the item it was highlighting moved upwards.

AppKit/CPOutlineView.j

@@ -367,6 +367,7 @@ CPOutlineViewDropOnItemIndex = -1;
     [self _noteItemWillCollapse:anItem];
     // Update selections:
     // * Deselect items inside the collapsed item.
+    // * Shift row selections below the collapsed item so that the same logical items remain selected.
     var collapseTopIndex = [self rowForItem:anItem],
         topLevel = [self levelForRow:collapseTopIndex],
         collapseEndIndex = collapseTopIndex;
@@ -377,14 +378,31 @@ CPOutlineViewDropOnItemIndex = -1;
     var collapseRange = CPMakeRange(collapseTopIndex + 1, collapseEndIndex - collapseTopIndex);
     if (collapseRange.length)
     {
-        var selection = [self selectedRowIndexes];
+        var selection = [self selectedRowIndexes],
+            didChange = NO;
+
         if ([selection intersectsIndexesInRange:collapseRange])
         {
-            [self _noteSelectionIsChanging];
             [selection removeIndexesInRange:collapseRange];
+            [self _noteSelectionIsChanging];
+            didChange = YES;
             // Will call _noteSelectionDidChange
             [self _setSelectedRowIndexes:selection];
         }
+
+        // Shift any selected rows below upwards.
+        if ([selection intersectsIndexesInRange:CPMakeRange(collapseEndIndex + 1, _itemsForRows.length)])
+        {
+            // Notify if that wasn't already done above.
+            if (!didChange)
+                [self _noteSelectionIsChanging];
+            didChange = YES;
+
+            [selection shiftIndexesStartingAtIndex:collapseEndIndex + 1 by:-collapseRange.length];
+        }
+
+        if (didChange)
+            [self _setSelectedRowIndexes:selection];
     }
     itemInfo.isExpanded = NO;
 

Tests/AppKit/CPOutlineViewTest.j

@@ -24,6 +24,9 @@
     [outlineView expandItem:nil expandChildren:YES];
 }
 
+/*!
+    Test that entries load and hide correctly.
+*/
 - (void)testCollapse
 {
     // By default all rows should be visible.
@@ -63,6 +66,42 @@
     [self assert:".1.1" equals:[outlineView itemAtRow:[outlineView selectedRow]] message:".1.1 selection should remain"];
 }
 
+/*!
+    Test that the selection stays on the same item below a collapse.
+*/
+- (void)testCollapseWithSelectionBelow
+{
+    // [".1", ".1.1", ".1.2", ".1.2.1", ".1.2.2", ".2", ".3", ".3.1"]
+    var preSelection = [CPIndexSet indexSet];
+    [preSelection addIndex:[outlineView rowForItem:".1.1"]];
+    [preSelection addIndex:[outlineView rowForItem:".3.1"]];
+
+    [outlineView selectRowIndexes:preSelection byExtendingSelection:NO];
+
+    [outlineView collapseItem:".1.2"];
+
+    var afterSelection = [outlineView selectedRowIndexes];
+
+    [self assert:2 equals:[afterSelection count] message:"selections should remain"];
+    [self assert:".1.1" equals:[outlineView itemAtRow:[afterSelection firstIndex]] message:".1.1 selection should remain"];
+    [self assert:".3.1" equals:[outlineView itemAtRow:[afterSelection lastIndex]] message:".3.1 selection should remain"];
+
+    // Collapse where one selection disappears and one shifts.
+    preSelection = [CPIndexSet indexSet];
+    [preSelection addIndex:[outlineView rowForItem:".1.1"]];
+    [preSelection addIndex:[outlineView rowForItem:".3"]];
+
+    [outlineView selectRowIndexes:preSelection byExtendingSelection:NO];
+
+    [outlineView collapseItem:".1"];
+
+    afterSelection = [outlineView selectedRowIndexes];
+
+    [self assert:1 equals:[afterSelection count] message:"1 selection should disappear"];
+
+    [self assert:".3" equals:[outlineView itemAtRow:[afterSelection firstIndex]] message:".3 selection should remain"];
+}
+
 @end
 
 @implementation TestDataSource : CPObject