Handle import, entrypoint and global params in included files. (#3395)

* Handle `import`, entrypoint and global params in included files.

* Fix language server.

* Extend `_createScopeForLegacyLookup` for `__include`.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

8 files changed, 200 insertions, 186 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 39f0b9096..396fdd297 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -6830,7 +6830,7 @@ namespace Slang
     void SemanticsVisitor::importFileDeclIntoScope(Scope* scope, FileDecl* fileDecl)
     {
         // Create a new sub-scope to wire the module
-       // into our lookup chain.
+        // into our lookup chain.
         auto subScope = getASTBuilder()->create<Scope>();
         subScope->containerDecl = fileDecl;
 
@@ -6882,10 +6882,10 @@ namespace Slang
         // We need to look for a module with the specified name
         // (whether it has already been loaded, or needs to
         // be loaded), and then put its declarations into
-        // the current scope.
+        // the module's scope.
 
         auto name = decl->moduleNameAndLoc.name;
-        auto scope = decl->scope;
+        auto scope = getModuleDecl(decl)->ownedScope;
 
         // Try to load a module matching the name
         auto importedModule = findOrImportModule(
diff --git a/source/slang/slang-check-shader.cpp b/source/slang/slang-check-shader.cpp
index 00c924982..b45107640 100644
--- a/source/slang/slang-check-shader.cpp
+++ b/source/slang/slang-check-shader.cpp
@@ -236,68 +236,74 @@ namespace Slang
         DiagnosticSink*         sink)
     {
         auto translationUnitSyntax = translationUnit->getModuleDecl();
+        FuncDecl* entryPointFuncDecl = nullptr;
 
-        // We will look up any global-scope declarations in the translation
-        // unit that match the name of our entry point.
-        Decl* firstDeclWithName = nullptr;
-        if (!translationUnitSyntax->getMemberDictionary().tryGetValue(name, firstDeclWithName))
+        for (auto globalScope = translationUnit->getModuleDecl()->ownedScope; globalScope; globalScope = globalScope->nextSibling)
         {
-            // If there doesn't appear to be any such declaration, then we are done.
-
-            sink->diagnose(translationUnitSyntax, Diagnostics::entryPointFunctionNotFound, name);
+            if (globalScope->containerDecl != translationUnitSyntax && globalScope->containerDecl->parentDecl != translationUnitSyntax)
+                continue; // Skip scopes that aren't part of the current module.
 
-            return nullptr;
-        }
-
-        // We found at least one global-scope declaration with the right name,
-        // but (1) it might not be a function, and (2) there might be
-        // more than one function.
-        //
-        // We'll walk the linked list of declarations with the same name,
-        // to see what we find. Along the way we'll keep track of the
-        // first function declaration we find, if any:
-        FuncDecl* entryPointFuncDecl = nullptr;
-        for (auto ee = firstDeclWithName; ee; ee = ee->nextInContainerWithSameName)
-        {
-            // Is this declaration a function?
-            if (auto funcDecl = as<FuncDecl>(ee))
+            // We will look up any global-scope declarations in the translation
+            // unit that match the name of our entry point.
+            Decl* firstDeclWithName = nullptr;
+            if (!globalScope->containerDecl->getMemberDictionary().tryGetValue(name, firstDeclWithName))
             {
-                // Skip non-primary declarations, so that
-                // we don't give an error when an entry
-                // point is forward-declared.
-                if (!isPrimaryDecl(funcDecl))
-                    continue;
+                // If there doesn't appear to be any such declaration, then we are done with this scope.
+                continue;
+            }
 
-                // is this the first one we've seen?
-                if (!entryPointFuncDecl)
-                {
-                    // If so, this is a candidate to be
-                    // the entry point function.
-                    entryPointFuncDecl = funcDecl;
-                }
-                else
+            // We found at least one global-scope declaration with the right name,
+            // but (1) it might not be a function, and (2) there might be
+            // more than one function.
+            //
+            // We'll walk the linked list of declarations with the same name,
+            // to see what we find. Along the way we'll keep track of the
+            // first function declaration we find, if any:
+            for (auto ee = firstDeclWithName; ee; ee = ee->nextInContainerWithSameName)
+            {
+                // Is this declaration a function?
+                if (auto funcDecl = as<FuncDecl>(ee))
                 {
-                    // Uh-oh! We've already seen a function declaration with this
-                    // name before, so the whole thing is ambiguous. We need
-                    // to diagnose and bail out.
-
-                    sink->diagnose(translationUnitSyntax, Diagnostics::ambiguousEntryPoint, name);
+                    // Skip non-primary declarations, so that
+                    // we don't give an error when an entry
+                    // point is forward-declared.
+                    if (!isPrimaryDecl(funcDecl))
+                        continue;
 
-                    // List all of the declarations that the user *might* mean
-                    for (auto ff = firstDeclWithName; ff; ff = ff->nextInContainerWithSameName)
+                    // is this the first one we've seen?
+                    if (!entryPointFuncDecl)
                     {
-                        if (auto candidate = as<FuncDecl>(ff))
+                        // If so, this is a candidate to be
+                        // the entry point function.
+                        entryPointFuncDecl = funcDecl;
+                    }
+                    else
+                    {
+                        // Uh-oh! We've already seen a function declaration with this
+                        // name before, so the whole thing is ambiguous. We need
+                        // to diagnose and bail out.
+
+                        sink->diagnose(translationUnitSyntax, Diagnostics::ambiguousEntryPoint, name);
+
+                        // List all of the declarations that the user *might* mean
+                        for (auto ff = firstDeclWithName; ff; ff = ff->nextInContainerWithSameName)
                         {
-                            sink->diagnose(candidate, Diagnostics::entryPointCandidate, candidate->getName());
+                            if (auto candidate = as<FuncDecl>(ff))
+                            {
+                                sink->diagnose(candidate, Diagnostics::entryPointCandidate, candidate->getName());
+                            }
                         }
-                    }
 
-                    // Bail out.
-                    return nullptr;
+                        // Bail out.
+                        return nullptr;
+                    }
                 }
             }
         }
 
+        if (!entryPointFuncDecl)
+            sink->diagnose(translationUnitSyntax, Diagnostics::entryPointFunctionNotFound, name);
+
         return entryPointFuncDecl;
     }
 
@@ -524,86 +530,13 @@ namespace Slang
         auto translationUnit = entryPointReq->getTranslationUnit();
         auto linkage = compileRequest->getLinkage();
         auto sink = compileRequest->getSink();
-        auto translationUnitSyntax = translationUnit->getModuleDecl();
 
         auto entryPointName = entryPointReq->getName();
-
-        // We will look up any global-scope declarations in the translation
-        // unit that match the name of our entry point.
-        Decl* firstDeclWithName = nullptr;
-        if( !translationUnitSyntax->getMemberDictionary().tryGetValue(entryPointName, firstDeclWithName))
-        {
-            // If there doesn't appear to be any such declaration, then
-            // we need to diagnose it as an error, and then bail out.
-            sink->diagnose(translationUnitSyntax, Diagnostics::entryPointFunctionNotFound, entryPointName);
-            return nullptr;
-        }
-
-        // We found at least one global-scope declaration with the right name,
-        // but (1) it might not be a function, and (2) there might be
-        // more than one function.
-        //
-        // We'll walk the linked list of declarations with the same name,
-        // to see what we find. Along the way we'll keep track of the
-        // first function declaration we find, if any:
-        //
-        FuncDecl* entryPointFuncDecl = nullptr;
-        for(auto ee = firstDeclWithName; ee; ee = ee->nextInContainerWithSameName)
-        {
-            // We want to support the case where the declaration is
-            // a generic function, so we will automatically
-            // unwrap any outer `GenericDecl` we find here.
-            //
-            auto decl = ee;
-            if(auto genericDecl = as<GenericDecl>(decl))
-                decl = genericDecl->inner;
-
-            // Is this declaration a function?
-            if (auto funcDecl = as<FuncDecl>(decl))
-            {
-                // Skip non-primary declarations, so that
-                // we don't give an error when an entry
-                // point is forward-declared.
-                if (!isPrimaryDecl(funcDecl))
-                    continue;
-
-                // is this the first one we've seen?
-                if (!entryPointFuncDecl)
-                {
-                    // If so, this is a candidate to be
-                    // the entry point function.
-                    entryPointFuncDecl = funcDecl;
-                }
-                else
-                {
-                    // Uh-oh! We've already seen a function declaration with this
-                    // name before, so the whole thing is ambiguous. We need
-                    // to diagnose and bail out.
-
-                    sink->diagnose(translationUnitSyntax, Diagnostics::ambiguousEntryPoint, entryPointName);
-
-                    // List all of the declarations that the user *might* mean
-                    for (auto ff = firstDeclWithName; ff; ff = ff->nextInContainerWithSameName)
-                    {
-                        if (auto candidate = as<FuncDecl>(ff))
-                        {
-                            sink->diagnose(candidate, Diagnostics::entryPointCandidate, candidate->getName());
-                        }
-                    }
-
-                    // Bail out.
-                    return nullptr;
-                }
-            }
-        }
+        FuncDecl* entryPointFuncDecl = findFunctionDeclByName(translationUnit->getModule(), entryPointName, sink);
 
         // Did we find a function declaration in our search?
         if(!entryPointFuncDecl)
         {
-            // If not, then we need to diagnose the error.
-            // For convenience, we will point to the first
-            // declaration with the right name, that wasn't a function.
-            sink->diagnose(firstDeclWithName, Diagnostics::entryPointSymbolNotAFunction, entryPointName);
             return nullptr;
         }
 
@@ -673,8 +606,6 @@ namespace Slang
 
     void Module::_collectShaderParams()
     {
-        auto moduleDecl = m_moduleDecl;
-
         // We are going to walk the global declarations in the body of the
         // module, and use those to build up our lists of:
         //
@@ -688,70 +619,86 @@ namespace Slang
         // will keep a set of the modules we've already
         // seen and processed.
         //
-        HashSet<Module*> requiredModuleSet;
 
-        for( auto globalDecl : moduleDecl->members )
+        // We need to use a work list to traverse through all global scopes,
+        // including the top level `moduleDecl` and all the included `FileDecl`s.
+
+        List<ContainerDecl*> workList;
+        workList.add(m_moduleDecl);
+
+        HashSet<Module*> requiredModuleSet;
+        for (Index i = 0; i < workList.getCount(); i++)
         {
-            if(auto globalVar = as<VarDecl>(globalDecl))
+            auto moduleDecl = workList[i];
+            for (auto globalDecl : moduleDecl->members)
             {
-                // We do not want to consider global variable declarations
-                // that don't represents shader parameters. This includes
-                // things like `static` globals and `groupshared` variables.
-                //
-                if(!isGlobalShaderParameter(globalVar))
-                    continue;
+                if (auto globalVar = as<VarDecl>(globalDecl))
+                {
+                    // We do not want to consider global variable declarations
+                    // that don't represents shader parameters. This includes
+                    // things like `static` globals and `groupshared` variables.
+                    //
+                    if (!isGlobalShaderParameter(globalVar))
+                        continue;
 
-                // At this point we know we have a global shader parameter.
+                    // At this point we know we have a global shader parameter.
 
-                ShaderParamInfo shaderParamInfo;
-                shaderParamInfo.paramDeclRef = makeDeclRef(globalVar);
+                    ShaderParamInfo shaderParamInfo;
+                    shaderParamInfo.paramDeclRef = makeDeclRef(globalVar);
 
-                // We need to consider what specialization parameters
-                // are introduced by this shader parameter. This step
-                // fills in fields on `shaderParamInfo` so that we
-                // can assocaite specialization arguments supplied later
-                // with the correct parameter.
-                //
-                _collectExistentialSpecializationParamsForShaderParam(
-                    getLinkage()->getASTBuilder(),
-                    shaderParamInfo,
-                    m_specializationParams,
-                    makeDeclRef(globalVar));
+                    // We need to consider what specialization parameters
+                    // are introduced by this shader parameter. This step
+                    // fills in fields on `shaderParamInfo` so that we
+                    // can assocaite specialization arguments supplied later
+                    // with the correct parameter.
+                    //
+                    _collectExistentialSpecializationParamsForShaderParam(
+                        getLinkage()->getASTBuilder(),
+                        shaderParamInfo,
+                        m_specializationParams,
+                        makeDeclRef(globalVar));
 
-                m_shaderParams.add(shaderParamInfo);
-            }
-            else if( auto globalGenericParam = as<GlobalGenericParamDecl>(globalDecl) )
-            {
-                // A global generic type parameter declaration introduces
-                // a suitable specialization parameter.
-                //
-                SpecializationParam specializationParam;
-                specializationParam.flavor = SpecializationParam::Flavor::GenericType;
-                specializationParam.loc = globalGenericParam->loc;
-                specializationParam.object = globalGenericParam;
-                m_specializationParams.add(specializationParam);
-            }
-            else if( auto globalGenericValueParam = as<GlobalGenericValueParamDecl>(globalDecl) )
-            {
-                // A global generic type parameter declaration introduces
-                // a suitable specialization parameter.
-                //
-                SpecializationParam specializationParam;
-                specializationParam.flavor = SpecializationParam::Flavor::GenericValue;
-                specializationParam.loc = globalGenericValueParam->loc;
-                specializationParam.object = globalGenericValueParam;
-                m_specializationParams.add(specializationParam);
-            }
-            else if( auto importDecl = as<ImportDecl>(globalDecl) )
-            {
-                // An `import` declaration creates a requirement dependency
-                // from this module to another module.
-                //
-                auto importedModule = getModule(importDecl->importedModuleDecl);
-                if(!requiredModuleSet.contains(importedModule))
+                    m_shaderParams.add(shaderParamInfo);
+                }
+                else if (auto globalGenericParam = as<GlobalGenericParamDecl>(globalDecl))
                 {
-                    requiredModuleSet.add(importedModule);
-                    m_requirements.add(importedModule);
+                    // A global generic type parameter declaration introduces
+                    // a suitable specialization parameter.
+                    //
+                    SpecializationParam specializationParam;
+                    specializationParam.flavor = SpecializationParam::Flavor::GenericType;
+                    specializationParam.loc = globalGenericParam->loc;
+                    specializationParam.object = globalGenericParam;
+                    m_specializationParams.add(specializationParam);
+                }
+                else if (auto globalGenericValueParam = as<GlobalGenericValueParamDecl>(globalDecl))
+                {
+                    // A global generic type parameter declaration introduces
+                    // a suitable specialization parameter.
+                    //
+                    SpecializationParam specializationParam;
+                    specializationParam.flavor = SpecializationParam::Flavor::GenericValue;
+                    specializationParam.loc = globalGenericValueParam->loc;
+                    specializationParam.object = globalGenericValueParam;
+                    m_specializationParams.add(specializationParam);
+                }
+                else if (auto importDecl = as<ImportDecl>(globalDecl))
+                {
+                    // An `import` declaration creates a requirement dependency
+                    // from this module to another module.
+                    //
+                    auto importedModule = getModule(importDecl->importedModuleDecl);
+                    if (!requiredModuleSet.contains(importedModule))
+                    {
+                        requiredModuleSet.add(importedModule);
+                        m_requirements.add(importedModule);
+                    }
+                }
+                else if (auto fileDecl = as<FileDecl>(globalDecl))
+                {
+                    // If we see a `FileDecl`, we need to recursively look into its
+                    // scope.
+                    workList.add(fileDecl);
                 }
             }
         }
@@ -1392,11 +1339,17 @@ namespace Slang
         //
         for( auto module : getModuleDependencies() )
         {
-            Scope* moduleScope = astBuilder->create<Scope>();
-            moduleScope->containerDecl = module->getModuleDecl();
+            for (auto srcScope = module->getModuleDecl()->ownedScope; srcScope; srcScope = srcScope->nextSibling)
+            {
+                if (srcScope->containerDecl != module->getModuleDecl() && srcScope->containerDecl->parentDecl != module->getModuleDecl())
+                    continue; // Skip scopes that is not part of current module.
+
+                Scope* moduleScope = astBuilder->create<Scope>();
+                moduleScope->containerDecl = srcScope->containerDecl;
 
-            moduleScope->nextSibling = scope->nextSibling;
-            scope->nextSibling = moduleScope;
+                moduleScope->nextSibling = scope->nextSibling;
+                scope->nextSibling = moduleScope;
+            }
         }
 
         return scope;
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 8f94b7e90..f0fd33e13 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -7304,6 +7304,11 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             // Variable declared at global/namespace scope? -> Global.
             return true;
         }
+        else if (as<FileDecl>(parent))
+        {
+            // Variable declared at file scope? -> Global.
+            return true;
+        }
         else if(as<AggTypeDeclBase>(parent))
         {
             if(decl->hasModifier<HLSLStaticModifier>())
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 2852dd43e..744033a0f 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -3367,6 +3367,25 @@ Linkage::IncludeResult Linkage::findAndIncludeFile(Module* module, TranslationUn
         return result;
     }
 
+    if (isInLanguageServer())
+    {
+        // HACK: When in language server mode, we will always load the currently opend file as a fresh module
+        // even if some previously opened file already references the current file via `import` or `include`.
+        // see comments in `WorkspaceVersion::getOrLoadModule()` for the reason behind this.
+        // An undesired outcome of this decision is that we could endup including the currently opened file itself
+        // via chain of `__include`s because the currently opened file will not have a true unique file system
+        // identity that allows it to be deduplicated correct. Therefore we insert a hack logic here to detect
+        // re-inclusion by just the file path.
+        // We can clean up this hack by making the language server truly support incremental checking so we can
+        // reuse the previously loaded module instead of needing to always start with a fresh copy.
+        //
+        for (auto file : translationUnit->getSourceFiles())
+        {
+            if (file->getPathInfo().hasFoundPath() && Path::equals(file->getPathInfo().foundPath, sourceFile->getPathInfo().foundPath))
+                return result;
+        }
+    }
+
     module->addFileDependency(sourceFile);
 
     // Create a transparent FileDecl to hold all children from the included file.
diff --git a/tests/language-feature/modules/import-in-include/a.slang b/tests/language-feature/modules/import-in-include/a.slang
new file mode 100644
index 000000000..e8bcd3ae8
--- /dev/null
+++ b/tests/language-feature/modules/import-in-include/a.slang
@@ -0,0 +1,3 @@
+implementing c;
+
+import helper;
diff --git a/tests/language-feature/modules/import-in-include/b.slang b/tests/language-feature/modules/import-in-include/b.slang
new file mode 100644
index 000000000..6d8ba30d5
--- /dev/null
+++ b/tests/language-feature/modules/import-in-include/b.slang
@@ -0,0 +1,14 @@
+implementing c;
+
+int f()
+{
+    return helperFunc();
+}
+
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    outputBuffer[0] = f();
+}
diff --git a/tests/language-feature/modules/import-in-include/c.slang b/tests/language-feature/modules/import-in-include/c.slang
new file mode 100644
index 000000000..5ff5e7ed6
--- /dev/null
+++ b/tests/language-feature/modules/import-in-include/c.slang
@@ -0,0 +1,14 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -vk -shaderobj -output-using-type
+
+// Test that a module imported from an __include'd file can be used from other files
+// in the same module, and that entry points and global parameters defined in an
+// included file can be correctly discovered by the reflection API.
+
+module c;
+
+__include b; // uses helper module.
+__include a; // imports helper module.
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+//CHECK: 1
diff --git a/tests/language-feature/modules/import-in-include/helper.slang b/tests/language-feature/modules/import-in-include/helper.slang
new file mode 100644
index 000000000..92c1bfb96
--- /dev/null
+++ b/tests/language-feature/modules/import-in-include/helper.slang
@@ -0,0 +1,6 @@
+module helper;
+
+public int helperFunc()
+{
+    return 1;
+}