[Windows] Tighten logic in MessageWindow::WindowProc

Replace switch statement with tighter logic during creation and
destruction. No functional changes.

Bug: none
Change-Id: Id2dc3c3a183ae62c3543cec65c42c913af332570
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4864940
Auto-Submit: Greg Thompson <[email protected]>
Reviewed-by: Will Harris <[email protected]>
Commit-Queue: Will Harris <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1206176}

base/win/message_window.cc

@@ -11,6 +11,7 @@
 
 #include "base/check.h"
 #include "base/check_op.h"
+#include "base/compiler_specific.h"
 #include "base/lazy_instance.h"
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
@@ -192,37 +193,36 @@ LRESULT CALLBACK MessageWindow::WindowProc(HWND hwnd,
   auto& message_window_map = MessageWindowMap::GetInstanceForCurrentThread();
   MessageWindow* self = message_window_map.Get(hwnd);
 
-  switch (message) {
-    // Set up the self before handling WM_CREATE.
-    case WM_CREATE: {
-      CREATESTRUCT* cs = reinterpret_cast<CREATESTRUCT*>(lparam);
-      self = reinterpret_cast<MessageWindow*>(cs->lpCreateParams);
+  // CreateWindow will send a WM_CREATE message during window creation.
+  if (UNLIKELY(!self && message == WM_CREATE)) {
+    CREATESTRUCT* const cs = reinterpret_cast<CREATESTRUCT*>(lparam);
+    self = reinterpret_cast<MessageWindow*>(cs->lpCreateParams);
 
-      // Make |hwnd| available to the message handler. At this point the control
-      // hasn't returned from CreateWindow() yet.
-      self->window_ = hwnd;
+    // Tell the MessageWindow instance the HWND that CreateWindow has produced.
+    self->window_ = hwnd;
 
-      // Store pointer to self to local map.
-      message_window_map.Insert(hwnd, *self);
-      break;
-    }
+    // Associate the MessageWindow instance with the HWND in the map.
+    message_window_map.Insert(hwnd, *self);
+  }
 
-    // Clear the map key to stop calling the self once WM_DESTROY is
-    // received.
-    case WM_DESTROY: {
-      message_window_map.Erase(hwnd);
-      break;
-    }
+  if (UNLIKELY(!self)) {
+    return DefWindowProc(hwnd, message, wparam, lparam);
   }
 
-  // Handle the message.
-  if (self) {
-    LRESULT message_result;
-    if (self->message_callback_.Run(message, wparam, lparam, &message_result))
-      return message_result;
+  LRESULT message_result = {};
+  if (!self->message_callback_.Run(message, wparam, lparam, &message_result)) {
+    message_result = DefWindowProc(hwnd, message, wparam, lparam);
+  }
+
+  if (UNLIKELY(message == WM_DESTROY)) {
+    // Tell the MessageWindow instance that it no longer has an HWND.
+    self->window_ = nullptr;
+
+    // Remove this HWND's MessageWindow from the map since it is going away.
+    message_window_map.Erase(hwnd);
   }
 
-  return DefWindowProc(hwnd, message, wparam, lparam);
+  return message_result;
 }
 
 }  // namespace win