capture/relay: Add capture interface classes (#4177)

* capture/relay: Add capture interface classes

Add `ModuleCapture` class for capturing `IModule`
  - The `IModule` can only be created from
    -- `ISession::loadModule`
    -- `ISession::loadModuleFromIRBlob`
    -- `ISession::loadModuleFromSource`
    -- `ISession::loadModuleFromSourceString`
    so, we create the `ModuleCapture` at those methods in `SessionCapture`
    class. We use a hash map to store a map from `IModule` to `ModuleCapture`
    to avoid creating new `ModuleCapture` when there is already an old one.
  - In `SessionCapture::getLoadedModule`, we will assert on not finding
    a `ModuleCapture` instance.

Add `EntryPointCapture` class for capturing `IEntryPoint`.
  - The `IEntryPoint` can only be created from:
    -- `IModule::findEntryPointByName`
    -- `IModule::findAndCheckEntryPoint`
    so, we create the `EntryPointCapture` at those methods in `ModuleCapture`.
    Similarly, we use a hash map to store a map from `IEntryPoint` to
    `EntryPointCapture`.
  - In `IModule::getDefinedEntryPoint`, we will assert on not finding
    a `EntryPointCapture` instance.

Add `CompositeComponentTypeCapture` class for capturing CompositeComponentType,
but since user is only exposed to `IComponentType`, so `CompositeComponentTypeCapture`
just inherits from `IComponentType`.
  - `CompositeComponentType` can only be created from:
    -- ISession::createCompositeComponentType
    so create it here.

Add `TypeConformanceCapture` class for capturing `ITypeConformance`.
  - The `ITypeConformance` can only be created from:
    -- `ISession::createTypeConformanceComponentType`
    so create it here.

In addition, because `EntryPointCapture` and `ModuleCapture` share a some
base class `IComponentType`, we generate the COM GUID for those two
classes to differentiate them.

* Fix the build issue

* Add nullptr check for output parameter

* define the SLANG_CAPTURE_ASSERT macro used in both debug and release build

1 files changed, 116 insertions, 25 deletions

diff --git a/source/slang-capture-replay/slang-session.cpp b/source/slang-capture-replay/slang-session.cpp
index 3d96af6d7..5cf029629 100644
--- a/source/slang-capture-replay/slang-session.cpp
+++ b/source/slang-capture-replay/slang-session.cpp
@@ -1,5 +1,8 @@
 #include "capture_utility.h"
 #include "slang-session.h"
+#include "slang-entrypoint.h"
+#include "slang-composite-component-type.h"
+#include "slang-type-conformance.h"
 
 namespace SlangCapture
 {
@@ -7,8 +10,8 @@ namespace SlangCapture
     SessionCapture::SessionCapture(slang::ISession* session)
         : m_actualSession(session)
     {
-        assert(m_actualSession);
-        slangCaptureLog(LogLevel::Verbose, "%s: %p\n", "Session", session);
+        SLANG_CAPTURE_ASSERT(m_actualSession);
+        slangCaptureLog(LogLevel::Verbose, "%s: %p\n", "SessionCapture create:", session);
     }
 
     SessionCapture::~SessionCapture()
@@ -26,7 +29,7 @@ namespace SlangCapture
 
     SLANG_NO_THROW slang::IGlobalSession* SessionCapture::getGlobalSession()
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IGlobalSession* pGlobalSession = m_actualSession->getGlobalSession();
         return pGlobalSession;
     }
@@ -35,9 +38,10 @@ namespace SlangCapture
         const char* moduleName,
         slang::IBlob**     outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IModule* pModule = m_actualSession->loadModule(moduleName, outDiagnostics);
-        return pModule;
+        ModuleCapture* pModuleCapture = getModuleCapture(pModule);
+        return static_cast<slang::IModule*>(pModuleCapture);
     }
 
     SLANG_NO_THROW slang::IModule* SessionCapture::loadModuleFromIRBlob(
@@ -46,9 +50,10 @@ namespace SlangCapture
         slang::IBlob* source,
         slang::IBlob** outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IModule* pModule = m_actualSession->loadModuleFromIRBlob(moduleName, path, source, outDiagnostics);
-        return pModule;
+        ModuleCapture* pModuleCapture = getModuleCapture(pModule);
+        return static_cast<slang::IModule*>(pModuleCapture);
     }
 
     SLANG_NO_THROW slang::IModule* SessionCapture::loadModuleFromSource(
@@ -57,9 +62,10 @@ namespace SlangCapture
         slang::IBlob* source,
         slang::IBlob** outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IModule* pModule = m_actualSession->loadModuleFromSource(moduleName, path, source, outDiagnostics);
-        return pModule;
+        ModuleCapture* pModuleCapture = getModuleCapture(pModule);
+        return static_cast<slang::IModule*>(pModuleCapture);
     }
 
     SLANG_NO_THROW slang::IModule* SessionCapture::loadModuleFromSourceString(
@@ -68,9 +74,10 @@ namespace SlangCapture
         const char* string,
         slang::IBlob** outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IModule* pModule = m_actualSession->loadModuleFromSourceString(moduleName, path, string, outDiagnostics);
-        return pModule;
+        ModuleCapture* pModuleCapture = getModuleCapture(pModule);
+        return static_cast<slang::IModule*>(pModuleCapture);
     }
 
     SLANG_NO_THROW SlangResult SessionCapture::createCompositeComponentType(
@@ -79,8 +86,26 @@ namespace SlangCapture
         slang::IComponentType**         outCompositeComponentType,
         ISlangBlob**                    outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
-        SlangResult result = m_actualSession->createCompositeComponentType(componentTypes, componentTypeCount, outCompositeComponentType, outDiagnostics);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
+
+        Slang::List<slang::IComponentType*> componentTypeList;
+
+        // get the actual component types from our capture wrappers
+        if(SLANG_OK != getActualComponentTypes(componentTypes, componentTypeCount, componentTypeList))
+        {
+            SLANG_CAPTURE_ASSERT(!"Failed to get actual component types");
+        }
+
+        SlangResult result = m_actualSession->createCompositeComponentType(
+                componentTypeList.getBuffer(), componentTypeCount, outCompositeComponentType, outDiagnostics);
+
+        if (SLANG_OK == result)
+        {
+            CompositeComponentTypeCapture* compositeComponentTypeCapture = new CompositeComponentTypeCapture(*outCompositeComponentType);
+            Slang::ComPtr<CompositeComponentTypeCapture> resultCapture(compositeComponentTypeCapture);
+            *outCompositeComponentType = resultCapture.detach();
+        }
+
         return result;
     }
 
