Make translation units in the same CompileReq visible to `import`. (#2184)

* Make translation unitts in the same CompileReq visible to `import`.

* Fix code review comments.

Co-authored-by: Yong He <yhe@nvidia.com>
Co-authored-by: Theresa Foley <10618364+tangent-vector@users.noreply.github.com>

10 files changed, 126 insertions, 27 deletions

diff --git a/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj b/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj
index fd24865a4..9da4294c9 100644
--- a/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj
+++ b/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj
@@ -287,6 +287,7 @@
     <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-rtti.cpp" />

     <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-short-list.cpp" />

     <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-string.cpp" />

+    <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-translation-unit-import.cpp" />

     <ClCompile Include="..\..\..\tools\unit-test\slang-unit-test.cpp" />

   </ItemGroup>

   <ItemGroup>

diff --git a/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj.filters b/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj.filters
index 7e06cf50c..5f935e3f7 100644
--- a/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj.filters
+++ b/build/visual-studio/slang-unit-test-tool/slang-unit-test-tool.vcxproj.filters
@@ -62,6 +62,9 @@
     <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-string.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

+    <ClCompile Include="..\..\..\tools\slang-unit-test\unit-test-translation-unit-import.cpp">

+      <Filter>Source Files</Filter>

+    </ClCompile>

     <ClCompile Include="..\..\..\tools\unit-test\slang-unit-test.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index dbbf7e973..8db157a37 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -5033,7 +5033,8 @@ namespace Slang
             getLinkage(),
             name,
             decl->moduleNameAndLoc.loc,
-            getSink());
+            getSink(),
+            m_shared->m_environmentModules);
 
         // If we didn't find a matching module, then bail out
         if (!importedModule)
diff --git a/source/slang/slang-check-impl.h b/source/slang/slang-check-impl.h
index 82a0e7453..f93db894f 100644
--- a/source/slang/slang-check-impl.h
+++ b/source/slang/slang-check-impl.h
@@ -210,6 +210,11 @@ namespace Slang
 
         DiagnosticSink* m_sink      = nullptr;
 
+            /// (optional) modules that comes from previously processed translation units in the
+            /// front-end request that are made visible to the module being checked. This allows
+            /// `import` to use them instead of trying to find the files in file system.
+        LoadedModuleDictionary* m_environmentModules = nullptr;
+
         DiagnosticSink* getSink()
         {
             return m_sink;
@@ -227,10 +232,12 @@ namespace Slang
         SharedSemanticsContext(
             Linkage*        linkage,
             Module*         module,
-            DiagnosticSink* sink)
+            DiagnosticSink* sink,
+            LoadedModuleDictionary* environmentModules = nullptr)
             : m_linkage(linkage)
             , m_module(module)
             , m_sink(sink)
+            , m_environmentModules(environmentModules)
         {}
 
         Session* getSession()
diff --git a/source/slang/slang-check.cpp b/source/slang/slang-check.cpp
index e52baaa84..56e7bd379 100644
--- a/source/slang/slang-check.cpp
+++ b/source/slang/slang-check.cpp
@@ -161,12 +161,14 @@ namespace Slang
     }
 
     void checkTranslationUnit(
-        TranslationUnitRequest* translationUnit)
+        TranslationUnitRequest* translationUnit,
+        LoadedModuleDictionary& loadedModules)
     {
         SharedSemanticsContext sharedSemanticsContext(
             translationUnit->compileRequest->getLinkage(),
             translationUnit->getModule(),
-            translationUnit->compileRequest->getSink());
+            translationUnit->compileRequest->getSink(),
+            &loadedModules);
 
         SemanticsDeclVisitorBase visitor(&sharedSemanticsContext);
 
diff --git a/source/slang/slang-check.h b/source/slang/slang-check.h
index 0309e5393..e84ad0bc5 100644
--- a/source/slang/slang-check.h
+++ b/source/slang/slang-check.h
@@ -18,22 +18,6 @@ namespace Slang
 
     class TranslationUnitRequest;
 
