Fixups for IR checkpoint

- The changes introduced a new path where we don't even go through the current "lowering" (really an AST-to-AST legalization pass), but this exposed a few issues I didn't anticipate:

- First, we needed to make sure to pass in the computed layout information when emitting the original program (since the layout info is no longer automatically attached to AST nodes)

- Second, we needed to take the sample-rate input checks that were being done in lowering before, and move them to the emit logic (which is really ugly, but I don't see a way around it for GLSL).

2 files changed, 74 insertions, 40 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 3bec17456..c8f3f06cb 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -48,6 +48,9 @@ struct SharedEmitContext
     // The entry point we are being asked to compile
     EntryPointRequest* entryPoint;
 
+    // The layout for the entry point
+    EntryPointLayout*   entryPointLayout;
+
     // The target language we want to generate code for
     CodeGenTarget target;
 
@@ -2283,8 +2286,32 @@ struct EmitVisitor
         emitName(expr->lookupResult2.getName());
     }
 
+    void setSampleRateFlag()
+    {
+        context->shared->entryPointLayout->flags |= EntryPointLayout::Flag::usesAnySampleRateInput;
+    }
+
+    void doSampleRateInputCheck(VarDeclBase* decl)
+    {
+        if (decl->HasModifier<HLSLSampleModifier>())
+        {
+            setSampleRateFlag();
+        }
+    }
+
+    void doSampleRateInputCheck(Name* name)
+    {
+        auto text = getText(name);
+        if (text == "gl_SampleID")
+        {
+            setSampleRateFlag();
+        }
+    }
+
     void visitVarExpr(VarExpr* varExpr, ExprEmitArg const& arg)
     {
+        doSampleRateInputCheck(varExpr->name);
+
         auto prec = kEOp_Atomic;
         auto outerPrec = arg.outerPrec;
         bool needClose = MaybeEmitParens(outerPrec, kEOp_Atomic);
@@ -2485,6 +2512,11 @@ struct EmitVisitor
         Emit("{\n");
         for( auto& token : stmt->tokens )
         {
+            if (token.type == TokenType::Identifier)
+            {
+                doSampleRateInputCheck(token.getName());
+            }
+
             emitTokenWithLocation(token);
         }
         Emit("}\n");
@@ -3560,6 +3592,15 @@ struct EmitVisitor
 
     void visitVarDeclBase(RefPtr<VarDeclBase> decl, DeclEmitArg const& arg)
     {
+        // Global variable? Check if it is a sample-rate input.
+        if (dynamic_cast<ModuleDecl*>(decl->ParentDecl))
+        {
+            if (decl->HasModifier<InModifier>())
+            {
+                doSampleRateInputCheck(decl);
+            }
+        }
+
         // Skip fields that have been tuple-ified and don't contribute
         // any fields of "ordinary" type.
         if (auto tupleFieldMod = decl->FindModifier<TupleFieldModifier>())
@@ -3783,6 +3824,29 @@ emitDeclImpl(decl, nullptr);
     }
 };
 
+
+EntryPointLayout* findEntryPointLayout(
+    ProgramLayout*      programLayout,
+    EntryPointRequest*  entryPointRequest)
+{
+    for( auto entryPointLayout : programLayout->entryPoints )
+    {
+        if(entryPointLayout->entryPoint->getName() != entryPointRequest->name)
+            continue;
+
+        if(entryPointLayout->profile != entryPointRequest->profile)
+            continue;
+
+        // TODO: can't easily filter on translation unit here...
+        // Ideally the `EntryPointRequest` should get filled in with a pointer
+        // the specific function declaration that represents the entry point.
+
+        return entryPointLayout.Ptr();
+    }
+
+    return nullptr;
+}
+
 String emitEntryPoint(
     EntryPointRequest*  entryPoint,
     ProgramLayout*      programLayout,
@@ -3795,6 +3859,13 @@ String emitEntryPoint(
     sharedContext.finalTarget = entryPoint->compileRequest->Target;
     sharedContext.entryPoint = entryPoint;
 
+    if (entryPoint)
+    {
+        sharedContext.entryPointLayout = findEntryPointLayout(
+            programLayout,
+            entryPoint);
+    }
+
     sharedContext.programLayout = programLayout;
 
     // Layout information for the global scope is either an ordinary
@@ -3916,7 +3987,9 @@ String emitEntryPoint(
         // along with whatever annotations we added along the way.
 
         sharedContext.program = translationUnitSyntax;
-        visitor.EmitDeclsInContainer(translationUnitSyntax);
+        visitor.EmitDeclsInContainerUsingLayout(
+            translationUnitSyntax,
+            globalStructLayout);
     }
 
     String code = sharedContext.sb.ProduceString();
diff --git a/source/slang/lower.cpp b/source/slang/lower.cpp
index 2c62a69a8..7a12fc3b7 100644
--- a/source/slang/lower.cpp
+++ b/source/slang/lower.cpp
@@ -886,8 +886,6 @@ struct LoweringVisitor
     LoweredExpr visitVarExpr(
         VarExpr* expr)
     {
-        doSampleRateInputCheck(expr->name);
-
         // If the expression didn't get resolved, we can leave it as-is
         if (!expr->declRef)
             return expr;
@@ -2206,14 +2204,6 @@ struct LoweringVisitor
         RefPtr<UnparsedStmt> loweredStmt = new UnparsedStmt();
         lowerStmtFields(loweredStmt, stmt);
 
-        for (auto token : stmt->tokens)
-        {
-            if (token.type == TokenType::Identifier)
-            {
-                doSampleRateInputCheck(token.getName());
-            }
-        }
-
         loweredStmt->tokens = stmt->tokens;
 
         addStmt(loweredStmt);
@@ -3369,28 +3359,6 @@ struct LoweringVisitor
         return SourceLanguage::Unknown;
     }
 
-    void setSampleRateFlag()
-    {
-        shared->entryPointLayout->flags |= EntryPointLayout::Flag::usesAnySampleRateInput;
-    }
-
-    void doSampleRateInputCheck(VarDeclBase* decl)
-    {
-        if (decl->HasModifier<HLSLSampleModifier>())
-        {
-            setSampleRateFlag();
-        }
-    }
-
-    void doSampleRateInputCheck(Name* name)
-    {
-        auto text = getText(name);
-        if (text == "gl_SampleIndex")
-        {
-            setSampleRateFlag();
-        }
-    }
-
     AggTypeDecl* isStructType(RefPtr<Type> type)
     {
         if (type->As<BasicExpressionType>()) return nullptr;
@@ -3440,14 +3408,8 @@ struct LoweringVisitor
     LoweredDecl visitVariable(
         Variable* decl)
     {
-        // Global variable? Check if it is a sample-rate input.
         if (dynamic_cast<ModuleDecl*>(decl->ParentDecl))
         {
-            if (decl->HasModifier<InModifier>())
-            {
-                doSampleRateInputCheck(decl);
-            }
-
             auto varLayout = tryToFindLayout(decl);
             if (varLayout)
             {
@@ -3902,7 +3864,6 @@ struct LoweringVisitor
             }
             else if (ns == "sv_sampleindex")
             {
-                setSampleRateFlag();
                 globalVarExpr = createGLSLBuiltinRef("gl_SampleID", getIntType());
             }
             else if (ns == "sv_stencilref")