Initial implementation for decl-tree reflection API (#4666)

* Initial implementation for decl-tree reflection API

This patch adds Slang API methods for walking all the declarations in the AST.

We expose this functionality through an abstract `DeclReflection` class that can be a type, function or a variable declaration.
We also provide ways to cast the decl to a `FunctionReflection`, `TypeReflection` or `VariableReflection` and traverse through the child nodes (for instance, a struct type will have component variable declarations)

This patch also adds `ISlangInternal` as an internal COM interface to allow us to cast IGlobalSession to the internal Session pointer while bypassing any wrappers (such as the capture interface)

* Update slang.h

* Remove `ISlangInternal` (its causing a diamond pattern w.r.t `ISlangUnknown`) and use `ComPtr` for proper ref management.

* Update unit-test-decl-tree-reflection.cpp

* Change `FunctionDeclBase` to use `DeclRef` instead of directly using the decl.

* Update slang-reflection-api.cpp

7 files changed, 200 insertions, 28 deletions

diff --git a/source/slang-capture-replay/slang-global-session.cpp b/source/slang-capture-replay/slang-global-session.cpp
index b8dfd7d40..d8f1e361d 100644
--- a/source/slang-capture-replay/slang-global-session.cpp
+++ b/source/slang-capture-replay/slang-global-session.cpp
@@ -28,11 +28,23 @@ namespace SlangCapture
         m_actualGlobalSession->release();
     }
 
-    ISlangUnknown* GlobalSessionCapture::getInterface(const Guid& guid)
+    SLANG_NO_THROW SlangResult SLANG_MCALL GlobalSessionCapture::queryInterface(SlangUUID const& uuid, void** outObject) 
     {
-        if(guid == ISlangUnknown::getTypeGuid() || guid == IGlobalSession::getTypeGuid())
-            return asExternal(this);
-        return nullptr;
+        if (uuid == Session::getTypeGuid())
+        {
+            // no add-ref here, the query will cause the inner session to handle the add-ref.
+            this->m_actualGlobalSession->queryInterface(uuid, outObject);
+            return SLANG_OK;
+        }
+
+        if (uuid == ISlangUnknown::getTypeGuid() && uuid == IGlobalSession::getTypeGuid())
+        {
+            addReference();
+            *outObject = static_cast<slang::IGlobalSession*>(this);
+            return SLANG_OK;
+        }
+
+        return SLANG_E_NO_INTERFACE;
     }
 
     SLANG_NO_THROW SlangResult SLANG_MCALL GlobalSessionCapture::createSession(slang::SessionDesc const&  desc, slang::ISession** outSession)
diff --git a/source/slang-capture-replay/slang-global-session.h b/source/slang-capture-replay/slang-global-session.h
index ae451505f..f478e60db 100644
--- a/source/slang-capture-replay/slang-global-session.h
+++ b/source/slang-capture-replay/slang-global-session.h
@@ -17,9 +17,10 @@ namespace SlangCapture
             explicit GlobalSessionCapture(slang::IGlobalSession* session);
             virtual ~GlobalSessionCapture();
 
-            SLANG_REF_OBJECT_IUNKNOWN_ALL
+            SLANG_REF_OBJECT_IUNKNOWN_ADD_REF
+            SLANG_REF_OBJECT_IUNKNOWN_RELEASE
 
-            ISlangUnknown* getInterface(const Guid& guid);
+            SLANG_NO_THROW SlangResult SLANG_MCALL queryInterface(SlangUUID const& uuid, void** outObject) SLANG_OVERRIDE;
 
             // slang::IGlobalSession
             SLANG_NO_THROW SlangResult SLANG_MCALL createSession(slang::SessionDesc const&  desc, slang::ISession** outSession) override;
diff --git a/source/slang-capture-replay/slang-module.cpp b/source/slang-capture-replay/slang-module.cpp
index 8a1a80126..273faa59d 100644
--- a/source/slang-capture-replay/slang-module.cpp
+++ b/source/slang-capture-replay/slang-module.cpp
@@ -27,6 +27,14 @@ namespace SlangCapture
             return nullptr;
     }
 
