API support for IR modules - setting references to libs, and returns lib from compilation (#1105)

* Added RiffReadHelper

* Move type to fourCC in Chunk simplifies some code.

* Make MemoryArena able to track external blocks.
Allow ownership of Data to vary.
Changed IR serialization to use moved allocations to avoid copies.

As it turns out all of the array writes could use unowned data, but doing so requires the IRData to stay in scope longer than IRSerialData, which it does at the moment - but perhaps needs better naming or a control for the feature.

* Write out slang-module container.

* WIP on -r option.
Loading modules - with -r.

* Making the serialized-module run (without using imported module).

* Split compiling module from the test.

* Separate module compilation with a function working.

* Remove serialization test as not used.

* Fix warning on gcc.

* Updated test to have types across module boundary.

* Allow entry point declaration.
A test that tries to build with just an entry point declaration and a module.

* Try to make link work with multiple modules.

* Multi module linking first pass working.

* Multi module test working with -module-name option

* Added feature to repro manifest of approximation of command line that was used.

* Use isDefinition - for determining to add decorations to entry point lowering.

* Added support for repo-file-system.h
More precise control of CacheFileSystem.
Allow RelativeFileSystem to strip paths optionally.
Use canonical paths in PathInfo cache.
Fix bug in -D options for command line output of StateSerailizeUtil

* Add missing slang-options.h

* Fix bug in bit slang-state-serialize.cpp with bit removal.

* Added documentation around -repro-file-system
Added spLoadReproAsFileSystem function.

* Fix warning.

* spAddLibraryReference

* * Add support for slang-lib extension
* Container output when using -no-codegen option

* Use the m_containerFormat to determine if the module container is constructed.
Store the result in a blob. This allows for potential access via the API.
Write the blob if a filename is set.
Use m_ prefix for container variables.

* Added spGetContainerCode.
Made spGetCompileRequestCode work.

6 files changed, 139 insertions, 39 deletions

diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index 92759ed0f..74075bb51 100644
--- a/source/slang/slang-compiler.cpp
+++ b/source/slang/slang-compiler.cpp
@@ -75,6 +75,7 @@
 
 namespace Slang
 {
+
 #define SLANG_CODE_GEN_TARGETS(x) \
     x("unknown", Unknown) \
     x("none", None) \
@@ -2198,9 +2199,8 @@ SlangResult dissassembleDXILUsingDXC(
         }
     }
 
-    static SlangResult _writeContainerFile(
-        EndToEndCompileRequest* endToEndReq,
-        Stream* stream)
+    
+    SlangResult EndToEndCompileRequest::writeContainerToStream(Stream* stream)
     {
         RiffContainer container;
 
@@ -2212,9 +2212,9 @@ SlangResult dissassembleDXILUsingDXC(
             // Module list
             RiffContainer::ScopeChunk listScope(&container, RiffContainer::Chunk::Kind::List, IRSerialBinary::kSlangModuleListFourCc);
 
-            auto linkage = endToEndReq->getLinkage();
-            auto sink = endToEndReq->getSink();
-            auto frontEndReq = endToEndReq->getFrontEndReq();
+            auto linkage = getLinkage();
+            auto sink = getSink();
+            auto frontEndReq = getFrontEndReq();
 
             IRSerialWriter::OptionFlags optionFlags = 0;
 
@@ -2239,7 +2239,7 @@ SlangResult dissassembleDXILUsingDXC(
                 SLANG_RETURN_ON_FAIL(IRSerialWriter::writeContainer(serialData, compressionType, &container));
             }
 
-            auto program = endToEndReq->getSpecializedGlobalAndEntryPointsComponentType();
+            auto program = getSpecializedGlobalAndEntryPointsComponentType();
 
             // TODO: in the case where we have specialization, we might need
             // to serialize IR related to `program`...
@@ -2263,20 +2263,57 @@ SlangResult dissassembleDXILUsingDXC(
         return SLANG_OK;
     }
 
-        /// Write out a "container" file with the stuff that has
-        /// been compiled as part of this request.
-        ///
-    static void _writeContainerFile(
-        EndToEndCompileRequest* endToEndReq,
-        const String& fileName)
+    SlangResult EndToEndCompileRequest::maybeCreateContainer()
+    {
+        switch (m_containerFormat)
+        {
+            case ContainerFormat::SlangModule:
+            {
+                m_containerBlob.setNull();
+
+                OwnedMemoryStream stream(FileAccess::Write);
+                SlangResult res = writeContainerToStream(&stream);
+                if (SLANG_FAILED(res))
+                {
+                    getSink()->diagnose(SourceLoc(), Diagnostics::unableToCreateModuleContainer);
+                    return res;
+                }
+
+                // Need to turn into a blob
+                RefPtr<ListBlob> blob(new ListBlob);
+                // Swap the streams contents into the blob
+                stream.swapContents(blob->m_data);
+                m_containerBlob = blob;
+
+                return res;
+            }
+            default: break;
+        }
+        return SLANG_OK;
+    }
+
+    SlangResult EndToEndCompileRequest::maybeWriteContainer(const String& fileName)
     {
+        // If there is no container, or filename, don't write anything 
+        if (fileName.getLength() == 0 || !m_containerBlob)
+        {
+            return SLANG_OK;
+        }
+
         FileStream stream(fileName, FileMode::Create, FileAccess::Write, FileShare::ReadWrite);
-        if (SLANG_FAILED(_writeContainerFile(endToEndReq, &stream)))
+        try
         {
-            endToEndReq->getSink()->diagnose(SourceLoc(), Diagnostics::unableToWriteModuleContainer, fileName);
+            stream.write(m_containerBlob->getBufferPointer(), m_containerBlob->getBufferSize());
         }
+        catch (IOException&)
+        {
+            // Unable to write
+            return SLANG_FAIL;
+        }
+        return SLANG_OK;
     }
 
+
     static void _generateOutput(
         BackEndCompileRequest* compileRequest,
         EndToEndCompileRequest* endToEndReq)
@@ -2322,10 +2359,8 @@ SlangResult dissassembleDXILUsingDXC(
                 }
             }
 
-            if (compileRequest->containerOutputPath.getLength() != 0)
-            {
-                _writeContainerFile(compileRequest, compileRequest->containerOutputPath);
-            }
+            compileRequest->maybeCreateContainer();
+            compileRequest->maybeWriteContainer(compileRequest->m_containerOutputPath);
         }
     }
 
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index dd90ff7c6..860195a90 100644
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1699,20 +1699,16 @@ namespace Slang
         EndToEndCompileRequest(
             Linkage* linkage);
 
