7 files changed, 198 insertions, 1 deletions
diff --git a/chrome/browser/cocoa/find_bar_bridge.mm b/chrome/browser/cocoa/find_bar_bridge.mm
index 4b4a4cd..ee760e5 100644
--- a/chrome/browser/cocoa/find_bar_bridge.mm
+++ b/chrome/browser/cocoa/find_bar_bridge.mm
@@ -67,5 +67,5 @@ void FindBarBridge::SetDialogPosition(const gfx::Rect& new_pos,
 }
 
 void FindBarBridge::RestoreSavedFocus() {
-  // http://crbug.com/12657
+  [cocoa_controller_ restoreSavedFocus];
 }
diff --git a/chrome/browser/cocoa/find_bar_cocoa_controller.h b/chrome/browser/cocoa/find_bar_cocoa_controller.h
index 13837de..221f859 100644
--- a/chrome/browser/cocoa/find_bar_cocoa_controller.h
+++ b/chrome/browser/cocoa/find_bar_cocoa_controller.h
@@ -6,11 +6,13 @@
 
 #import "chrome/browser/cocoa/find_bar_cocoa_controller.h"
 
+#include "base/scoped_nsobject.h"
 #include "base/string16.h"
 
 class BrowserWindowCocoa;
 class FindBarBridge;
 class FindNotificationDetails;
+@class FocusTracker;
 
 // A controller for the find bar in the browser window.  Manages
 // updating the state of the find bar and provides a target for the
@@ -27,6 +29,8 @@ class FindNotificationDetails;
 
   // Needed to call methods on FindBarController.
   FindBarBridge* findBarBridge_;  // weak
+
+  scoped_nsobject<FocusTracker> focusTracker_;
 };
 
 // Initializes a new FindBarCocoaController.
@@ -48,6 +52,7 @@ class FindNotificationDetails;
 - (void)showFindBar;
 - (void)hideFindBar;
 - (void)setFocusAndSelection;
+- (void)restoreSavedFocus;
 - (void)setFindText:(const string16&)findText;
 
 - (void)clearResults:(const FindNotificationDetails&)results;
diff --git a/chrome/browser/cocoa/find_bar_cocoa_controller.mm b/chrome/browser/cocoa/find_bar_cocoa_controller.mm
index b0f42c1..4b6e89e 100644
--- a/chrome/browser/cocoa/find_bar_cocoa_controller.mm
+++ b/chrome/browser/cocoa/find_bar_cocoa_controller.mm
@@ -12,6 +12,7 @@
 #include "chrome/browser/cocoa/browser_window_cocoa.h"
 #import "chrome/browser/cocoa/find_bar_cocoa_controller.h"
 #import "chrome/browser/cocoa/find_bar_bridge.h"
+#import "chrome/browser/cocoa/focus_tracker.h"
 #import "chrome/browser/cocoa/tab_strip_controller.h"
 #include "chrome/browser/tab_contents/tab_contents.h"
 
@@ -120,6 +121,14 @@
 // Methods from FindBar
 - (void)showFindBar {
   [[self view] setHidden:NO];
+
+  // Save the currently-focused view.  |[self view]| is in the view
+  // hierarchy by now.  showFindBar can be called even when the
+  // findbar is already open, so do not overwrite an already saved
+  // view.
+  if (!focusTracker_.get())
+    focusTracker_.reset(
+        [[FocusTracker alloc] initWithWindow:[[self view] window]]);
 }
 
 - (void)hideFindBar {
@@ -136,6 +145,15 @@
 
 }
 
+- (void)restoreSavedFocus {
+  if (!(focusTracker_.get() &&
+        [focusTracker_ restoreFocusInWindow:[[self view] window]])) {
+    // Fall back to giving focus to the tab contents.
+    findBarBridge_->GetFindBarController()->tab_contents()->Focus();
+  }
+  focusTracker_.reset(nil);
+}
+
 - (void)setFindText:(const string16&)findText {
   [findText_ setStringValue:base::SysUTF16ToNSString(findText)];
 }
@@ -177,6 +195,11 @@
     [resultsLabel_ setStringValue:@""];
   }
 
+  // If we found any results, reset the focus tracker, so we always
+  // restore focus to the tab contents.
+  if (result.number_of_matches() > 0)
+    focusTracker_.reset(nil);
+
   // Resize |resultsLabel_| to completely contain its string and right-justify
   // it within |findText_|.  sizeToFit may shrink the frame vertically, which we
   // don't want, so we save the original vertical positioning.
