Fix __init() functions that returns an existing value (#3866)

Fixes the issue #3671

* The __init constructors are not expected to return a value like other member
  functions, but must construct a new value and return the struct type or none.

* This patch enables this behavior in the IR lowering without complaining about
  illegal situations where the user returns an invalid type or none at all.
   Translate    ordinary struct `return ...;` to `this = ...; return this;`
   Translate NonCopyableType struct `return ...;` to `return this;`

* This patch also fixes the issue with type checking when __init()
  returns a void that mismatches the base type of the struct/ class
   Translate    ordinary struct `return;` to `return this;`
   Translate NonCopyableType struct `return;` to `return;`

* Add end-to-end test and compile only tests to check the above behavior.

6 files changed, 366 insertions, 3 deletions

diff --git a/source/slang/slang-check-stmt.cpp b/source/slang/slang-check-stmt.cpp
index 729b24d35..f7b085579 100644
--- a/source/slang/slang-check-stmt.cpp
+++ b/source/slang/slang-check-stmt.cpp
@@ -392,7 +392,7 @@ namespace Slang
         auto function = getParentFunc();
         if (!stmt->expression)
         {
-            if (function && !function->returnType.equals(m_astBuilder->getVoidType()))
+            if (function && !function->returnType.equals(m_astBuilder->getVoidType()) && !as<ConstructorDecl>(function))
             {
                 getSink()->diagnose(stmt, Diagnostics::returnNeedsExpression);
             }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 6a6ca7c64..3f776c1a1 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -586,6 +586,10 @@ struct IRGenContext
     // (For use by functions that returns non-copyable types)
     LoweredValInfo returnDestination;
 
+    // A reference to the Function decl to identify the parent function
+    // that contains the Inst.
+    FunctionDeclBase* funcDecl;
+
     bool includeDebugInfo = false;
 
     explicit IRGenContext(SharedIRGenContext* inShared, ASTBuilder* inAstBuilder)
@@ -6012,12 +6016,15 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
     {
         startBlockIfNeeded(stmt);
 
+        // Check if this return is within a constructor.
+        auto constructorDecl = as<ConstructorDecl>(context->funcDecl);
+
         // A `return` statement turns into a `return` instruction,
         // but we have two kinds of `return`: one for returning
         // a (non-`void`) value, and one for returning "no value"
         // (which effectively returns a value of type `void`).
         //
-        if( auto expr = stmt->expression )
+        if (auto expr = stmt->expression)
         {
             if (context->returnDestination.flavor != LoweredValInfo::Flavor::None)
             {
@@ -6027,6 +6034,16 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
                 return;
             }
 
+            if (constructorDecl)
+            {
+                // If this function is a constructor, but returns a value, rewrite it as
+                // this = val;
+                // return this;
+                lowerRValueExprWithDestination(context, context->thisVal, expr);
+                getBuilder()->emitReturn(getSimpleVal(context, context->thisVal));
+                return;
+            }
+
             // If the AST `return` statement had an expression, then we
             // need to lower it to the IR at this point, both to
             // compute its value and (in case we are returning a
@@ -6061,6 +6078,24 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
             // with no value, which can only occur in a function with
             // a `void` result type.
             //
+            if (constructorDecl)
+            {
+                // If this `return` is within a NonCopyableType or an ordinary constructor,
+                // then we must either simply return or `return` the instance respectively.
+                if (context->returnDestination.flavor != LoweredValInfo::Flavor::None)
+                {
+                    // If we have a NonCopyableType constructor of the form
+                    //   void ctor(inout this) { return; }
+                    getBuilder()->emitReturn();
+                }
+                else
+                {
+                    // If we have an ordinary constructor of the form
+                    //   Type ctor() { return; }
+                    getBuilder()->emitReturn(getSimpleVal(context, context->thisVal));
+                }
+                return;
+            }
             getBuilder()->emitReturn();
         }
     }
@@ -9298,7 +9333,8 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
     {
 
         IRGeneric* outerGeneric = nullptr;
-        
+        subContext->funcDecl = decl;
+
         if (auto derivativeRequirement = as<DerivativeRequirementDecl>(decl))
             outerGeneric = emitOuterGenerics(subContext, derivativeRequirement->originalRequirementDecl, derivativeRequirement->originalRequirementDecl);
         else
diff --git a/tests/diagnostics/ctor-noncopyable-retval-legal.slang b/tests/diagnostics/ctor-noncopyable-retval-legal.slang
new file mode 100644
index 000000000..d4cfc5030
--- /dev/null
+++ b/tests/diagnostics/ctor-noncopyable-retval-legal.slang
@@ -0,0 +1,103 @@
+//TEST:SIMPLE(filecheck=CHECK):
+
+// Test to catch compilation time errors with return values in constructors within a NonCopyable struct.
+// A constructor from the callsite's point of view is a function that returns the struct type.
+// A constructor from the inside the body is treated as a function that modifies `this` and not return.
+
+// In the following test-cases where the constructor either returns an expr or no value,
+// within a condition for early exit or not, it automatically gets resolved as
+// return *this; or return;
+
+[__NonCopyableType] struct S
+{
+    float v;    
+}
+
+[__NonCopyableType] struct S1 : S
+{
+    __init()
+    {
+        this.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return this;
+    }
+}
+
+[__NonCopyableType] struct S4 : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            this.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return this;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+[__NonCopyableType] struct S2 : S
+{
+    __init()
+    {
+        S2 t;
+        t.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+[__NonCopyableType] struct S5 : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            S5 t;
+            t.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+[__NonCopyableType] struct S3 : S
+{
+    __init()
+    {
+        S3 t;
+        t.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return t;
+    }
+}
+
+[__NonCopyableType] struct S6 : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            S6 t;
+            t.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return t;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+void main()
+{
+    S1 s1;
+    S2 s2;
+    S3 s3;
+    S4 s4 = S4(0);
+    S5 s5 = S5(0);
+    S6 s6 = S6(0);
+}
diff --git a/tests/diagnostics/ctor-ordinary-retval-legal.slang b/tests/diagnostics/ctor-ordinary-retval-legal.slang
new file mode 100644
index 000000000..f2d1765ff
--- /dev/null
+++ b/tests/diagnostics/ctor-ordinary-retval-legal.slang
@@ -0,0 +1,104 @@
+//TEST:SIMPLE(filecheck=CHECK):
+
+// Test to catch compilation time errors with return values in constructors within an ordinary struct.
+// A constructor from the callsite's point of view is a function that returns the struct type.
+// A constructor from the inside the body is treated as a function that modifies `this` and not return.
+
+// In the following test-cases where the constructor either returns an expr or no value,
+// within a condition for early exit or not, it automatically gets resolved as
+// this = t;
+// return *this; or return;
+
+struct S
+{
+    float v;    
+}
+
+struct S1a : S
+{
+    __init()
+    {
+        this.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return this;
+    }
+}
+
+struct S1b : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            this.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return this;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+struct S2a : S
+{
+    __init()
+    {
+        S2a t;
+        t.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+struct S2b : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            S2b t;
+            t.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+struct S3a : S
+{
+    __init()
+    {
+        S3a t;
+        t.v = 5;
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return t;
+    }
+}
+
+struct S3b : S
+{
+    __init(float u)
+    {
+        if (u != 0)
+        {
+            S3b t;
+            t.v = u;
+            // CHECK-NOT: ([[# @LINE+1]]): error
+            return t;
+        }
+        // CHECK-NOT: ([[# @LINE+1]]): error
+        return;
+    }
+}
+
+void main()
+{
+    S1a S1a;
+    S2a S2a;
+    S3a S3a;
+    S1b S1b = S1b(0);
+    S2b S2b = S2b(0);
+    S3b S3b = S3b(0);
+}
diff --git a/tests/language-feature/interfaces/constructor-noncopyable-return.slang b/tests/language-feature/interfaces/constructor-noncopyable-return.slang
new file mode 100644
index 000000000..9f6c7355a
--- /dev/null
+++ b/tests/language-feature/interfaces/constructor-noncopyable-return.slang
@@ -0,0 +1,60 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj -output-using-type
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+// Note: This test won't produce expected result on vulkan.
+//DISABLED_TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -vk -shaderobj -output-using-type
+
+// Executable test to check the correctness of return values from the constructors within a
+// NonCopyableType struct.
+
+// A constructor from the callsite's point of view is a function that returns the struct type.
+// A constructor from the inside the body is treated as a function that modifies `this` and not return.
+
+// In the following test-cases where the constructor either returns an expr or no value,
+// such as `return ...` within a condition for early exit or not,
+// It automatically gets resolved as
+// `return this` if it returns the struct type, else `return`
+RWStructuredBuffer<float> outputBuffer;
+
+[__NonCopyableType] struct Impl
+{
+    float x;
+    __init(float v)
+    {
+        if (v > 0)
+            {
+            this.x = 2 * v;
+            return this;
+        }
+        else
+        {
+            this.x = 3 * v;
+            return;
+        }
+    }
+
+    __init(int ival)
+    {
+        float val = ival;
+        Impl v = Impl(val);
+        return v;
+    }
+
+    __init()
+    {
+        float val = 2.0;
+        return Impl(val);
+    }
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint id : SV_DispatchThreadID)
+{
+    // CHECK: 4.000000
+    // CHECK: 2.000000
+    // CHECK: -3.000000
+    // CHECK: -6.000000
+    outputBuffer[id] = Impl().x;
+    outputBuffer[id + 1] = Impl(1).x;
+    outputBuffer[id + 2] = Impl(-1).x;
+    outputBuffer[id + 3] = Impl(-2.0).x;
+}
diff --git a/tests/language-feature/interfaces/constructor-return.slang b/tests/language-feature/interfaces/constructor-return.slang
new file mode 100644
index 000000000..85f474c57
--- /dev/null
+++ b/tests/language-feature/interfaces/constructor-return.slang
@@ -0,0 +1,60 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj -output-using-type
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -vk -shaderobj -output-using-type
+
+// Executable test to check the correctness of return values from the constructors within an
+// ordinary struct.
+
+// A constructor from the callsite's point of view is a function that returns the struct type.
+// A constructor from the inside the body is treated as a function that modifies `this` and not return.
+
+// In the following test-cases where the constructor either returns an expr or no value,
+// such as `return ...` within a condition for early exit or not,
+// It automatically gets resolved as
+// `this = t;` along with
+// `return this` if it returns the struct type, else `return`
+RWStructuredBuffer<float> outputBuffer;
+
+struct Impl
+{
+    float x;
+    __init(float v)
+    {
+        if (v > 0)
+        {
+            this.x = 2 * v;
+            return this;
+        }
+        else
+        {
+            this.x = 3 * v;
+            return;
+        }
+    }
+
+    __init(int ival)
+    {
+        float val = ival;
+        Impl v = Impl(val);
+        return v;
+    }
+
+    __init()
+    {
+        float val = 2.0;
+        return Impl(val);
+    }
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint id : SV_DispatchThreadID)
+{
+    // CHECK: 4.000000
+    // CHECK: 2.000000
+    // CHECK: -3.000000
+    // CHECK: -6.000000
+    outputBuffer[id] = Impl().x;
+    outputBuffer[id + 1] = Impl(1).x;
+    outputBuffer[id + 2] = Impl(-1).x;
+    outputBuffer[id + 3] = Impl(-2.0).x;
+}