+    SLANG_NO_THROW slang::DeclReflection* ModuleCapture::getModuleReflection()
+    {
+        // No need to capture this call as it is just a query.
+        slangCaptureLog(LogLevel::Verbose, "%s\n", __PRETTY_FUNCTION__);
+        slang::DeclReflection* res = (slang::DeclReflection*)m_actualModule->getModuleReflection();
+        return res;
+    }
+
     SLANG_NO_THROW SlangResult ModuleCapture::findEntryPointByName(
         char const*     name,
         slang::IEntryPoint**   outEntryPoint)
diff --git a/source/slang-capture-replay/slang-module.h b/source/slang-capture-replay/slang-module.h
index de63b7967..94539532c 100644
--- a/source/slang-capture-replay/slang-module.h
+++ b/source/slang-capture-replay/slang-module.h
@@ -83,6 +83,7 @@ namespace SlangCapture
             uint32_t compilerOptionEntryCount,
             slang::CompilerOptionEntry* compilerOptionEntries,
             ISlangBlob** outDiagnostics = nullptr) override;
+        virtual SLANG_NO_THROW slang::DeclReflection* getModuleReflection() override;
 
         slang::IModule* getActualModule() const { return m_actualModule; }
     private:
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 35b2fcf44..2409cedfb 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -282,6 +282,7 @@ namespace Slang
         HashSet<SourceFile*> m_fileSet;
     };
 
+
     class EntryPoint;
 
     class ComponentType;
@@ -1051,7 +1052,7 @@ namespace Slang
 
         SLANG_NO_THROW slang::FunctionReflection* SLANG_MCALL getFunctionReflection() SLANG_OVERRIDE
         {
-            return (slang::FunctionReflection*)m_funcDeclRef.getDecl();
+            return (slang::FunctionReflection*)m_funcDeclRef.declRefBase;
         }
     protected:
         void acceptVisitor(ComponentTypeVisitor* visitor, SpecializationInfo* specializationInfo) SLANG_OVERRIDE;
@@ -1476,6 +1477,8 @@ namespace Slang
 
         virtual void buildHash(DigestBuilder<SHA1>& builder) SLANG_OVERRIDE;
 
+        virtual slang::DeclReflection* getModuleReflection() SLANG_OVERRIDE;
+
         void setDigest(SHA1::Digest const& digest) { m_digest = digest; }
         SHA1::Digest computeDigest();
 
@@ -3079,9 +3082,11 @@ namespace Slang
     class Session : public RefObject, public slang::IGlobalSession
     {
     public:
-        SLANG_REF_OBJECT_IUNKNOWN_ALL
+        SLANG_COM_INTERFACE(0xd6b767eb, 0xd786, 0x4343, { 0x2a, 0x8c, 0x6d, 0xa0, 0x3d, 0x5a, 0xb4, 0x4a })
 
-        ISlangUnknown* getInterface(const Guid& guid);
+        SLANG_NO_THROW SlangResult SLANG_MCALL queryInterface(SlangUUID const& uuid, void** outObject) SLANG_OVERRIDE;
+        SLANG_REF_OBJECT_IUNKNOWN_ADD_REF
+        SLANG_REF_OBJECT_IUNKNOWN_RELEASE
 
         // slang::IGlobalSession 
         SLANG_NO_THROW SlangResult SLANG_MCALL createSession(slang::SessionDesc const&  desc, slang::ISession** outSession) override;
@@ -3287,9 +3292,11 @@ SLANG_FORCE_INLINE slang::IGlobalSession* asExternal(Session* session)
     return static_cast<slang::IGlobalSession*>(session);
 }
 
