Some Slang API additions (#1195)

* Some Slang API additions

These are additions to the public Slang API that came up while I was trying to write an example to demonstrate GPU printing. The additions aren't strictly necessary for the example, but I found these to be missing services when writing the code the way I wanted to.

The main public changes are:

* There is a new distinct `IEntryPoint` interface which inherit from `IComponentType` (much like `IModule` before). For now this doesn't expose any new functions on top of `IComponentType`, but I expect it to do so eventually.

* It is now possible to get the `IModule` for a specific translation unit in a compile request with `spCompileRequest_getModule`. Even for a compile request that had only one translation unit this is *not* the same object as gets returned by `spCompileRequest_getProgram`. The latter should probably be called `*_getLinkedProgram` because it returns a composite component type that links the module with everything it `import`s.

* An `IModule` can look up entry points declared in the module. Eventually a module should support looking up most of the declarations in the module (e.g., types) by name, but entry points are an obvious case.

* A new `link()` operation is added to `IComponentType`. It is possible to have component types that have unsatisfied dependencies, such that trying to generate kernel code from them will fail. The `link()` operation tries to produce a new composite component type that combines a component with its dependencies, to enable code generation. The implementation of end-to-end compilation was using a function like this internally, but it hadn't been exposed to the API.

Notes on the implementation:

* The list of entry points declared in a given translation unit has moved from `TranslationUnitRequest` to the `Module` inside of it.

* `EntryPoint` now has to do a song and dance much like `Module` to both inherit from `ComponentType` and support the `IEntryPoint` interface.

* The `Session* m_session` member in `Linkage` (in terms of public API, this is the `slang::ISession` holding a pointer to the `slang::IGlobalSession`) has been changed to a `RefPtr`. Without this change an application can't just hold onto a `ComPtr<slang::ISession>`; they also need to retain the `IGlobalSession` or things will crash. The new behavior seems more correct, but I worry that it might introduce a leak.

* The `asInternal` operation for `IComponentType` had to be updated to not just perform a cast. A type like `Module` has two `IComponentType` sub-objects, and only one of these is at the same address as the `ComponentType` base.

* Similarly, the `Module::getInterface` logic was changed to fall back to `Super::getInterface` for all the cases other than `IID_IModule`, so that it would be guaranteed to return the `IComponentType` at the same address as `ComponentType` in response to a `queryInterface`.

* Fixes for memory retain cycles

As part of the earlier change, I made the `Linkage` type hold a `RefPtr` to the `Session`.
The motivation there is it lets a user hang onto just a `slang::ISession` without having to also retain the `slang::IGlobalSession` for no immediately apparent reason.

There are two problems that this surfaced, one pre-existing and one new.

The new problem was that `Session` already held a `RefPtr<Linkage> m_builtinLinkage` for the linkage that holds the stdlib code.
I solved that problem by splitting the parent pointer in `Linkage` into two pointers: a raw pointer that is used to actually locate the parent session, and a ref-counted one that can be used to *optionally* retain the parent session.
The builtin linkage is then set up to explicitly not retain its parent, thus breaking the cycle.

The second problem was a pre-existing one, where every `ComponentType` was holding a retained pointer to its parent `Linkage`, but in turn the `Linkage` was holding retained pointers to many `ComponentTypes` (and subclasses thereof).
For this case I used the more expedient fix of making the parent pointer into a raw pointer, and figuring that it is a reasonable rule to expect user to retain the `Linkage` (aka `slang::ISession`) that owns a component type if they want to be able to use the component type.

I might need/want to investigate a better fix for the latter issue, but for now this seems to clean up the issues I was seeing in the tests. Fingers crossed.

5 files changed, 239 insertions, 22 deletions

diff --git a/source/slang/slang-check-shader.cpp b/source/slang/slang-check-shader.cpp
index 611f6be58..da0313049 100644
--- a/source/slang/slang-check-shader.cpp
+++ b/source/slang/slang-check-shader.cpp
@@ -828,7 +828,7 @@ namespace Slang
                     // but didn't specify any groups (since the current
                     // compilation API doesn't allow for grouping).
                     //
-                    entryPointReq->getTranslationUnit()->entryPoints.add(entryPoint);
+                    entryPointReq->getTranslationUnit()->module->_addEntryPoint(entryPoint);
 
                     outUnspecializedEntryPoints.add(entryPoint);
                     allComponentTypes.add(entryPoint);
@@ -907,7 +907,7 @@ namespace Slang
                     // group, so that its entry-point parameters lay out
                     // independent of the others.
                     //
-                    translationUnit->entryPoints.add(entryPoint);
+                    translationUnit->module->_addEntryPoint(entryPoint);
 
                     outUnspecializedEntryPoints.add(entryPoint);
                     allComponentTypes.add(entryPoint);
@@ -1264,7 +1264,7 @@ namespace Slang
         if(sink->GetErrorCount())
             return nullptr;
 
-        return unspecializedEntryPoint->specialize(
+        return ((ComponentType*) unspecializedEntryPoint)->specialize(
             args.getBuffer(),
             args.getCount(),
             sink);
diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index d0fca1ab5..65c48d2f2 100644
--- a/source/slang/slang-compiler.cpp
+++ b/source/slang/slang-compiler.cpp
@@ -196,6 +196,16 @@ namespace Slang
     // EntryPoint
     //
 
+    static const Guid IID_IEntryPoint = SLANG_UUID_IEntryPoint;
+
+    ISlangUnknown* EntryPoint::getInterface(const Guid& guid)
+    {
+        if(guid == IID_IEntryPoint)
+            return static_cast<slang::IEntryPoint*>(this);
+
+        return Super::getInterface(guid);
+    }
+
     RefPtr<EntryPoint> EntryPoint::create(
         Linkage*            linkage,
         DeclRef<FuncDecl>   funcDeclRef,
@@ -2059,6 +2069,15 @@ SlangResult dissassembleDXILUsingDXC(
         if( result.format != ResultFormat::None )
             return result;
 
+        // If we haven't yet computed a layout for this target
+        // program, we need to make sure that is done before
+        // code generation.
+        //
+        if( !getOrCreateIRModuleForLayout(sink) )
+        {
+            return result;
+        }
+
         RefPtr<BackEndCompileRequest> backEndRequest = new BackEndCompileRequest(
             m_program->getLinkage(),
             sink,
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 2783197af..2c58cb901 100644
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -286,6 +286,9 @@ namespace Slang
             SlangInt                        specializationArgCount,
             slang::IComponentType**         outSpecializedComponentType,
             ISlangBlob**                    outDiagnostics) SLANG_OVERRIDE;
+        SLANG_NO_THROW SlangResult SLANG_MCALL link(
+            slang::IComponentType**         outLinkedComponentType,
+            ISlangBlob**                    outDiagnostics) SLANG_OVERRIDE;
 
             /// Get the linkage (aka "session" in the public API) for this component type.
         Linkage* getLinkage() { return m_linkage; }
@@ -465,7 +468,7 @@ namespace Slang
         ComponentType(Linkage* linkage);
 
     private:
-        RefPtr<Linkage> m_linkage;
+        Linkage* m_linkage;
 
         // Cache of target-specific programs for each target.
         Dictionary<TargetRequest*, RefPtr<TargetProgram>> m_targetPrograms;
@@ -637,9 +640,61 @@ namespace Slang
         /// `getName()` and `getProfile()` methods should be expected to
         /// return useful data on pass-through entry points.
         ///
-    class EntryPoint : public ComponentType
+    class EntryPoint : public ComponentType, public slang::IEntryPoint
     {
+        typedef ComponentType Super;
+
     public:
+        SLANG_REF_OBJECT_IUNKNOWN_ALL
+
+        ISlangUnknown* getInterface(const Guid& guid);
+
+
+        // Forward `IComponentType` methods
+
+        SLANG_NO_THROW slang::ISession* SLANG_MCALL getSession() SLANG_OVERRIDE
+        {
+            return Super::getSession();
+        }
+
+        SLANG_NO_THROW slang::ProgramLayout* SLANG_MCALL getLayout(
+            SlangInt        targetIndex,
+            slang::IBlob**  outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::getLayout(targetIndex, outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL getEntryPointCode(
+            SlangInt        entryPointIndex,
+            SlangInt        targetIndex,
+            slang::IBlob**  outCode,
+            slang::IBlob**  outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::getEntryPointCode(entryPointIndex, targetIndex, outCode, outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL specialize(
+            slang::SpecializationArg const* specializationArgs,
+            SlangInt                        specializationArgCount,
+            slang::IComponentType**         outSpecializedComponentType,
+            ISlangBlob**                    outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::specialize(
+                specializationArgs,
+                specializationArgCount,
+                outSpecializedComponentType,
+                outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL link(
+            slang::IComponentType**         outLinkedComponentType,
+            ISlangBlob**                    outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::link(
+                outLinkedComponentType,
+                outDiagnostics);
+        }
+
             /// Create an entry point that refers to the given function.
         static RefPtr<EntryPoint> create(
             Linkage*            linkage,
@@ -837,6 +892,27 @@ namespace Slang
                 outDiagnostics);
         }
 
+        SLANG_NO_THROW SlangResult SLANG_MCALL link(
+            slang::IComponentType**         outLinkedComponentType,
+            ISlangBlob**                    outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::link(
+                outLinkedComponentType,
+                outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL findEntryPointByName(
+            char const*             name,
+            slang::IEntryPoint**     outEntryPoint) SLANG_OVERRIDE
+        {
+            ComPtr<slang::IEntryPoint> entryPoint(findEntryPointByName(UnownedStringSlice(name)));
+            if((!entryPoint))
+                return SLANG_FAIL;
+
+            *outEntryPoint = entryPoint.detach();
+            return SLANG_OK;
+        }
+
         //
 
             /// Create a module (initially empty).
@@ -912,6 +988,11 @@ namespace Slang
             List<ExpandedSpecializationArg> existentialArgs;
         };
 
+        RefPtr<EntryPoint> findEntryPointByName(UnownedStringSlice const& name);
+
+        List<RefPtr<EntryPoint>> const& getEntryPoints() { return m_entryPoints; }
+        void _addEntryPoint(EntryPoint* entryPoint);
+
     protected:
         void acceptVisitor(ComponentTypeVisitor* visitor, SpecializationInfo* specializationInfo) SLANG_OVERRIDE;
 
@@ -937,6 +1018,26 @@ namespace Slang
 
         // List of filesystem paths this module depends on
         FilePathDependencyList m_filePathDependencyList;
+
+        // Entry points that were defined in thsi module
+        //
+        // Note: the entry point defined in the module are *not*
+        // part of the memory image/layout of the module when
+        // it is considered as an IComponentType. This can be
+        // a bit confusing, but if all the entry points in the
+        // module were automatically linked into the component
+        // type, we'd need a way to access just the global
+        // scope of the module without the entry points, in
+        // case we wanted to link a single entry point against
+        // the global scope. The `Module` type provides exactly
+        // that "module without its entry points" unit of
+        // granularity for linking.
+        //
+        // This list only exists for lookup purposes, so that
+        // the user can find an existing entry-point function
+        // that was defined as part of the module.
+        //
+        List<RefPtr<EntryPoint>> m_entryPoints;
     };
     typedef Module LoadedModule;
 
@@ -963,7 +1064,9 @@ namespace Slang
         void addSourceFile(SourceFile* sourceFile);
 
         // The entry points associated with this translation unit
-        List<RefPtr<EntryPoint>> entryPoints;
+        List<RefPtr<EntryPoint>> const& getEntryPoints() { return module->getEntryPoints(); }
+
+        void _addEntryPoint(EntryPoint* entryPoint) { module->_addEntryPoint(entryPoint); }
 
         // Preprocessor definitions to use for this translation unit only
         // (whereas the ones on `compileRequest` will be shared)
@@ -1258,9 +1361,19 @@ namespace Slang
         // Modules that have been read in with the -r option
         List<RefPtr<IRModule>> m_libModules;
 
+        void _stopRetainingParentSession()
+        {
+            m_retainedSession = nullptr;
+        }
+
     private:
+            /// The global Slang library session that this linkage is a child of
         Session* m_session = nullptr;
 
+        RefPtr<Session> m_retainedSession;
+
+
+
             /// Tracks state of modules currently being loaded.
             ///
             /// This information is used to diagnose cases where
@@ -1975,7 +2088,7 @@ namespace Slang
             /// Get the downstream compiler prelude
         const String& getDownstreamCompilerPrelude(PassThroughMode mode) { return m_downstreamCompilerPreludes[int(mode)]; }
 
-        Session();
+        void init();
 
         void addBuiltinSource(
             RefPtr<Scope> const&    scope,
@@ -2058,10 +2171,7 @@ SLANG_FORCE_INLINE slang::IModule* asExternal(Module* module)
     return static_cast<slang::IModule*>(module);
 }
 
-SLANG_FORCE_INLINE ComponentType* asInternal(slang::IComponentType* componentType)
-{
-    return static_cast<ComponentType*>(componentType);
-}
+ComponentType* asInternal(slang::IComponentType* inComponentType);
 
 SLANG_FORCE_INLINE slang::IComponentType* asExternal(ComponentType* componentType)
 {
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 7cf5215c9..14cc68b16 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -6800,7 +6800,7 @@ IRModule* generateIRForTranslationUnit(
 
     // First, ensure that all entry points have been emitted,
     // in case they require special handling.
-    for (auto entryPoint : translationUnit->entryPoints)
+    for (auto entryPoint : translationUnit->getEntryPoints())
     {
         lowerFrontEndEntryPointToIR(context, entryPoint);
     }
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index fce90d612..bfc77b2e3 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -38,13 +38,15 @@
 namespace Slang {
 
 // Allocate static const storage for the various interface IDs that the Slang API needs to expose
-static const Guid IID_ISlangUnknown = SLANG_UUID_ISlangUnknown;
-static const Guid IID_ISlangBlob    = SLANG_UUID_ISlangBlob;
-static const Guid IID_ISession = SLANG_UUID_ISession;
-static const Guid IID_IGlobalSession = SLANG_UUID_IGlobalSession;
-static const Guid IID_IModule = SLANG_UUID_IModule;
+static const Guid IID_IComponentType    = SLANG_UUID_IComponentType;
+static const Guid IID_IEntryPoint       = SLANG_UUID_IEntryPoint;
+static const Guid IID_IGlobalSession    = SLANG_UUID_IGlobalSession;
+static const Guid IID_IModule           = SLANG_UUID_IModule;
+static const Guid IID_ISession          = SLANG_UUID_ISession;
+static const Guid IID_ISlangBlob        = SLANG_UUID_ISlangBlob;
+static const Guid IID_ISlangUnknown     = SLANG_UUID_ISlangUnknown;
 
-Session::Session()
+void Session::init()
 {
     ::memset(m_downstreamCompilerLocators, 0, sizeof(m_downstreamCompilerLocators));
     DownstreamCompilerUtil::setDefaultLocators(m_downstreamCompilerLocators);
@@ -77,6 +79,16 @@ Session::Session()
 
     m_builtinLinkage = new Linkage(this);
 
+    // Because the `Session` retains the builtin `Linkage`,
+    // we need to make sure that the parent pointer inside
+    // `Linkage` doesn't create a retain cycle.
+    //
+    // This operation ensures that the parent pointer will
+    // just be a raw pointer, so that the builtin linkage
+    // doesn't keep the parent session alive.
+    //
+    m_builtinLinkage->_stopRetainingParentSession();
+
     // Initialize representations of some very basic types:
     initializeTypes();
 
@@ -437,6 +449,7 @@ Profile getEffectiveProfile(EntryPoint* entryPoint, TargetRequest* target)
 
 Linkage::Linkage(Session* session)
     : m_session(session)
+    , m_retainedSession(session)
     , m_sourceManager(&m_defaultSourceManager)
 {
     getNamePool()->setRootNamePool(session->getRootNamePool());
@@ -1704,9 +1717,9 @@ Module::Module(Linkage* linkage)
 
 ISlangUnknown* Module::getInterface(const Guid& guid)
 {
-    if(guid == IID_ISlangUnknown || guid == IID_IModule)
+    if(guid == IID_IModule)
         return asExternal(this);
-    return nullptr;
+    return Super::getInterface(guid);
 }
 
 void Module::addModuleDependency(Module* module)
@@ -1725,14 +1738,51 @@ void Module::setModuleDecl(ModuleDecl* moduleDecl)
     m_moduleDecl = moduleDecl;
 }
 
-// ComponentType
+RefPtr<EntryPoint> Module::findEntryPointByName(UnownedStringSlice const& name)
+{
+    // TODO: We should consider having this function be expanded to be able
+    // to look up and validate possible entry-point functions in teh module
+    // even if they were not marked with `[shader(...)]` in the source code.
+    //
+    // With such a change the function would probably need to accept a stage
+    // to use and a sink to write validation errors to.
+
+    for(auto entryPoint : m_entryPoints)
+    {
+        if(entryPoint->getName()->text.getUnownedSlice() == name)
+            return entryPoint;
+    }
+
+    return nullptr;
+}
+
+void Module::_addEntryPoint(EntryPoint* entryPoint)
+{
+    m_entryPoints.add(entryPoint);
+}
 
-static const Guid IID_IComponentType = SLANG_UUID_IComponentType;
+
+// ComponentType
 
 ComponentType::ComponentType(Linkage* linkage)
     : m_linkage(linkage)
 {}
 
+ComponentType* asInternal(slang::IComponentType* inComponentType)
+{
+    // Note: we use a `queryInterface` here instead of just a `static_cast`
+    // to ensure that the `IComponentType` we get is the preferred/canonical
+    // one, which shares its address with the `ComponentType`.
+    //
+    // TODO: An alternative choice here would be to have a "magic" IID that
+    // we pass into `queryInterface` that returns the `ComponentType` directly
+    // (without even `addRef`-ing it).
+    //
+    ComPtr<slang::IComponentType> componentType;
+    inComponentType->queryInterface(IID_IComponentType, (void**) componentType.writeRef());
+    return static_cast<ComponentType*>(componentType.get());
+}
+
 ISlangUnknown* ComponentType::getInterface(Guid const& guid)
 {
     if(guid == IID_ISlangUnknown
@@ -1873,6 +1923,28 @@ SLANG_NO_THROW SlangResult SLANG_MCALL ComponentType::specialize(
     return SLANG_OK;
 }
 
+RefPtr<ComponentType> fillRequirements(
+    ComponentType* inComponentType);
+
+SLANG_NO_THROW SlangResult SLANG_MCALL ComponentType::link(
+    slang::IComponentType**         outLinkedComponentType,
+    ISlangBlob**                    outDiagnostics)
+{
+    // TODO: It should be possible for `fillRequirements` to fail,
+    // in cases where we have a dependency that can't be automatically
+    // resolved.
+    //
+    SLANG_UNUSED(outDiagnostics);
+
+    auto linked = fillRequirements(this);
+    if(!linked)
+        return SLANG_FAIL;
+
+    *outLinkedComponentType = ComPtr<slang::IComponentType>(linked).detach();
+    return SLANG_OK;
+}
+
+
     /// Visitor used by `ComponentType::enumerateModules`
 struct EnumerateModulesVisitor : ComponentTypeVisitor
 {
@@ -2488,6 +2560,7 @@ Session::~Session()
 SLANG_API SlangSession* spCreateSession(const char*)
 {
     Slang::RefPtr<Slang::Session> session(new Slang::Session());
+    session->init();
     // Will be returned with a refcount of 1
     return asExternal(session.detach());
 }
@@ -2500,6 +2573,7 @@ SLANG_API SlangResult slang_createGlobalSession(
         return SLANG_E_NOT_IMPLEMENTED;
 
     Slang::Session* globalSession = new Slang::Session();
+    globalSession->init();
     Slang::ComPtr<slang::IGlobalSession> result(Slang::asExternal(globalSession));
     *outGlobalSession = result.detach();
     return SLANG_OK;
@@ -3487,6 +3561,20 @@ SLANG_API SlangResult spCompileRequest_getProgram(
     return SLANG_OK;
 }
 
+SLANG_API SlangResult spCompileRequest_getModule(
+    SlangCompileRequest*    request,
+    SlangInt                translationUnitIndex,
+    slang::IModule**        outModule)
+{
+    if( !request ) return SLANG_ERROR_INVALID_PARAMETER;
+    auto req = Slang::asInternal(request);
+
+    auto module = req->getFrontEndReq()->getTranslationUnit(translationUnitIndex)->getModule();
+
+    *outModule = Slang::ComPtr<slang::IModule>(module).detach();
+    return SLANG_OK;
+}
+
 SLANG_API SlangResult spCompileRequest_getEntryPoint(
     SlangCompileRequest*    request,
     SlangInt                entryPointIndex,