[PATCH 2/3] user32: Add internal reference counting to menu structs

[Andrew Eikum]

This is needed so the client can release the HMENU handle, but we can
still reference the menu structure internally. The HMENU is still used
all over the place internally, and those have to be fixed to point to
the raw struct so things don't fail if the HMENU has been released.

---
 dlls/user32/menu.c       | 48 +++++++++++++++++++++++++++++++++++++++---------
 dlls/user32/tests/menu.c |  6 ++++--
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/dlls/user32/menu.c b/dlls/user32/menu.c
index 0aa7744..6f064de 100644
--- a/dlls/user32/menu.c
+++ b/dlls/user32/menu.c
@@ -97,6 +97,7 @@ typedef struct {
     BOOL        bScrolling;   /* Scroll arrows are active */
     UINT        nScrollPos;   /* Current scroll position */
     UINT        nTotalHeight; /* Total height of menu items inside menu */
+    LONG        refcount;
     /* ------------ MENUINFO members ------ */
     DWORD	dwStyle;	/* Extended menu style */
     UINT	cyMax;		/* max height of the whole menu, 0 is screen height */
@@ -179,6 +180,7 @@ static BOOL fEndMenu = FALSE;
 DWORD WINAPI DrawMenuBarTemp(HWND hwnd, HDC hDC, LPRECT lprect, HMENU hMenu, HFONT hFont);
 
 static BOOL SetMenuItemInfo_common( MENUITEM *, const MENUITEMINFOW *, BOOL);
+static BOOL MENU_ReleaseMenu( POPUPMENU *lppop );
 
 /*********************************************************************
  * menu class descriptor
@@ -3943,6 +3945,8 @@ BOOL WINAPI RemoveMenu( HMENU hMenu, UINT nPos, UINT wFlags )
     if (!(menu = MENU_GetMenu(hMenu))) return FALSE;
 
       /* Remove item */
+    if ((item->fType & MF_POPUP) && item->hSubMenu)
+        MENU_ReleaseMenu(MENU_GetMenu(item->hSubMenu));
 
     MENU_FreeItemData( item );
 
@@ -4086,6 +4090,7 @@ HMENU WINAPI CreateMenu(void)
     if (!(menu = HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(*menu) ))) return 0;
     menu->FocusedItem = NO_SELECTED_ITEM;
     menu->bTimeToHide = FALSE;
+    menu->refcount = 1;
 
     if (!(hMenu = alloc_user_handle( &menu->obj, USER_MENU ))) HeapFree( GetProcessHeap(), 0, menu );
 
@@ -4095,17 +4100,17 @@ HMENU WINAPI CreateMenu(void)
 }
 
 
-/**********************************************************************
- *         DestroyMenu    (USER32.@)
- */
-BOOL WINAPI DestroyMenu( HMENU hMenu )
+/* return value indicates whether the menu was actually destroyed */
+static BOOL MENU_ReleaseMenu( POPUPMENU *lppop )
 {
-    LPPOPUPMENU lppop;
+    LONG ref;
 
-    TRACE("(%p)\n", hMenu);
+    if (!lppop)
+        return FALSE;
 
-    if (!(lppop = free_user_handle( hMenu, USER_MENU ))) return FALSE;
-    if (lppop == OBJ_OTHER_PROCESS) return FALSE;
+    ref = InterlockedDecrement(&lppop->refcount);
+    if (ref)
+        return FALSE;
 
     /* DestroyMenu should not destroy system menu popup owner */
     if ((lppop->wFlags & (MF_POPUP | MF_SYSMENU)) == MF_POPUP && lppop->hWnd)
@@ -4120,12 +4125,34 @@ BOOL WINAPI DestroyMenu( HMENU hMenu )
         MENUITEM *item = lppop->items;
         for (i = lppop->nItems; i > 0; i--, item++)
         {
-            if (item->fType & MF_POPUP) DestroyMenu(item->hSubMenu);
+            if (item->fType & MF_POPUP){
+                if (!MENU_ReleaseMenu(MENU_GetMenu(item->hSubMenu)))
+                    /* release HMENU if it hasn't already been released */
+                    DestroyMenu(item->hSubMenu);
+            }
             MENU_FreeItemData( item );
         }
         HeapFree( GetProcessHeap(), 0, lppop->items );
     }
     HeapFree( GetProcessHeap(), 0, lppop );
+
+    return TRUE;
+}
+
+/**********************************************************************
+ *         DestroyMenu    (USER32.@)
+ */
+BOOL WINAPI DestroyMenu( HMENU hMenu )
+{
+    POPUPMENU *lppop;
+
+    TRACE("(%p)\n", hMenu);
+
+    if (!(lppop = free_user_handle( hMenu, USER_MENU ))) return FALSE;
+    if (lppop == OBJ_OTHER_PROCESS) return FALSE;
+
+    MENU_ReleaseMenu(lppop);
+
     return TRUE;
 }
 
@@ -4775,6 +4802,8 @@ static BOOL SetMenuItemInfo_common(MENUITEM * menu,
 	menu->wID = lpmii->wID;
 
     if (lpmii->fMask & MIIM_SUBMENU) {
+	if ((menu->fType & MF_POPUP) && menu->hSubMenu)
+	    MENU_ReleaseMenu(MENU_GetMenu(menu->hSubMenu));
 	menu->hSubMenu = lpmii->hSubMenu;
 	if (menu->hSubMenu) {
 	    POPUPMENU *subMenu = MENU_GetMenu(menu->hSubMenu);
@@ -4786,6 +4815,7 @@ static BOOL SetMenuItemInfo_common(MENUITEM * menu,
                 }
 		subMenu->wFlags |= MF_POPUP;
 		menu->fType |= MF_POPUP;
+		InterlockedIncrement(&subMenu->refcount);
 	    } else {
                 SetLastError( ERROR_INVALID_PARAMETER);
                 return FALSE;
diff --git a/dlls/user32/tests/menu.c b/dlls/user32/tests/menu.c
index 593292d..f4dc2ec 100644
--- a/dlls/user32/tests/menu.c
+++ b/dlls/user32/tests/menu.c
@@ -1854,8 +1854,10 @@ static void test_menu_search_bycommand( void )
     ok (info.wID == (UINT_PTR)hmenuSub, "IDs differ for the popup menu\n");
     ok (!strcmp(info.dwTypeData, "Case 2 MenuItem 2"), "Returned item has wrong label (%s)\n", info.dwTypeData);
 
-    DestroyMenu( hmenu );
-    DestroyMenu( hmenuSub );
+    rc = DestroyMenu( hmenu );
+    ok(rc == TRUE, "main menu destroy should have succeeded\n");
+    rc = DestroyMenu( hmenuSub );
+    ok(rc == FALSE, "submenu should have already been destroyed\n");
 
     /* 
         Case 3: Menu containing a popup menu which in turn 
-- 
1.9.2