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>

6 files changed, 62 insertions, 26 deletions

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(