Replace PeekMessage for TSF awareness

base/message_pump_win.cc

Index: base/message_pump_win.cc
diff --git a/base/message_pump_win.cc b/base/message_pump_win.cc
index 2b2a10e07534afcf7c7dd8de14567bd4c794d5cd..18c4394ff66c125908984ce6d09dd63d85ea2f3d 100644
--- a/base/message_pump_win.cc
+++ b/base/message_pump_win.cc
@@ -4,14 +4,89 @@
 
 #include "base/message_pump_win.h"
 
+#include <imm.h>
 #include <math.h>
+#include <msctf.h>
+#include <Windows.h>
 
 #include "base/debug/trace_event.h"
 #include "base/message_loop.h"
 #include "base/metrics/histogram.h"
 #include "base/process_util.h"
+#include "base/win/metro.h"
+#include "base/win/scoped_comptr.h"
 #include "base/win/wrapped_window_proc.h"
 
+namespace base {
+
+static const wchar_t kWndClass[] = L"Chrome_MessagePumpWindow";
+
+// Message sent to get an additional time slice for pumping (processing) another
+// task (a series of such messages creates a continuous task pump).
+static const int kMsgHaveWork = WM_USER + 1;
+
+class MessagePumpTSFInternal {
+ public:
+  TfClientId client_id;
+  bool is_initialized;
+  base::win::ScopedComPtr<ITfThreadMgr> thread_mgr;
+  base::win::ScopedComPtr<ITfMessagePump> message_pump;
+  base::win::ScopedComPtr<ITfKeystrokeMgr> keystroke_mgr;
+  MessagePumpTSFInternal()
+      : client_id(TF_CLIENTID_NULL), is_initialized(true) {
+    is_initialized &= SUCCEEDED(thread_mgr.CreateInstance(CLSID_TF_ThreadMgr));
+    if (!thread_mgr)
+      return;
+
+    is_initialized &= SUCCEEDED(message_pump.QueryFrom(thread_mgr));
+    is_initialized &= SUCCEEDED(keystroke_mgr.QueryFrom(thread_mgr));
+    // When activate succeeded, |client_id| is set valid value.
+    is_initialized &= SUCCEEDED(thread_mgr->Activate(&client_id));
+  }
+  ~MessagePumpTSFInternal() {
+    if (thread_mgr && client_id != TF_CLIENTID_NULL)
+      thread_mgr->Deactivate();
+  }
+
+  bool ForwardKeyMessageForTSF(const MSG& msg) {
+    if (!is_initialized)
+      return false;
+
+    if(msg.message == WM_KEYDOWN) {
+      BOOL eaten = FALSE;
+      HRESULT hr = keystroke_mgr->TestKeyDown(msg.wParam, msg.lParam, &eaten);
+      if (FAILED(hr) && !eaten)
+        return false;
+      eaten = FALSE;
+      hr = keystroke_mgr->KeyDown(msg.wParam, msg.lParam, &eaten);
+      return (SUCCEEDED(hr) && eaten);
+    }
+
+    if(msg.message == WM_KEYUP) {
+      BOOL eaten = FALSE;
+      HRESULT hr = keystroke_mgr->TestKeyUp(msg.wParam, msg.lParam, &eaten);
+      if (FAILED(hr) && !eaten)
+        return false;
+      eaten = FALSE;
+      hr = keystroke_mgr->KeyUp(msg.wParam, msg.lParam, &eaten);
+      return (SUCCEEDED(hr) && eaten);
+    }
+
+    return false;
+  }
+
+  bool PeekMessageForTSF(MSG* msg, HWND hwnd, UINT wmin, UINT wmax, UINT wmsg) {
+    if (!is_initialized)
+      // FallBack.
+      return !!::PeekMessage(msg, hwnd, wmin, wmax, wmsg);
+
+    BOOL result = FALSE;
+    if(FAILED(message_pump->PeekMessage(msg, hwnd, wmin, wmax, wmsg, &result)))
+      result = FALSE;
+    return !!result;
+  }
+};
+
 namespace {
 
 enum MessageLoopProblems {
@@ -23,14 +98,6 @@ enum MessageLoopProblems {
 
 }  // namespace
 
-namespace base {
-
-static const wchar_t kWndClass[] = L"Chrome_MessagePumpWindow";
-
-// Message sent to get an additional time slice for pumping (processing) another
-// task (a series of such messages creates a continuous task pump).
-static const int kMsgHaveWork = WM_USER + 1;
-
 //-----------------------------------------------------------------------------
 // MessagePumpWin public:
 
@@ -97,9 +164,14 @@ int MessagePumpWin::GetCurrentDelay() const {
 
 MessagePumpForUI::MessagePumpForUI() : instance_(NULL) {
   InitMessageWnd();
+
+  if (base::win::IsTsfAwareRequired())
+    tsf_message_pump_.reset(new MessagePumpTSFInternal);
 }
 
 MessagePumpForUI::~MessagePumpForUI() {
+  if (base::win::IsTsfAwareRequired())
+    tsf_message_pump_.reset(NULL);
   DestroyWindow(message_hwnd_);
   UnregisterClass(kWndClass, instance_);
 }
@@ -182,7 +254,7 @@ void MessagePumpForUI::PumpOutPendingPaintMessages() {
   int peek_count;
   for (peek_count = 0; peek_count < kMaxPeekCount; ++peek_count) {
     MSG msg;
-    if (!PeekMessage(&msg, NULL, 0, 0, PM_REMOVE | PM_QS_PAINT))
+    if (!PeekMessageInternal(&msg, NULL, 0, 0, PM_REMOVE | PM_QS_PAINT))
       break;
     ProcessMessageHelper(msg);
     if (state_->should_quit)  // Handle WM_QUIT.
@@ -304,7 +376,8 @@ void MessagePumpForUI::WaitForWork() {
     MSG msg = {0};
     DWORD queue_status = GetQueueStatus(QS_MOUSE);
     if (HIWORD(queue_status) & QS_MOUSE &&
-       !PeekMessage(&msg, NULL, WM_MOUSEFIRST, WM_MOUSELAST, PM_NOREMOVE)) {
+       !PeekMessageInternal(&msg, NULL, WM_MOUSEFIRST, WM_MOUSELAST,
+                            PM_NOREMOVE)) {
       WaitMessage();
     }
     return;
@@ -361,7 +434,7 @@ bool MessagePumpForUI::ProcessNextWindowsMessage() {
     sent_messages_in_queue = true;
 
   MSG msg;
-  if (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
+  if (PeekMessageInternal(&msg, NULL, 0, 0, PM_REMOVE))
     return ProcessMessageHelper(msg);
 
   return sent_messages_in_queue;
@@ -387,12 +460,19 @@ bool MessagePumpForUI::ProcessMessageHelper(const MSG& msg) {
 
   WillProcessMessage(msg);
 
-  if (state_->dispatcher) {
-    if (!state_->dispatcher->Dispatch(msg))
-      state_->should_quit = true;
-  } else {
-    TranslateMessage(&msg);
-    DispatchMessage(&msg);
+  bool is_key_dispatched_by_tsf = false;
+  if (base::win::IsTsfAwareRequired()) {
+    is_key_dispatched_by_tsf = tsf_message_pump_->ForwardKeyMessageForTSF(msg);
+  }
+
+  if (!is_key_dispatched_by_tsf) {
+    if (state_->dispatcher) {
+      if (!state_->dispatcher->Dispatch(msg))
+        state_->should_quit = true;
+    } else {
+      TranslateMessage(&msg);
+      DispatchMessage(&msg);
+    }
   }
 
   DidProcessMessage(msg);
@@ -416,10 +496,11 @@ bool MessagePumpForUI::ProcessPumpReplacementMessage() {
   // This is to ensure that these messages are peeked out by the OS modal loop.
   if (MessageLoop::current()->os_modal_loop()) {
     // We only peek out WM_PAINT and WM_TIMER here for reasons mentioned above.
-    have_message = PeekMessage(&msg, NULL, WM_PAINT, WM_PAINT, PM_REMOVE) ||
-                   PeekMessage(&msg, NULL, WM_TIMER, WM_TIMER, PM_REMOVE);
+    have_message =
+      PeekMessageInternal(&msg, NULL, WM_PAINT, WM_PAINT, PM_REMOVE) ||
+      PeekMessageInternal(&msg, NULL, WM_TIMER, WM_TIMER, PM_REMOVE);
   } else {
-    have_message = (0 != PeekMessage(&msg, NULL, 0, 0, PM_REMOVE));
+    have_message = PeekMessageInternal(&msg, NULL, 0, 0, PM_REMOVE);
   }
 
   DCHECK(!have_message || kMsgHaveWork != msg.message ||
@@ -441,6 +522,14 @@ bool MessagePumpForUI::ProcessPumpReplacementMessage() {
   return ProcessMessageHelper(msg);
 }
 
+bool MessagePumpForUI::PeekMessageInternal(
+    MSG* msg, HWND hwnd, UINT wmin, UINT wmax, UINT wmsg) {
+  if (base::win::IsTsfAwareRequired())
+    return tsf_message_pump_->PeekMessageForTSF(msg, hwnd, wmin, wmax, wmsg);
+
+  return !!::PeekMessage(msg, hwnd, wmin, wmax, wmsg);
+}
+
 //-----------------------------------------------------------------------------
 // MessagePumpForIO public: