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>

1 files changed, 26 insertions, 4 deletions

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(