-    void checkTranslationUnit(
-        TranslationUnitRequest* translationUnit);
-
-    // Look for a module that matches the given name:
-    // either one we've loaded already, or one we
-    // can find vai the search paths available to us.
-    //
-    // Needed by import declaration checking.
-    //
-    // TODO: need a better location to declare this.
-    RefPtr<Module> findOrImportModule(
-        Linkage*            linkage,
-        Name*               name,
-        SourceLoc const&    loc,
-        DiagnosticSink*     sink);
-
     bool isGlobalShaderParameter(VarDeclBase* decl);
     bool isFromStdLib(Decl* decl);
 
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 9cd2b0e63..1a0e61aee 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1569,7 +1569,11 @@ namespace Slang
                 Slang::getHashCode(elementType), Slang::getHashCode(containerType));
         }
     };
-    
+
+        /// A dictionary of currently loaded modules. Used by `findOrImportModule` to
+        /// lookup additional loaded modules.
+    typedef Dictionary<Name*, Module*> LoadedModuleDictionary;
+
         /// A context for loading and re-using code modules.
     class Linkage : public RefObject, public slang::ISession
     {
@@ -1764,7 +1768,8 @@ namespace Slang
         RefPtr<Module> findOrImportModule(
             Name*               name,
             SourceLoc const&    loc,
-            DiagnosticSink*     sink);
+            DiagnosticSink*     sink,
+            const LoadedModuleDictionary* loadedModules = nullptr);
 
         SourceManager* getSourceManager()
         {
@@ -2916,6 +2921,21 @@ namespace Slang
         double m_downstreamCompileTime = 0.0;
     };
 