-SLANG_FORCE_INLINE Session* asInternal(slang::IGlobalSession* session)
+SLANG_FORCE_INLINE ComPtr<Session> asInternal(slang::IGlobalSession* session)
 {
-    return static_cast<Session*>(session);
+    Slang::Session* internalSession = nullptr;
+    session->queryInterface(SLANG_IID_PPV_ARGS(&internalSession));
+    return ComPtr<Session>(INIT_ATTACH, static_cast<Session*>(internalSession));
 }
 
 SLANG_FORCE_INLINE slang::ISession* asExternal(Linkage* linkage)
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index d0d09f814..9e08330e4 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -65,14 +65,15 @@ static inline SlangReflectionVariable* convert(Decl* var)
     return (SlangReflectionVariable*) var;
 }
 
-static inline FunctionDeclBase* convert(SlangReflectionFunction* func)
+static inline DeclRef<FunctionDeclBase> convert(SlangReflectionFunction* func)
 {
-    return (FunctionDeclBase*)func;
+    DeclRefBase* declBase = (DeclRefBase*)func;
+    return DeclRef<FunctionDeclBase>(declBase);
 }
 
-static inline SlangReflectionFunction* convert(FunctionDeclBase* func)
+static inline SlangReflectionFunction* convert(DeclRef<FunctionDeclBase> func)
 {
-    return (SlangReflectionFunction*)func;
+    return (SlangReflectionFunction*)func.declRefBase;
 }
 
 static inline VarLayout* convert(SlangReflectionVariableLayout* var)