@@ -90,7 +115,7 @@ namespace SlangCapture
         SlangInt                        specializationArgCount,
         ISlangBlob**                    outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::TypeReflection* pTypeReflection = m_actualSession->specializeType(type, specializationArgs, specializationArgCount, outDiagnostics);
         return pTypeReflection;
     }
@@ -101,7 +126,7 @@ namespace SlangCapture
         slang::LayoutRules     rules,
         ISlangBlob**    outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::TypeLayoutReflection* pTypeLayoutReflection = m_actualSession->getTypeLayout(type, targetIndex, rules, outDiagnostics);
         return pTypeLayoutReflection;
     }
@@ -111,14 +136,14 @@ namespace SlangCapture
         slang::ContainerType containerType,
         ISlangBlob** outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::TypeReflection* pTypeReflection = m_actualSession->getContainerType(elementType, containerType, outDiagnostics);
         return pTypeReflection;
     }
 
     SLANG_NO_THROW slang::TypeReflection* SessionCapture::getDynamicType()
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::TypeReflection* pTypeReflection = m_actualSession->getDynamicType();
         return pTypeReflection;
     }
@@ -127,7 +152,7 @@ namespace SlangCapture
         slang::TypeReflection* type,
         ISlangBlob** outNameBlob)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         SlangResult result = m_actualSession->getTypeRTTIMangledName(type, outNameBlob);
         return result;
     }
@@ -137,7 +162,7 @@ namespace SlangCapture
         slang::TypeReflection* interfaceType,
         ISlangBlob** outNameBlob)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         SlangResult result = m_actualSession->getTypeConformanceWitnessMangledName(type, interfaceType, outNameBlob);
         return result;
     }