-        // What container format are we being asked to generate?
-        //
-        // Note: This field is unused except by the options-parsing
-        // logic; it exists to support wriiting out binary modules
-        // once that feature is ready.
-        //
-        ContainerFormat containerFormat = ContainerFormat::None;
+            // What container format are we being asked to generate?
+            // If it's set to a format, the container blob will be calculated during compile
+        ContainerFormat m_containerFormat = ContainerFormat::None;
 
-        // Path to output container to
-        //
-        // Note: This field exists to support wriiting out binary modules
-        // once that feature is ready.
-        //
-        String containerOutputPath;
+            /// Where the container blob is stored. This is calculated as part of compile if m_containerFormat is set to
+            /// a supported format. 
+        ComPtr<ISlangBlob> m_containerBlob;
+
+            // Path to output container to
+        String m_containerOutputPath;
 
         // Should we just pass the input to another compiler?
         PassThroughMode passThrough = PassThroughMode::None;
@@ -1757,6 +1753,15 @@ namespace Slang
         };
         Dictionary<TargetRequest*, RefPtr<TargetInfo>> targetInfos;
 
+            /// Writes the modules in a container to the stream
+        SlangResult writeContainerToStream(Stream* stream);
+        
+            /// If a container format has been specified produce a container (stored in m_containerBlob)
+        SlangResult maybeCreateContainer();
+            /// If a container has been constructed and the filename/path has contents will try to write
+            /// the container contents to the file
+        SlangResult maybeWriteContainer(const String& fileName);
+
         Linkage* getLinkage() { return m_linkage; }
 
         int addEntryPoint(
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 405590111..485f463e8 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -112,6 +112,9 @@ DIAGNOSTIC(    81, Error, parametersAfterLoadReproIgnored, "parameters after -lo
 DIAGNOSTIC(    82, Error, unableToWriteReproFile, "unable to write repro file '%0'");
 DIAGNOSTIC(    83, Error, unableToWriteModuleContainer, "unable to write module container '%0'");
 DIAGNOSTIC(    84, Error, unableToReadModuleContainer, "unable to read module container '%0'");
+DIAGNOSTIC(    85, Error, unableToAddReferenceToModuleContainer, "unable to add a reference to a module container");
+DIAGNOSTIC(    86, Error, unableToCreateModuleContainer, "unable to create module container");
+
 
 //
 // 1xxxx - Lexical anaylsis
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index ad1d27114..89e5b8c14 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -398,10 +398,10 @@ struct OptionsParser
 
 #undef CASE
 
-        else if (path.endsWith(".slang-module"))
+        else if (path.endsWith(".slang-module") || path.endsWith(".slang-lib"))
         {
             spSetOutputContainerFormat(compileRequest, SLANG_CONTAINER_FORMAT_SLANG_MODULE);
-            requestImpl->containerOutputPath = path;
+            requestImpl->m_containerOutputPath = path;
         }
         else
         {
diff --git a/source/slang/slang-state-serialize.cpp b/source/slang/slang-state-serialize.cpp
index 369c433ca..f8baa23bb 100644
--- a/source/slang/slang-state-serialize.cpp
+++ b/source/slang/slang-state-serialize.cpp
@@ -350,7 +350,7 @@ static bool _isStorable(const PathInfo::Type type)
 
         dst->debugInfoLevel = linkage->debugInfoLevel;
         dst->optimizationLevel = linkage->optimizationLevel;
-        dst->containerFormat = request->containerFormat;
+        dst->containerFormat = request->m_containerFormat;
         dst->passThroughMode = request->passThrough;
 
 
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 946d7337b..ab1b0e489 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -1271,6 +1271,9 @@ SlangResult EndToEndCompileRequest::executeActionsInner()
         m_specializedGlobalAndEntryPointsComponentType = getUnspecializedGlobalAndEntryPointsComponentType();
         m_specializedEntryPoints = getFrontEndReq()->getUnspecializedEntryPoints();
 
+        SLANG_RETURN_ON_FAIL(maybeCreateContainer());
+        SLANG_RETURN_ON_FAIL(maybeWriteContainer(m_containerOutputPath));
+
         return SLANG_OK;
     }
 
@@ -2913,7 +2916,7 @@ SLANG_API void spSetOutputContainerFormat(
     SlangContainerFormat    format)
 {
     auto req = Slang::asInternal(request);
-    req->containerFormat = Slang::ContainerFormat(format);
+    req->m_containerFormat = Slang::ContainerFormat(format);
 }
 
 
@@ -3039,6 +3042,33 @@ SLANG_API void spSetDefaultModuleName(
 }
 
 
+SLANG_API SlangResult spAddLibraryReference(
+    SlangCompileRequest*    request,
+    const void* libData,
+    size_t libDataSize)
+{
+    using namespace Slang;
+    auto req = Slang::asInternal(request);
+
+    // We need to deserialize and add the modules
+    MemoryStreamBase fileStream(FileAccess::Read, libData, libDataSize);
+
+    // Read all of the contained modules
+    List<RefPtr<IRModule>> irModules;
+    if (SLANG_FAILED(IRSerialReader::readStreamModules(&fileStream, req->getSession(), req->getFrontEndReq()->getSourceManager(), irModules)))
+    {
+        req->getSink()->diagnose(SourceLoc(), Diagnostics::unableToAddReferenceToModuleContainer);
+        return SLANG_FAIL;
+    }
+
+    // TODO(JS): May be better to have a ITypeComponent that encapsulates a collection of modules
+    // For now just add to the linkage
+    auto linkage = req->getLinkage();
+    linkage->m_libModules.addRange(irModules.getBuffer(), irModules.getCount());
+
+    return SLANG_OK;
+}
+
 SLANG_API void spTranslationUnit_addPreprocessorDefine(
     SlangCompileRequest*    request,
     int                     translationUnitIndex,
@@ -3497,14 +3527,41 @@ SLANG_API char const* spGetEntryPointSource(
 }
 
 SLANG_API void const* spGetCompileRequestCode(
-    SlangCompileRequest*    request,
+    SlangCompileRequest*    inRequest,
     size_t*                 outSize)
 {
-    SLANG_UNUSED(request);
-    SLANG_UNUSED(outSize);
+    using namespace Slang;
+    auto request = asInternal(inRequest);
+
+    if (request->m_containerBlob)
+    {
+        *outSize = request->m_containerBlob->getBufferSize();
+        return request->m_containerBlob->getBufferPointer();
+    }
+
+    // Container blob does not have any contents
+    *outSize = 0;
     return nullptr;
 }
 
+SLANG_API SlangResult spGetContainerCode(
+    SlangCompileRequest*    inRequest,
+    ISlangBlob**            outBlob)
+{
+    using namespace Slang;
+    auto request = asInternal(inRequest);
+
+    ISlangBlob* containerBlob = request->m_containerBlob;
+    if (containerBlob)
+    {
+        containerBlob->addRef();
+        *outBlob = containerBlob;
+        return SLANG_OK;
+    }
+
+    return SLANG_FAIL;
+}
+
 SLANG_API SlangResult spLoadRepro(
     SlangCompileRequest* inRequest,
     ISlangFileSystem* fileSystem,