layers: Fix several bugs in WSI extension handling in layers

Layers weren't intercepting the WSI device extensions proerly, or weren't
NOT intercepting them if the extension wasn't enabled.  Fixes to APIDump,
Generic, ScreenShot, MemTracker, ObjectTracker
diff --git a/layers/mem_tracker.cpp b/layers/mem_tracker.cpp
index 7da466b..f60754b 100644
--- a/layers/mem_tracker.cpp
+++ b/layers/mem_tracker.cpp
@@ -1366,6 +1366,16 @@
 static void createDeviceRegisterExtensions(const VkDeviceCreateInfo* pCreateInfo, VkDevice device)
 {
     layer_data *my_device_data = get_my_data_ptr(get_dispatch_key(device), layer_data_map);
+    VkLayerDispatchTable *pDisp = get_dispatch_table(mem_tracker_device_table_map, device);
+    PFN_vkGetDeviceProcAddr gpa = pDisp->GetDeviceProcAddr;
+    pDisp->GetSurfacePropertiesKHR = (PFN_vkGetSurfacePropertiesKHR) gpa(device, "vkGetSurfacePropertiesKHR");
+    pDisp->GetSurfaceFormatsKHR = (PFN_vkGetSurfaceFormatsKHR) gpa(device, "vkGetSurfaceFormatsKHR");
+    pDisp->GetSurfacePresentModesKHR = (PFN_vkGetSurfacePresentModesKHR) gpa(device, "vkGetSurfacePresentModesKHR");
+    pDisp->CreateSwapchainKHR = (PFN_vkCreateSwapchainKHR) gpa(device, "vkCreateSwapchainKHR");
+    pDisp->DestroySwapchainKHR = (PFN_vkDestroySwapchainKHR) gpa(device, "vkDestroySwapchainKHR");
+    pDisp->GetSwapchainImagesKHR = (PFN_vkGetSwapchainImagesKHR) gpa(device, "vkGetSwapchainImagesKHR");
+    pDisp->AcquireNextImageKHR = (PFN_vkAcquireNextImageKHR) gpa(device, "vkAcquireNextImageKHR");
+    pDisp->QueuePresentKHR = (PFN_vkQueuePresentKHR) gpa(device, "vkQueuePresentKHR");
     my_device_data->wsi_enabled = false;
     for (uint32_t i = 0; i < pCreateInfo->extensionCount; i++) {
         if (strcmp(pCreateInfo->ppEnabledExtensionNames[i], VK_EXT_KHR_DEVICE_SWAPCHAIN_EXTENSION_NAME) == 0)
diff --git a/layers/object_track.h b/layers/object_track.h
index ad997d2..694aa93 100644
--- a/layers/object_track.h
+++ b/layers/object_track.h
@@ -178,6 +178,16 @@
 static void createDeviceRegisterExtensions(const VkDeviceCreateInfo* pCreateInfo, VkDevice device)
 {
     layer_data *my_device_data = get_my_data_ptr(get_dispatch_key(device), layer_data_map);
+    VkLayerDispatchTable *pDisp = get_dispatch_table(ObjectTracker_device_table_map, device);
+    PFN_vkGetDeviceProcAddr gpa = pDisp->GetDeviceProcAddr;
+    pDisp->GetSurfacePropertiesKHR = (PFN_vkGetSurfacePropertiesKHR) gpa(device, "vkGetSurfacePropertiesKHR");
+    pDisp->GetSurfaceFormatsKHR = (PFN_vkGetSurfaceFormatsKHR) gpa(device, "vkGetSurfaceFormatsKHR");
+    pDisp->GetSurfacePresentModesKHR = (PFN_vkGetSurfacePresentModesKHR) gpa(device, "vkGetSurfacePresentModesKHR");
+    pDisp->CreateSwapchainKHR = (PFN_vkCreateSwapchainKHR) gpa(device, "vkCreateSwapchainKHR");
+    pDisp->DestroySwapchainKHR = (PFN_vkDestroySwapchainKHR) gpa(device, "vkDestroySwapchainKHR");
+    pDisp->GetSwapchainImagesKHR = (PFN_vkGetSwapchainImagesKHR) gpa(device, "vkGetSwapchainImagesKHR");
+    pDisp->AcquireNextImageKHR = (PFN_vkAcquireNextImageKHR) gpa(device, "vkAcquireNextImageKHR");
+    pDisp->QueuePresentKHR = (PFN_vkQueuePresentKHR) gpa(device, "vkQueuePresentKHR");
     my_device_data->wsi_enabled = false;
     for (uint32_t i = 0; i < pCreateInfo->extensionCount; i++) {
         if (strcmp(pCreateInfo->ppEnabledExtensionNames[i], VK_EXT_KHR_DEVICE_SWAPCHAIN_EXTENSION_NAME) == 0)
@@ -640,6 +650,7 @@
         layer_data *my_device_data = get_my_data_ptr(get_dispatch_key(*pDevice), layer_data_map);
         my_device_data->report_data = layer_debug_report_create_device(my_instance_data->report_data, *pDevice);
         create_obj(*pDevice, *pDevice, VK_OBJECT_TYPE_DEVICE);
+        createDeviceRegisterExtensions(pCreateInfo, *pDevice);
     }
 
     loader_platform_thread_unlock_mutex(&objLock);
diff --git a/layers/screenshot.cpp b/layers/screenshot.cpp
index e8d4b60..bc2b820 100644
--- a/layers/screenshot.cpp
+++ b/layers/screenshot.cpp
@@ -312,6 +312,15 @@
 {
     uint32_t i;
     VkLayerDispatchTable *pDisp  = get_dispatch_table(screenshot_device_table_map, device);
+    PFN_vkGetDeviceProcAddr gpa = pDisp->GetDeviceProcAddr;
+    pDisp->GetSurfacePropertiesKHR = (PFN_vkGetSurfacePropertiesKHR) gpa(device, "vkGetSurfacePropertiesKHR");
+    pDisp->GetSurfaceFormatsKHR = (PFN_vkGetSurfaceFormatsKHR) gpa(device, "vkGetSurfaceFormatsKHR");
+    pDisp->GetSurfacePresentModesKHR = (PFN_vkGetSurfacePresentModesKHR) gpa(device, "vkGetSurfacePresentModesKHR");
+    pDisp->CreateSwapchainKHR = (PFN_vkCreateSwapchainKHR) gpa(device, "vkCreateSwapchainKHR");
+    pDisp->DestroySwapchainKHR = (PFN_vkDestroySwapchainKHR) gpa(device, "vkDestroySwapchainKHR");
+    pDisp->GetSwapchainImagesKHR = (PFN_vkGetSwapchainImagesKHR) gpa(device, "vkGetSwapchainImagesKHR");
+    pDisp->AcquireNextImageKHR = (PFN_vkAcquireNextImageKHR) gpa(device, "vkAcquireNextImageKHR");
+    pDisp->QueuePresentKHR = (PFN_vkQueuePresentKHR) gpa(device, "vkQueuePresentKHR");
     deviceExtMap[pDisp].wsi_enabled = false;
     for (i = 0; i < pCreateInfo->extensionCount; i++) {
         if (strcmp(pCreateInfo->ppEnabledExtensionNames[i], VK_EXT_KHR_DEVICE_SWAPCHAIN_EXTENSION_NAME) == 0)
@@ -671,7 +680,7 @@
         return (PFN_vkVoidFunction) vkCreateCommandPool;
 
     VkLayerDispatchTable *pDisp =  get_dispatch_table(screenshot_device_table_map, dev);
-    if (deviceExtMap.size() == 0 || deviceExtMap[pDisp].wsi_enabled)
+    if (deviceExtMap.size() != 0 && deviceExtMap[pDisp].wsi_enabled)
     {
         if (!strcmp(funcName, "vkCreateSwapchainKHR"))
             return (PFN_vkVoidFunction) vkCreateSwapchainKHR;
diff --git a/vk-generate.py b/vk-generate.py
index 9491ec1..c127ea2 100755
--- a/vk-generate.py
+++ b/vk-generate.py
@@ -124,7 +124,7 @@
             for proto in self.protos:
                 if proto.name == "CreateInstance" or proto.name == "GetGlobalExtensionProperties" or proto.name == "GetGlobalLayerProperties" or proto.params[0].ty == "VkInstance" or (proto.params[0].ty == "VkPhysicalDevice" and proto.name != "CreateDevice"):
                     continue
-                if proto.name != "GetDeviceProcAddr":
+                if proto.name != "GetDeviceProcAddr" and 'KHR' not in proto.name:
                     stmts.append("table->%s = (PFN_vk%s) gpa(baseDevice, \"vk%s\");" %
                         (proto.name, proto.name, proto.name))
             func.append("static inline void %s_initialize_dispatch_table(VkLayerDispatchTable *table,"
@@ -143,7 +143,7 @@
                     continue
                 if proto.name == "CreateDevice":
                     continue
-                if proto.name != "GetInstanceProcAddr":
+                if proto.name != "GetInstanceProcAddr" and 'KHR' not in proto.name:
                     stmts.append("table->%s = (PFN_vk%s) gpa(baseInstance, \"vk%s\");" %
                           (proto.name, proto.name, proto.name))
             func.append("static inline void %s_init_instance_dispatch_table(VkLayerInstanceDispatchTable *table,"
diff --git a/vk-layer-generate.py b/vk-layer-generate.py
index a22d6ff..a7971ed 100755
--- a/vk-layer-generate.py
+++ b/vk-layer-generate.py
@@ -435,7 +435,7 @@
                 extra_space = ""
                 for (ext_enable, ext_list) in extensions:
                     if 0 != len(ext_enable):
-                        func_body.append('    if (deviceExtMap.size() == 0 || deviceExtMap[pDisp].%s)' % ext_enable)
+                        func_body.append('    if (deviceExtMap.size() != 0 && deviceExtMap[pDisp].%s)' % ext_enable)
                         func_body.append('    {')
                         extra_space = "    "
                     for ext_name in ext_list:
@@ -593,6 +593,15 @@
         gen_header.append('{')
         gen_header.append('    uint32_t i;')
         gen_header.append('    VkLayerDispatchTable *pDisp  = device_dispatch_table(device);')
+        gen_header.append('    PFN_vkGetDeviceProcAddr gpa = pDisp->GetDeviceProcAddr;')
+        gen_header.append('    pDisp->GetSurfacePropertiesKHR = (PFN_vkGetSurfacePropertiesKHR) gpa(device, "vkGetSurfacePropertiesKHR");')
+        gen_header.append('    pDisp->GetSurfaceFormatsKHR = (PFN_vkGetSurfaceFormatsKHR) gpa(device, "vkGetSurfaceFormatsKHR");')
+        gen_header.append('    pDisp->GetSurfacePresentModesKHR = (PFN_vkGetSurfacePresentModesKHR) gpa(device, "vkGetSurfacePresentModesKHR");')
+        gen_header.append('    pDisp->CreateSwapchainKHR = (PFN_vkCreateSwapchainKHR) gpa(device, "vkCreateSwapchainKHR");')
+        gen_header.append('    pDisp->DestroySwapchainKHR = (PFN_vkDestroySwapchainKHR) gpa(device, "vkDestroySwapchainKHR");')
+        gen_header.append('    pDisp->GetSwapchainImagesKHR = (PFN_vkGetSwapchainImagesKHR) gpa(device, "vkGetSwapchainImagesKHR");')
+        gen_header.append('    pDisp->AcquireNextImageKHR = (PFN_vkAcquireNextImageKHR) gpa(device, "vkAcquireNextImageKHR");')
+        gen_header.append('    pDisp->QueuePresentKHR = (PFN_vkQueuePresentKHR) gpa(device, "vkQueuePresentKHR");\n')
         gen_header.append('    deviceExtMap[pDisp].wsi_enabled = false;')
         gen_header.append('    for (i = 0; i < pCreateInfo->extensionCount; i++) {')
         gen_header.append('        if (strcmp(pCreateInfo->ppEnabledExtensionNames[i], VK_EXT_KHR_DEVICE_SWAPCHAIN_EXTENSION_NAME) == 0)')
@@ -675,8 +684,10 @@
     def generate_body(self):
         self.layer_name = "Generic"
         extensions=[('wsi_enabled', 
-                     ['vkCreateSwapchainKHR', 'vkDestroySwapchainKHR',
-                      'vkGetSwapchainImagesKHR', 'vkQueuePresentKHR'])]
+                     ['vkGetSurfacePropertiesKHR', 'vkGetSurfaceFormatsKHR',
+                      'vkGetSurfacePresentModesKHR', 'vkCreateSwapchainKHR',
+                      'vkDestroySwapchainKHR', 'vkGetSwapchainImagesKHR',
+                      'vkAcquireNextImageKHR', 'vkQueuePresentKHR'])]
         body = [self._generate_layer_initialization(True),
                 self._generate_dispatch_entrypoints("VK_LAYER_EXPORT"),
                 self._gen_create_msg_callback(),
@@ -773,6 +784,15 @@
         header_txt.append('{')
         header_txt.append('    uint32_t i;')
         header_txt.append('    VkLayerDispatchTable *pDisp  = device_dispatch_table(device);')
+        header_txt.append('    PFN_vkGetDeviceProcAddr gpa = pDisp->GetDeviceProcAddr;')
+        header_txt.append('    pDisp->GetSurfacePropertiesKHR = (PFN_vkGetSurfacePropertiesKHR) gpa(device, "vkGetSurfacePropertiesKHR");')
+        header_txt.append('    pDisp->GetSurfaceFormatsKHR = (PFN_vkGetSurfaceFormatsKHR) gpa(device, "vkGetSurfaceFormatsKHR");')
+        header_txt.append('    pDisp->GetSurfacePresentModesKHR = (PFN_vkGetSurfacePresentModesKHR) gpa(device, "vkGetSurfacePresentModesKHR");')
+        header_txt.append('    pDisp->CreateSwapchainKHR = (PFN_vkCreateSwapchainKHR) gpa(device, "vkCreateSwapchainKHR");')
+        header_txt.append('    pDisp->DestroySwapchainKHR = (PFN_vkDestroySwapchainKHR) gpa(device, "vkDestroySwapchainKHR");')
+        header_txt.append('    pDisp->GetSwapchainImagesKHR = (PFN_vkGetSwapchainImagesKHR) gpa(device, "vkGetSwapchainImagesKHR");')
+        header_txt.append('    pDisp->AcquireNextImageKHR = (PFN_vkAcquireNextImageKHR) gpa(device, "vkAcquireNextImageKHR");')
+        header_txt.append('    pDisp->QueuePresentKHR = (PFN_vkQueuePresentKHR) gpa(device, "vkQueuePresentKHR");\n')
         header_txt.append('    deviceExtMap[pDisp].wsi_enabled = false;')
         header_txt.append('    for (i = 0; i < pCreateInfo->extensionCount; i++) {')
         header_txt.append('        if (strcmp(pCreateInfo->ppEnabledExtensionNames[i], VK_EXT_KHR_DEVICE_SWAPCHAIN_EXTENSION_NAME) == 0)')
@@ -928,8 +948,8 @@
                     if p.name == proto.params[y].name:
                         cp = True
             (pft, pfi) = self._get_printf_params(p.ty, p.name, cp, cpp=True)
-            if p.name == "pSwapchain":
-                log_func += '%s = " << %s->handle << ", ' % (p.name, p.name)
+            if p.name == "pSwapchain" or p.name == "pSwapchainImages":
+                log_func += '%s = " << ((%s == NULL) ? 0 : %s->handle) << ", ' % (p.name, p.name, p.name)
             elif p.name == "swapchain":
                 log_func += '%s = " << %s.handle << ", ' % (p.name, p.name)
             else:
@@ -1085,8 +1105,10 @@
     def generate_body(self):
         self.layer_name = "APIDump"
         extensions=[('wsi_enabled',
-                    ['vkCreateSwapchainKHR', 'vkDestroySwapchainKHR',
-                    'vkGetSwapchainImagesKHR', 'vkQueuePresentKHR'])]
+                     ['vkGetSurfacePropertiesKHR', 'vkGetSurfaceFormatsKHR',
+                      'vkGetSurfacePresentModesKHR', 'vkCreateSwapchainKHR',
+                      'vkDestroySwapchainKHR', 'vkGetSwapchainImagesKHR',
+                      'vkAcquireNextImageKHR', 'vkQueuePresentKHR'])]
         body = [self.generate_init(),
                 self._generate_dispatch_entrypoints("VK_LAYER_EXPORT"),
                 self._generate_layer_gpa_function(extensions)]
@@ -1597,8 +1619,10 @@
     def generate_body(self):
         self.layer_name = "ObjectTracker"
         extensions=[('wsi_enabled',
-                    ['vkCreateSwapchainKHR', 'vkDestroySwapchainKHR',
-                     'vkGetSwapchainImagesKHR', 'vkQueuePresentKHR'])]
+                     ['vkGetSurfacePropertiesKHR', 'vkGetSurfaceFormatsKHR',
+                      'vkGetSurfacePresentModesKHR', 'vkCreateSwapchainKHR',
+                      'vkDestroySwapchainKHR', 'vkGetSwapchainImagesKHR',
+                      'vkAcquireNextImageKHR', 'vkQueuePresentKHR'])]
         body = [self.generate_maps(),
                 self.generate_procs(),
                 self.generate_destroy_instance(),