@@ -147,7 +172,7 @@ namespace SlangCapture
         slang::TypeReflection* interfaceType,
         uint32_t*              outId)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         SlangResult result = m_actualSession->getTypeConformanceWitnessSequentialID(type, interfaceType, outId);
         return result;
     }
@@ -159,38 +184,104 @@ namespace SlangCapture
         SlangInt conformanceIdOverride,
         ISlangBlob** outDiagnostics)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
+
         SlangResult result = m_actualSession->createTypeConformanceComponentType(type, interfaceType, outConformance, conformanceIdOverride, outDiagnostics);
+
+        if (SLANG_OK != result)
+        {
+            TypeConformanceCapture* conformanceCapture = new TypeConformanceCapture(*outConformance);
+            Slang::ComPtr<TypeConformanceCapture> resultCapture(conformanceCapture);
+            *outConformance = resultCapture.detach();
+        }
+
         return result;
     }
 
     SLANG_NO_THROW SlangResult SessionCapture::createCompileRequest(
         SlangCompileRequest**   outCompileRequest)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         SlangResult result = m_actualSession->createCompileRequest(outCompileRequest);
         return result;
     }
 
     SLANG_NO_THROW SlangInt SessionCapture::getLoadedModuleCount()
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         SlangInt count = m_actualSession->getLoadedModuleCount();
         return count;
     }
 
     SLANG_NO_THROW slang::IModule* SessionCapture::getLoadedModule(SlangInt index)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         slang::IModule* pModule = m_actualSession->getLoadedModule(index);
+
+        if (pModule)
+        {
+            ModuleCapture* moduleCapture = m_mapModuleToCapture.tryGetValue(pModule);
+            if (!moduleCapture)
+            {
+                SLANG_CAPTURE_ASSERT(!"Module not found in mapModuleToCapture");
+            }
+            return static_cast<slang::IModule*>(moduleCapture);
+        }
+
         return pModule;
     }
 
     SLANG_NO_THROW bool SessionCapture::isBinaryModuleUpToDate(const char* modulePath, slang::IBlob* binaryModuleBlob)
     {
-        slangCaptureLog(LogLevel::Verbose, "%s\n", __func__);
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
         bool result = m_actualSession->isBinaryModuleUpToDate(modulePath, binaryModuleBlob);
         return result;
     }
 
+    ModuleCapture* SessionCapture::getModuleCapture(slang::IModule* module)
+    {
+        ModuleCapture* moduleCapture = nullptr;
+        moduleCapture = m_mapModuleToCapture.tryGetValue(module);
+        if (!moduleCapture)
+        {
+            moduleCapture = new ModuleCapture(module);
+            Slang::ComPtr<ModuleCapture> result(moduleCapture);
+            m_mapModuleToCapture.add(module, *result.detach());
+        }
+        return moduleCapture;
+    }
+
+    SlangResult SessionCapture::getActualComponentTypes(
+            slang::IComponentType* const*   componentTypes,
+            SlangInt                        componentTypeCount,
+            List<slang::IComponentType*>&   outActualComponentTypes)
+    {
+        for (SlangInt i = 0; i < componentTypeCount; i++)
+        {
+            slang::IComponentType* const& componentType = componentTypes[i];
+            void* outObj = nullptr;
+
+            if (componentType->queryInterface(ModuleCapture::getTypeGuid(), &outObj) == SLANG_OK)
+            {
+                ModuleCapture* moduleCapture = static_cast<ModuleCapture*>(outObj);
+                outActualComponentTypes.add(moduleCapture->getActualModule());
+            }
+            else if (componentType->queryInterface(EntryPointCapture::getTypeGuid(), &outObj) == SLANG_OK)
+            {
+                EntryPointCapture* entrypointCapture = static_cast<EntryPointCapture*>(outObj);
+                outActualComponentTypes.add(entrypointCapture->getActualEntryPoint());
+            }
+            // will fall back to the actual component type, it means that we didn't capture this type.
+            else
+            {
+                outActualComponentTypes.add(componentType);
+            }
+        }
+
+        if (componentTypeCount == outActualComponentTypes.getCount())
+        {
+            return SLANG_OK;
+        }
+        return SLANG_FAIL;
+    }
 } // namespace SlangCapture