diff --git a/chrome/browser/cocoa/focus_tracker.h b/chrome/browser/cocoa/focus_tracker.h
new file mode 100644
index 0000000..f828979
--- /dev/null
+++ b/chrome/browser/cocoa/focus_tracker.h
@@ -0,0 +1,28 @@
+// Copyright (c) 2009 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.
+
+#import <Cocoa/Cocoa.h>
+
+#include "base/scoped_nsobject.h"
+
+// A class that handles saving and restoring focus.  An instance of
+// this class snapshots the currently focused view when it is
+// constructed, and callers can use restoreFocus to return focus to
+// that view.  FocusTracker will not restore focus to views that are
+// no longer in the view hierarchy or are not in the correct window.
+
+@interface FocusTracker : NSObject {
+ @private
+  scoped_nsobject<NSView> focusedView_;
+}
+
+// |window| is the window that we are saving focus for.  This
+// method snapshots the currently focused view.
+- (id)initWithWindow:(NSWindow*)window;
+
+// Attempts to restore focus to the snapshotted view.  Returns YES if
+// focus was restored.  Will not restore focus if the view is no
+// longer in the view hierarchy under |window|.
+- (BOOL)restoreFocusInWindow:(NSWindow*)window;
+@end
diff --git a/chrome/browser/cocoa/focus_tracker.mm b/chrome/browser/cocoa/focus_tracker.mm
new file mode 100644
index 0000000..ecbb864
--- /dev/null
+++ b/chrome/browser/cocoa/focus_tracker.mm
@@ -0,0 +1,46 @@
+// Copyright (c) 2009 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.
+
+#import "chrome/browser/cocoa/focus_tracker.h"
+
+#include "base/logging.h"
+
+@implementation FocusTracker
+
+- (id)initWithWindow:(NSWindow*)window {
+  if ((self = [super init])) {
+    NSResponder* current_focus = [window firstResponder];
+
+    // Special case NSTextViews, because they are removed from the
+    // view hierarchy when their text field does not have focus.  If
+    // an NSTextView is the current first responder, save a pointer to
+    // its NSTextField delegate instead.
+    if ([current_focus isKindOfClass:[NSTextView class]]) {
+      id delegate = [(NSTextView*)current_focus delegate];
+      if ([delegate isKindOfClass:[NSTextField class]])
+        current_focus = delegate;
+      else
+        current_focus = nil;
+    }
+
+    if ([current_focus isKindOfClass:[NSView class]]) {
+      NSView* current_focus_view = (NSView*)current_focus;
+      focusedView_.reset([current_focus_view retain]);
+    }
+  }
+
+  return self;
+}
+
+- (BOOL)restoreFocusInWindow:(NSWindow*)window {
+  if (!focusedView_.get())
+    return NO;
+
+  if ([focusedView_ window] && [focusedView_ window] == window)
+    return [window makeFirstResponder:focusedView_.get()];
+
+  return NO;
+}
+
+@end
diff --git a/chrome/browser/cocoa/focus_tracker_unittest.mm b/chrome/browser/cocoa/focus_tracker_unittest.mm
new file mode 100644
index 0000000..65ab91e
--- /dev/null
+++ b/chrome/browser/cocoa/focus_tracker_unittest.mm
@@ -0,0 +1,92 @@
+// Copyright (c) 2009 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.
+
+#import <Cocoa/Cocoa.h>
+
+#include "base/scoped_nsobject.h"
+#import "chrome/browser/cocoa/cocoa_test_helper.h"
+#import "chrome/browser/cocoa/focus_tracker.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "testing/platform_test.h"
+
+namespace {
+
+class FocusTrackerTest : public PlatformTest {
+ public:
+  virtual void SetUp() {
+    PlatformTest::SetUp();
+
+    viewA_.reset([[NSView alloc] initWithFrame:NSZeroRect]);
+    viewB_.reset([[NSView alloc] initWithFrame:NSZeroRect]);
+    [helper_.contentView() addSubview:viewA_.get()];
+    [helper_.contentView() addSubview:viewB_.get()];
+  }
+
+ protected:
+  CocoaTestHelper helper_;
+  scoped_nsobject<NSView> viewA_;
+  scoped_nsobject<NSView> viewB_;
+};
+
+TEST_F(FocusTrackerTest, SaveRestore) {
+  NSWindow* window = helper_.window();
+  ASSERT_TRUE([window makeFirstResponder:viewA_.get()]);
+  FocusTracker* tracker =
+    [[[FocusTracker alloc] initWithWindow:window] autorelease];
+
+  // Give focus to |viewB_|, then try and restore it to view1.
+  ASSERT_TRUE([window makeFirstResponder:viewB_.get()]);
+  EXPECT_TRUE([tracker restoreFocusInWindow:window]);
+  EXPECT_EQ(viewA_.get(), [window firstResponder]);
+}
+
+TEST_F(FocusTrackerTest, SaveRestoreWithTextView) {
+  // Valgrind will complain if the text field has zero size.
+  NSRect frame = NSMakeRect(0, 0, 100, 20);
+  NSWindow* window = helper_.window();
+  NSTextField* text =
+    [[[NSTextField alloc] initWithFrame:frame] autorelease];
+  [helper_.contentView() addSubview:text];
+
+  ASSERT_TRUE([window makeFirstResponder:text]);
+  FocusTracker* tracker =
+    [[[FocusTracker alloc] initWithWindow:window] autorelease];
+
+  // Give focus to |viewB_|, then try and restore it to the text field.
+  ASSERT_TRUE([window makeFirstResponder:viewB_.get()]);
+  EXPECT_TRUE([tracker restoreFocusInWindow:window]);
+  EXPECT_TRUE([[window firstResponder] isKindOfClass:[NSTextView class]]);
+}
+
+TEST_F(FocusTrackerTest, DontRestoreToViewNotInWindow) {
+  NSWindow* window = helper_.window();
+  NSView* view3 = [[[NSView alloc] initWithFrame:NSZeroRect] autorelease];
+  [helper_.contentView() addSubview:view3];
+
+  ASSERT_TRUE([window makeFirstResponder:view3]);
+  FocusTracker* tracker =
+    [[[FocusTracker alloc] initWithWindow:window] autorelease];
+
+  // Give focus to |viewB_|, then remove view3 from the hierarchy and try
+  // to restore focus.  The restore should fail.
+  ASSERT_TRUE([window makeFirstResponder:viewB_.get()]);
+  [view3 removeFromSuperview];
+  EXPECT_FALSE([tracker restoreFocusInWindow:window]);
+}
+
+TEST_F(FocusTrackerTest, DontRestoreFocusToViewInDifferentWindow) {
+  NSWindow* window = helper_.window();
+  ASSERT_TRUE([window makeFirstResponder:viewA_.get()]);
+  FocusTracker* tracker =
+    [[[FocusTracker alloc] initWithWindow:window] autorelease];
+
+  // Give focus to |viewB_|, then try and restore focus in a different
+  // window.  It is ok to pass a nil NSWindow here because we only use
+  // it for direct comparison.
+  ASSERT_TRUE([window makeFirstResponder:viewB_.get()]);
+  EXPECT_FALSE([tracker restoreFocusInWindow:nil]);
+}
+
+
+}  // namespace
diff --git a/chrome/chrome.gyp b/chrome/chrome.gyp
index eb799d7..bdf8863 100644
--- a/chrome/chrome.gyp
+++ b/chrome/chrome.gyp
@@ -959,6 +959,8 @@
         'browser/cocoa/find_bar_view.mm',
         'browser/cocoa/first_run_dialog.h',
         'browser/cocoa/first_run_dialog.mm',
+        'browser/cocoa/focus_tracker.h',
+        'browser/cocoa/focus_tracker.mm',
         'browser/cocoa/fullscreen_window.h',
         'browser/cocoa/fullscreen_window.mm',
         'browser/cocoa/gradient_button_cell.h',
@@ -4031,6 +4033,7 @@
         'browser/cocoa/find_bar_bridge_unittest.mm',
         'browser/cocoa/find_bar_cocoa_controller_unittest.mm',
         'browser/cocoa/find_bar_view_unittest.mm',
+        'browser/cocoa/focus_tracker_unittest.mm',
         'browser/cocoa/fullscreen_window_unittest.mm',
         'browser/cocoa/hung_renderer_controller_unittest.mm',
         'browser/cocoa/infobar_container_controller_unittest.mm',