@@ -773,7 +774,7 @@ SLANG_API SlangReflectionFunction* spReflection_FindFunctionByName(SlangReflecti
     {
         auto result = program->findDeclFromString(name, &sink);
         if (auto funcDeclRef = result.as<FunctionDeclBase>())
-            return (SlangReflectionFunction*)funcDeclRef.getDecl();
+            return convert(funcDeclRef);
     }
     catch (...)
     {
@@ -2621,6 +2622,9 @@ SLANG_API SlangReflectionModifier* spReflectionVariable_FindModifier(SlangReflec
     case SLANG_MODIFIER_DIFFERENTIABLE:
         modifier = var->findModifier<DifferentiableAttribute>();
         break;
+    case SLANG_MODIFIER_MUTATING:
+        modifier = var->findModifier<MutatingAttribute>();
+        break;
     default:
         return nullptr;
     }
@@ -2647,6 +2651,17 @@ SLANG_API SlangReflectionUserAttribute* spReflectionVariable_FindUserAttributeBy
     return findUserAttributeByName(asInternal(session), varDecl, name);
 }
 
+SLANG_API bool spReflectionVariable_HasDefaultValue(SlangReflectionVariable* inVar)
+{
+    auto decl = convert(inVar);
+    if (auto varDecl = as<VarDeclBase>(decl))
+    {
+        return varDecl->initExpr != nullptr;
+    }
+    
+    return false;
+}
+
 // Variable Layout Reflection
 
 SLANG_API SlangReflectionVariable* spReflectionVariableLayout_GetVariable(SlangReflectionVariableLayout* inVarLayout)
@@ -2793,11 +2808,18 @@ SLANG_API SlangStage spReflectionVariableLayout_getStage(
 
 // Function Reflection
 
+SLANG_API SlangReflectionDecl* spReflectionFunction_asDecl(SlangReflectionFunction* inFunc)
+{
+    auto func = convert(inFunc);
+    if (!func) return nullptr;
+    return (SlangReflectionDecl*)func.getDecl();
+}
+
 SLANG_API char const* spReflectionFunction_GetName(SlangReflectionFunction* inFunc)
 {
     auto func = convert(inFunc);
     if (!func) return nullptr;
-    return getText(func->getName()).getBuffer();
+    return getText(func.getDecl()->getName()).getBuffer();
 }
 
 SLANG_API SlangReflectionType* spReflectionFunction_GetResultType(SlangReflectionFunction* inFunc)
@@ -2805,42 +2827,146 @@ SLANG_API SlangReflectionType* spReflectionFunction_GetResultType(SlangReflectio
     auto func = convert(inFunc);
     if (!func) return nullptr;
 
-    return convert(func->returnType.type);
+    auto rawType = func.getDecl()->returnType.type;
+    auto astBuilder = rawType->getASTBuilderForReflection();
+
+    return convert((Type*)rawType->substitute(astBuilder, SubstitutionSet(func.declRefBase)));
+}
+
+SLANG_API SlangReflectionModifier* spReflectionFunction_FindModifier(SlangReflectionFunction* inFunc, SlangModifierID modifierID)
+{
+    auto funcDeclRef = convert(inFunc);
+    auto varRefl = convert(funcDeclRef.getDecl());
+    if (!varRefl) return nullptr;
+
+    return spReflectionVariable_FindModifier(varRefl, modifierID);
 }
 
 SLANG_API unsigned int spReflectionFunction_GetUserAttributeCount(SlangReflectionFunction* inFunc)
 {
     auto func = convert(inFunc);
     if (!func) return 0;
-    return getUserAttributeCount(func);
+    return getUserAttributeCount(func.getDecl());
 }
 
 SLANG_API SlangReflectionUserAttribute* spReflectionFunction_GetUserAttribute(SlangReflectionFunction* inFunc, unsigned int index)
 {
     auto func = convert(inFunc);
     if (!func) return nullptr;
-    return getUserAttributeByIndex(func, index);
+    return getUserAttributeByIndex(func.getDecl(), index);
 }
 
 SLANG_API SlangReflectionUserAttribute* spReflectionFunction_FindUserAttributeByName(SlangReflectionFunction* inFunc, SlangSession* session, char const* name)
 {
     auto func = convert(inFunc);
     if (!func) return nullptr;
-    return findUserAttributeByName(asInternal(session), func, name);
+    return findUserAttributeByName(asInternal(session), func.getDecl(), name);
 }
 
 SLANG_API unsigned int spReflectionFunction_GetParameterCount(SlangReflectionFunction* inFunc)
 {
     auto func = convert(inFunc);
     if (!func) return 0;
-    return (unsigned int)func->getParameters().getCount();
+    return (unsigned int)func.getDecl()->getParameters().getCount();
 }
 
 SLANG_API SlangReflectionVariable* spReflectionFunction_GetParameter(SlangReflectionFunction* inFunc, unsigned int index)
 {
     auto func = convert(inFunc);
     if (!func) return nullptr;
-    return convert(as<Decl>(func->getParameters()[index]));
+    return convert(as<Decl>(func.getDecl()->getParameters()[index]));
+}
+
+// Abstract decl reflection
+
+SLANG_API unsigned int spReflectionDecl_getChildrenCount(SlangReflectionDecl* parentDecl)
+{
+    Decl* decl = (Decl*)parentDecl;
+    if (as<ContainerDecl>(decl))
+    {
+        return (unsigned int)as<ContainerDecl>(decl)->members.getCount();
+    }
+    
+    return 0;
+}
+
+SLANG_API SlangReflectionDecl* spReflectionDecl_getChild(SlangReflectionDecl* parentDecl, unsigned int index)
+{
+    Decl* decl = (Decl*)parentDecl;
+    if (auto containerDecl = as<ContainerDecl>(decl))
+    {
+        if (containerDecl->members.getCount() > index)
+            return (SlangReflectionDecl*)containerDecl->members[index];
+    }
+
+    return nullptr;
+}
+
+SLANG_API SlangDeclKind spReflectionDecl_getKind(SlangReflectionDecl* decl)
+{
+    Decl* slangDecl = (Decl*)decl;
+    if (as<StructDecl>(slangDecl))
+    {
+        return SLANG_DECL_KIND_STRUCT;
+    }
+    else if (as<VarDeclBase>(slangDecl))
+    {
+        return SLANG_DECL_KIND_VARIABLE;
+    }
+    else if (as<GenericDecl>(slangDecl))
+    {
+        return SLANG_DECL_KIND_GENERIC;
+    }
+    else if (as<FunctionDeclBase>(slangDecl))
+    {
+        return SLANG_DECL_KIND_FUNC;
+    }
+    else if (as<ModuleDecl>(slangDecl))
+    {
+        return SLANG_DECL_KIND_MODULE;
+    }
+    else
+        return SLANG_DECL_KIND_UNSUPPORTED_FOR_REFLECTION;
+}
+
+SLANG_API SlangReflectionFunction* spReflectionDecl_castToFunction(SlangReflectionDecl* decl)
+{
+    Decl* slangDecl = (Decl*) decl;
+    if (auto funcDecl = as<FunctionDeclBase>(slangDecl))
+    {
+        return convert(DeclRef<FunctionDeclBase>(funcDecl->getDefaultDeclRef()));
+    }
+
+    // Improper cast
+    return nullptr;
+}
+
+SLANG_API SlangReflectionVariable* spReflectionDecl_castToVariable(SlangReflectionDecl* decl)
+{
+    Decl* slangDecl = (Decl*) decl;
+    if (auto varDecl = as<VarDeclBase>(slangDecl))
+    {
+        return (SlangReflectionVariable*) varDecl;
+    }
+
+    // Improper cast
+    return nullptr;
+
+}
+
+SLANG_API SlangReflectionType* spReflection_getTypeFromDecl(SlangSession* session, SlangReflectionDecl* decl)
+{
+    Decl* slangDecl = (Decl*)decl;
+    auto slangSession = asInternal(session);
+
+    ASTBuilder* builder = slangSession->getGlobalASTBuilder();
+    if (auto type = DeclRefType::create(builder, slangDecl->getDefaultDeclRef()))
+    {
+        return convert(type);
+    }
+
+    // Couldn't create a type from the decl
+    return nullptr;
 }
 
 // Shader Parameter Reflection
@@ -2906,7 +3032,7 @@ SLANG_API SlangReflectionFunction* spReflectionEntryPoint_GetFunction(SlangRefle
     auto entryPointLayout = convert(inEntryPoint);
     if (entryPointLayout)
     {
-        return (SlangReflectionFunction*)entryPointLayout->entryPoint.getDecl();
+        return convert(entryPointLayout->entryPoint);
     }
     return nullptr;
 }
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index f28d3cdeb..f440ce3cb 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -546,11 +546,23 @@ SlangResult Session::_readBuiltinModule(ISlangFileSystem* fileSystem, Scope* sco
     return SLANG_OK;
 }
 
-ISlangUnknown* Session::getInterface(const Guid& guid)
+SLANG_NO_THROW SlangResult SLANG_MCALL Session::queryInterface(SlangUUID const& uuid, void** outObject)
 {
-    if(guid == ISlangUnknown::getTypeGuid() || guid == IGlobalSession::getTypeGuid())
-        return asExternal(this);
-    return nullptr;
+    if (uuid == Session::getTypeGuid())
+    {
+        addReference();
+        *outObject = static_cast<Session*>(this);
+        return SLANG_OK;
+    }
+
+    if (uuid == ISlangUnknown::getTypeGuid() && uuid == IGlobalSession::getTypeGuid())
+    {
+        addReference();
+        *outObject = static_cast<slang::IGlobalSession*>(this);
+        return SLANG_OK;
+    }
+
+    return SLANG_E_NO_INTERFACE;
 }
 
 static size_t _getStructureSize(const uint8_t* src)
@@ -4070,6 +4082,11 @@ void Module::buildHash(DigestBuilder<SHA1>& builder)
     builder.append(computeDigest());
 }
 
+slang::DeclReflection* Module::getModuleReflection()
+{
+    return (slang::DeclReflection*)m_moduleDecl;
+}
+
 SHA1::Digest Module::computeDigest()
 {
     if (m_digest == SHA1::Digest())