Mac: Fix kill button and selection in task manager

chrome/browser/cocoa/task_manager_mac.mm

 // Copyright (c) 2010 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome/browser/cocoa/task_manager_mac.h"
 
 #include <algorithm>
 #include <vector>
 
 #include "app/l10n_util_mac.h"
@@ skipping 71 matching lines @@
 };
 
 }  // namespace
 
 @interface TaskManagerWindowController (Private)
 - (NSTableColumn*)addColumnWithId:(int)columnId visible:(BOOL)isVisible;
 - (void)setUpTableColumns;
 - (void)setUpTableHeaderContextMenu;
 - (void)toggleColumn:(id)sender;
 - (void)adjustSelectionAndEndProcessButton;
+- (void)deselectRows;
 @end
 
 ////////////////////////////////////////////////////////////////////////////////
 // TaskManagerWindowController implementation:
 
 @implementation TaskManagerWindowController
 
 - (id)initWithTaskManagerObserver:(TaskManagerMac*)taskManagerObserver {
   NSString* nibpath = [mac_util::MainAppBundle()
                         pathForResource:@"TaskManager"
                                  ofType:@"nib"];
   if ((self = [super initWithWindowNibPath:nibpath owner:self])) {
     taskManagerObserver_ = taskManagerObserver;
     taskManager_ = taskManagerObserver_->task_manager();
     model_ = taskManager_->model();
 
     if (g_browser_process && g_browser_process->local_state()) {
       size_saver_.reset([[WindowSizeAutosaver alloc]
           initWithWindow:[self window]
              prefService:g_browser_process->local_state()
                     path:prefs::kTaskManagerWindowPlacement
                    state:kSaveWindowRect]);
     }
     [self showWindow:self];
   }
   return self;
 }
 
 - (void)sortShuffleArray {
-  indexShuffle_.resize(model_->ResourceCount());
-  for (size_t i = 0; i < indexShuffle_.size(); ++i)
-    indexShuffle_[i] = i;
+  viewToModelMap_.resize(model_->ResourceCount());
+  for (size_t i = 0; i < viewToModelMap_.size(); ++i)
+    viewToModelMap_[i] = i;
 
-  std::sort(indexShuffle_.begin(), indexShuffle_.end(),
+  std::sort(viewToModelMap_.begin(), viewToModelMap_.end(),
             SortHelper(model_, currentSortDescriptor_.get()));
+
+  modelToViewMap_.resize(viewToModelMap_.size());
+  for (size_t i = 0; i < viewToModelMap_.size(); ++i)
+    modelToViewMap_[viewToModelMap_[i]] = i;
 }
 
 - (void)reloadData {
+  // Store old view indices, and the model indices they map to.
+  NSIndexSet* viewSelection = [tableView_ selectedRowIndexes];
+  std::vector<int> modelSelection;
+  for (NSUInteger i = [viewSelection lastIndex];
+       i != NSNotFound;
+       i = [viewSelection indexLessThanIndex:i]) {
+    modelSelection.push_back(viewToModelMap_[i]);
+  }
+
+  // Sort.
   [self sortShuffleArray];
+
+  // Use the model indices to get the new view indices of the selection, and
+  // set selection to that. This assumes that no rows were added or removed
+  // (in that case, the selection is cleared before -reloadData is called).
+  if (modelSelection.size() > 0)
+    DCHECK_EQ([tableView_ numberOfRows], model_->ResourceCount());
+  NSMutableIndexSet* indexSet = [NSMutableIndexSet indexSet];
+  for (size_t i = 0; i < modelSelection.size(); ++i)
+    [indexSet addIndex:modelToViewMap_[modelSelection[i]]];
+  [tableView_ selectRowIndexes:indexSet byExtendingSelection:NO];
+
   [tableView_ reloadData];
   [self adjustSelectionAndEndProcessButton];
 }
 
 - (IBAction)statsLinkClicked:(id)sender {
   TaskManager::GetInstance()->OpenAboutMemory();
 }
 
 - (IBAction)killSelectedProcesses:(id)sender {
   NSIndexSet* selection = [tableView_ selectedRowIndexes];
   for (NSUInteger i = [selection lastIndex];
        i != NSNotFound;
        i = [selection indexLessThanIndex:i]) {
-    taskManager_->KillProcess(i);
+    taskManager_->KillProcess(viewToModelMap_[i]);
   }
 }
 
 - (void)selectDoubleClickedTab:(id)sender {
   NSInteger row = [tableView_ clickedRow];
   if (row < 0)
     return;  // Happens e.g. if the table header is double-clicked.
   taskManager_->ActivateProcess(row);
 }
 
@@ skipping 136 matching lines @@
 - (void)adjustSelectionAndEndProcessButton {
   bool selectionContainsBrowserProcess = false;
 
   // If a row is selected, make sure that all rows belonging to the same process
   // are selected as well. Also, check if the selection contains the browser
   // process.
   NSIndexSet* selection = [tableView_ selectedRowIndexes];
   for (NSUInteger i = [selection lastIndex];
        i != NSNotFound;
        i = [selection indexLessThanIndex:i]) {
-    if (taskManager_->IsBrowserProcess(i))
+    int modelIndex = viewToModelMap_[i];
+    if (taskManager_->IsBrowserProcess(modelIndex))
       selectionContainsBrowserProcess = true;
 
-    std::pair<int, int> rangePair = model_->GetGroupRangeForResource(i);
-    NSRange range = NSMakeRange(rangePair.first, rangePair.second);
-    NSIndexSet* rangeIndexSet = [NSIndexSet indexSetWithIndexesInRange:range];
-    [tableView_ selectRowIndexes:rangeIndexSet byExtendingSelection:YES];
+    std::pair<int, int> rangePair =
+        model_->GetGroupRangeForResource(modelIndex);
+    NSMutableIndexSet* indexSet = [NSMutableIndexSet indexSet];
+    for (int j = 0; j < rangePair.second; ++j)
+      [indexSet addIndex:modelToViewMap_[rangePair.first + j]];
+    [tableView_ selectRowIndexes:indexSet byExtendingSelection:YES];
   }
 
   bool enabled = [selection count] > 0 && !selectionContainsBrowserProcess;
   [endProcessButton_ setEnabled:enabled];
 }
 
+- (void)deselectRows {
+  [tableView_ deselectAll:self];
+}
+
 // Table view delegate method.
 - (void)tableViewSelectionIsChanging:(NSNotification*)aNotification {
   [self adjustSelectionAndEndProcessButton];
 }
 
 - (void)windowWillClose:(NSNotification*)notification {
   if (taskManagerObserver_) {
     taskManagerObserver_->WindowWasClosed();
     taskManagerObserver_ = nil;
   }
   [self autorelease];
 }
 
 @end
 
 @implementation TaskManagerWindowController (NSTableDataSource)
 
 - (NSInteger)numberOfRowsInTableView:(NSTableView*)tableView {
   DCHECK(tableView == tableView_ || tableView_ == nil);
   return model_->ResourceCount();
 }
 
 - (NSString*)modelTextForRow:(int)row column:(int)columnId {
-  DCHECK_LT(static_cast<size_t>(row), indexShuffle_.size());
-  row = indexShuffle_[row];
+  DCHECK_LT(static_cast<size_t>(row), viewToModelMap_.size());
+  row = viewToModelMap_[row];
   switch (columnId) {
     case IDS_TASK_MANAGER_PAGE_COLUMN:  // Process
       return base::SysWideToNSString(model_->GetResourceTitle(row));
 
     case IDS_TASK_MANAGER_PRIVATE_MEM_COLUMN:  // Memory
       if (!model_->IsResourceFirstInGroup(row))
         return @"";
       return base::SysWideToNSString(model_->GetResourcePrivateMemory(row));
 
     case IDS_TASK_MANAGER_SHARED_MEM_COLUMN:  // Memory
@@ skipping 65 matching lines @@
   NSCell* cell = [tableColumn dataCellForRow:rowIndex];
 
   // Set the favicon and title for the task in the name column.
   if ([[tableColumn identifier] intValue] == IDS_TASK_MANAGER_PAGE_COLUMN) {
     DCHECK([cell isKindOfClass:[NSButtonCell class]]);
     NSButtonCell* buttonCell = static_cast<NSButtonCell*>(cell);
     NSString* title = [self modelTextForRow:rowIndex
                                     column:[[tableColumn identifier] intValue]];
     [buttonCell setTitle:title];
     [buttonCell setImage:
-        taskManagerObserver_->GetImageForRow(indexShuffle_[rowIndex])];
+        taskManagerObserver_->GetImageForRow(viewToModelMap_[rowIndex])];
     [buttonCell setRefusesFirstResponder:YES];  // Don't push in like a button.
     [buttonCell setHighlightsBy:NSNoCellMask];
   }
 
   return cell;
 }
 
 - (void)           tableView:(NSTableView*)tableView
     sortDescriptorsDidChange:(NSArray*)oldDescriptors {
   NSArray* newDescriptors = [tableView sortDescriptors];
@@ skipping 28 matching lines @@
     task_manager_->OnWindowClosed();
   }
   model_->RemoveObserver(this);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // TaskManagerMac, TaskManagerModelObserver implementation:
 
 void TaskManagerMac::OnModelChanged() {
   icon_cache_.OnModelChanged();
+  [window_controller_ deselectRows];
   [window_controller_ reloadData];
 }
 
 void TaskManagerMac::OnItemsChanged(int start, int length) {
   icon_cache_.OnItemsChanged(start, length);
   [window_controller_ reloadData];
 }
 
 void TaskManagerMac::OnItemsAdded(int start, int length) {
   icon_cache_.OnItemsAdded(start, length);
+  [window_controller_ deselectRows];
   [window_controller_ reloadData];
 }
 
 void TaskManagerMac::OnItemsRemoved(int start, int length) {
   icon_cache_.OnItemsRemoved(start, length);
+  [window_controller_ deselectRows];
   [window_controller_ reloadData];
 }
 
 NSImage* TaskManagerMac::GetImageForRow(int row) {
   return icon_cache_.GetImageForRow(row);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // TaskManagerMac, public:
 
 void TaskManagerMac::WindowWasClosed() {
   delete this;
   instance_ = NULL;
 }
 
 // static
 void TaskManagerMac::Show() {
   if (instance_) {
     // If there's a Task manager window open already, just activate it.
     [[instance_->window_controller_ window]
         makeKeyAndOrderFront:instance_->window_controller_];
   } else {
     instance_ = new TaskManagerMac(TaskManager::GetInstance());
     instance_->model_->StartUpdating();
   }
 }