+    void checkTranslationUnit(
+        TranslationUnitRequest* translationUnit, LoadedModuleDictionary& loadedModules);
+
+    // Look for a module that matches the given name:
+    // either one we've loaded already, or one we
+    // can find vai the search paths available to us.
+    //
+    // Needed by import declaration checking.
+    //
+    RefPtr<Module> findOrImportModule(
+        Linkage* linkage,
+        Name* name,
+        SourceLoc const& loc,
+        DiagnosticSink* sink,
+        const LoadedModuleDictionary* additionalLoadedModules);
 
 //
 // The following functions are utilties to convert between
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index aa3f3ac8c..2ddef9e3e 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -2002,11 +2002,21 @@ RefPtr<ComponentType> createSpecializedGlobalAndEntryPointsComponentType(
 
 void FrontEndCompileRequest::checkAllTranslationUnits()
 {
+    LoadedModuleDictionary loadedModules;
+
     // Iterate over all translation units and
     // apply the semantic checking logic.
     for( auto& translationUnit : translationUnits )
     {
-        checkTranslationUnit(translationUnit.Ptr());
+        checkTranslationUnit(translationUnit.Ptr(), loadedModules);
+
+        // Add the checked module to list of loadedModules so that they can be
+        // discovered by `findOrImportModule` when processing future `import` decls.
+        // TODO: this does not handle the case where a translation unit to discover
+        // another translation unit added later to the compilation request.
+        // We should output an error message when we detect such a case, or support
+        // this scenario with a recursive style checking.
+        loadedModules.Add(translationUnit->moduleName, translationUnit->getModule());
     }
     checkEntryPoints();
 }
@@ -2663,7 +2673,8 @@ bool Linkage::isBeingImported(Module* module)
 RefPtr<Module> Linkage::findOrImportModule(
     Name*               name,
     SourceLoc const&    loc,
-    DiagnosticSink*     sink)
+    DiagnosticSink*     sink,
+    const LoadedModuleDictionary*  loadedModules)
 {
     // Have we already loaded a module matching this name?
     //
@@ -2691,6 +2702,16 @@ RefPtr<Module> Linkage::findOrImportModule(
         return loadedModule;
     }
 
+    // If the user is providing an additional list of loaded modules, we find
+    // if the module being imported is in that list. This allows a translation
+    // unit to use previously checked translation units in the same
+    // FrontEndCompileRequest.
+    Module* previouslyLoadedModule = nullptr;
+    if (loadedModules && loadedModules->TryGetValue(name, previouslyLoadedModule))
+    {
+        return previouslyLoadedModule;
+    }
+
     // Derive a file name for the module, by taking the given
     // identifier, replacing all occurrences of `_` with `-`,
     // and then appending `.slang`.
@@ -4019,9 +4040,10 @@ RefPtr<Module> findOrImportModule(
     Linkage*            linkage,
     Name*               name,
     SourceLoc const&    loc,
-    DiagnosticSink*     sink)
+    DiagnosticSink*     sink,
+    const LoadedModuleDictionary* loadedModules)
 {
-    return linkage->findOrImportModule(name, loc, sink);
+    return linkage->findOrImportModule(name, loc, sink, loadedModules);
 }
 
 void Session::addBuiltinSource(
diff --git a/tools/slang-unit-test/unit-test-find-type-by-name.cpp b/tools/slang-unit-test/unit-test-find-type-by-name.cpp
index d12a0d91f..87be156db 100644
--- a/tools/slang-unit-test/unit-test-find-type-by-name.cpp
+++ b/tools/slang-unit-test/unit-test-find-type-by-name.cpp
@@ -1,4 +1,4 @@
-// unit-test-byte-encode.cpp
+// unit-test-find-type-by-name.cpp
 
 #include "../../slang.h"
 
diff --git a/tools/slang-unit-test/unit-test-translation-unit-import.cpp b/tools/slang-unit-test/unit-test-translation-unit-import.cpp
new file mode 100644
index 000000000..dfa75eb9f
--- /dev/null
+++ b/tools/slang-unit-test/unit-test-translation-unit-import.cpp
@@ -0,0 +1,59 @@
+// unit-test-translation-unit-import.cpp
+
+#include "../../slang.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "tools/unit-test/slang-unit-test.h"
+#include "../../slang-com-ptr.h"
+
+using namespace Slang;
+
+// Test that the API supports discovering previously checked translation unit in the same
+// FrontEndCompileRequest.
+SLANG_UNIT_TEST(translationUnitImport)
+{
+    // Source for the first translation unit.
+    const char* generatedSource =
+        "int f() {"
+        "   return 5;"
+        "};";
+
+    // Source for the second translation unit that imports the first translation unit.
+    // The import should succeed and `f` should be visible to this module.
+    const char* userSource =
+        R"(
+        import generatedUnit;
+
+        [shader("compute")]
+        [numthreads(4,1,1)]
+        void computeMain(
+            uint3 sv_dispatchThreadID : SV_DispatchThreadID,
+            uniform RWStructuredBuffer<int> buffer)
+        {
+            buffer[sv_dispatchThreadID.x] = f();
+        }
+        )";    
+    auto session = spCreateSession();
+    auto request = spCreateCompileRequest(session);
+    spAddCodeGenTarget(request, SLANG_HLSL);
+    int generatedTranslationUnitIndex = spAddTranslationUnit(request, SLANG_SOURCE_LANGUAGE_SLANG, "generatedUnit");
+    spAddTranslationUnitSourceString(
+        request, generatedTranslationUnitIndex, "generatedFile", generatedSource);
+    int entryPointTranslationUnitIndex = spAddTranslationUnit(request, SLANG_SOURCE_LANGUAGE_SLANG, "userUnit");
+    spAddTranslationUnitSourceString(
+        request, entryPointTranslationUnitIndex, "userFile", userSource);
+    spAddEntryPoint(request, entryPointTranslationUnitIndex, "computeMain", SLANG_STAGE_COMPUTE);
+
+    auto compileResult = spCompile(request);
+    SLANG_CHECK(compileResult == SLANG_OK);
+
+    Slang::ComPtr<ISlangBlob> outBlob;
+    spGetEntryPointCodeBlob(request, 0, 0, outBlob.writeRef());
+    SLANG_CHECK(outBlob && outBlob->getBufferSize() != 0);
+
+    spDestroyCompileRequest(request);
+    spDestroySession(session);
+}
+