Allow [mutating] methods on existential values (#2245)

The problematic case is when an `interface` has a `[mutating]` method:

    interface ICounter
    {
        [mutating] void increment();
    }

and code tries to invoke that method on a value of existential type:

    ICounter c = ...;
    c.increment();

We know that the existential value `c` is conceptually a tuple of:

* A concrete type `X`
* A witness that `X : ICounter`
* A value `v` of type `X`

We simply want to invoke `increment()` on the `v` part, using the `X : ICounter` witness table.
The catch that the compiler faces is that the variable `c` is mutable, so we need to be careful that we "snapshot" its value (the tuple `X, X:ICounter, v`) at a single point.
The snapshotting behavior is important when invoking a method that involves `This` or associated types in its signature, so we cannot get rid of it.

The snapshotting we do relies on the idea of a `LetExpr` AST node, which cannot be written in the input syntax.
A `LetExpr` introduces a variable binding (with an initial-value expression) and then evaluates a body expression in the context of that binding.
For a call site like `c.increment()` the front-end makes an intermediate copy of `c` and then "opens" that immutable value to get at the elements of the tuple `X`, `X : ICounter`, `v`.
The resulting AST after checking looks something like:

    ICounter c = ...;
    (let tmp = c in extractExistentialValue(tmp)).increment();

In that form it is more clear why the attempt to call `increment()` fails:

1. The binding `tmp` sure looks immutable
2. There is no logic in the compiler to make `extractExistentialValue(x)` be an l-value if `x` is
3. There is seemingly no logic to write back from `tmp` to `c` when the operation completes

Let us walk through those problems in order.

Item (1) turns out to be a bit of a non-issue.
Despite the way that I've written out `let` expressions above, the logic in `moveTemp()` in the compiler actually introduces a *mutable* binding.

Item (2) can be fixed for the purposes of semantic checking by modifying `openExistential()`.
Simplistically, we make the overall expression be an l-value if the operand is.

Item (3) is handled at the level of AST->IR lowering. Each kind of expression that can form an l-value needs to have a way to represent the "location" of that l-value in the `LoweredValInfo` type.
This change adds a case to handle the `extractExistentialVal` operation, by tracking both the extract value (of concrete type) and the underlying l-value (of existential type).

Where all of this comes crashing against reality a bit is that the scoping I've drawn for the `let` expressions above kind of doesn't work once we look at types.
The basic problem is that the *type* of the `(let tmp = c in ...)` expression is the concrete type `X` that was extracted from the existential.
That type can conceptually be written as `ExtractExistentialType(tmp)` which, notably, references `tmp`.
That means that we end up with AST expression nodes that reference the variable `tmp` *outside* of its scope.
Furthermore, those references to `tmp` can end up being lowered to IR *before* we have lowered the `let ...` expression itself.

Fixing the scoping issue turns out to be a major undertaking.
The first (and more obvious) issue is needing to address the scoping problem.
The solution I implemented includes a bit of refactoring to make all the `SemanticsVisitor` types better able to pass around the contextual scope-dependent state that might be needed during semantic checking, but really only adds a single piece of state.

The semantic-checking state used for checking expressions is bottlenecked so that there will (or at least *should*) always be an explicit representation of a "scope" that surrounds a complete expression (as opposed to a sub-expression).
When a `LetExpr` needs to be introduced, it is added to a pending list on the active scope, rather than being added locally.
Once the complete expression is checked, the resulting expression is wrapped up in the pending `LetExpr`s so that their scope is as broad as possible.

Technically this solution doesn't cover all cases. For example:

    interface ICell { associatedtype Content; Content getContent(); }
    ...
    ICell cell = ...;
    let content = cell.getContent();

In this case the type of `content` refers to the binding introduced by a `LetExpr` in the initial-value expression.
I am leaving such issues as a piece of future work, in the hopes that we can get at least a partial fix for the problem in place.
A future fix probably nees to extend the scoping even wider (e.g., by unwrapping the `LetExpr`s from the initial-value expression and turning them into distinct temporaries).

The second piece of the fix is that we need a way for the modified value of the extracted existential to be "written back" to the original location.

Well...

We are actually being a little slippery here, based on some logic in the compiler codebase that I guess Just Works.
When AST->IR lowering encounters a `LetExpr` that binds an l-value to a name, it actually ends up binding that name more or less as a *reference* to that l-value.
At this point the `let`-ness of `LetExpr` is very much in doubt: the binding can be mutable, and it can even be an *alias* of some location?!?

In any case, the result is that the AST->IR codegen logic implicitly handles the "write-back" because the `let`-bound temporary is actually an alias for the original location.
A more complete future fix might need to introduce a distinct case in `LoweredValInfo` to handle the case of copy of a mutable temporary.

1 files changed, 155 insertions, 5 deletions

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 4ff43b697..5cfb07c1c 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -80,7 +80,8 @@ struct SubscriptInfo : ExtendedValueInfo
 struct BoundStorageInfo;
 struct BoundMemberInfo;
 struct SwizzledLValueInfo;
-
+struct CopiedValInfo;
+struct ExtractedExistentialValInfo;
 
 // This type is our core representation of lowered values.
 // In the simple case, it just wraps an `IRInst*`.
@@ -113,6 +114,9 @@ struct LoweredValInfo
 
         // The result of applying swizzling to an l-value
         SwizzledLValue,
+
+        // The value extracted from an opened existential
+        ExtractedExistential,
     };
 
     union
@@ -185,6 +189,15 @@ struct LoweredValInfo
         SLANG_ASSERT(flavor == Flavor::SwizzledLValue);
         return (SwizzledLValueInfo*)ext;
     }
+
+    static LoweredValInfo extractedExistential(
+        ExtractedExistentialValInfo* extInfo);
+
+    ExtractedExistentialValInfo* getExtractedExistentialValInfo()
+    {
+        SLANG_ASSERT(flavor == Flavor::ExtractedExistential);
+        return (ExtractedExistentialValInfo*)ext;
+    }
 };
 
 // This case is used to indicate a reference to an AST-level
@@ -272,6 +285,27 @@ struct SwizzledLValueInfo : ExtendedValueInfo
     UInt            elementIndices[4];
 };
 
+// Represents the results of extractng a value of
+// some (statically unknown) concrete type from
+// an existential, in an l-value context.
+//
+struct ExtractedExistentialValInfo : ExtendedValueInfo
+{
+    // The extracted value
+    IRInst* extractedVal;
+
+    // The original existential value
+    LoweredValInfo existentialVal;
+
+    // The type of `existentialVal`
+    IRType* existentialType;
+
+    // The IR witness table for the conformance of
+    // the type of `extractedVal` to `existentialType`
+    //
+    IRInst* witnessTable;
+};
+
 LoweredValInfo LoweredValInfo::boundMember(
     BoundMemberInfo*    boundMemberInfo)
 {
@@ -308,6 +342,16 @@ LoweredValInfo LoweredValInfo::swizzledLValue(
     return info;
 }
 
+LoweredValInfo LoweredValInfo::extractedExistential(
+    ExtractedExistentialValInfo* extInfo)
+{
+    LoweredValInfo info;
+    info.flavor = Flavor::ExtractedExistential;
+    info.ext = extInfo;
+    return info;
+}
+
+
 // An "environment" for mapping AST declarations to IR values.
 //
 // This is required because in some cases we might lower the
@@ -935,6 +979,13 @@ top:
                 swizzleInfo->elementIndices));
         }
 
+    case LoweredValInfo::Flavor::ExtractedExistential:
+        {
+            auto info = lowered.getExtractedExistentialValInfo();
+
+            return LoweredValInfo::simple(info->extractedVal);
+        }
+
     default:
         SLANG_UNEXPECTED("unhandled value flavor");
         UNREACHABLE_RETURN(LoweredValInfo());
@@ -2108,6 +2159,7 @@ void addInArg(
     case LoweredValInfo::Flavor::SwizzledLValue:
     case LoweredValInfo::Flavor::BoundStorage:
     case LoweredValInfo::Flavor::BoundMember:
+    case LoweredValInfo::Flavor::ExtractedExistential:
         args.add(getSimpleVal(context, argVal));
         break;
 
@@ -2784,6 +2836,8 @@ static LoweredValInfo _emitCallToAccessor(
 template<typename Derived>
 struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
 {
+    static bool isLValueContext() { return Derived::_isLValueContext(); }
+
     IRGenContext* context;
 
     IRBuilder* getBuilder() { return context->irBuilder; }
@@ -2797,6 +2851,13 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
         return this->dispatch(expr);
     }
 
+    LoweredValInfo lowerSubExpr(Expr* expr, IRGenContext* subContext)
+    {
+        IRBuilderSourceLocRAII sourceLocInfo(getBuilder(), expr->loc);
+        Derived d;
+        d.context = subContext;
+        return d.dispatch(expr);
+    }
 
     LoweredValInfo visitVarExpr(VarExpr* expr)
     {
@@ -3802,8 +3863,20 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
 
     LoweredValInfo visitLetExpr(LetExpr* expr)
     {
-        // TODO: deal with the case where we might want to capture
-        // a reference to the bound value...
+        // Note: The semantics here are annoyingly subtle.
+        //
+        // If `expr->decl->initExpr` is an l-value, then we will set things up
+        // so that `expr->decl` is bound as an *alias* for that l-value.
+        //
+        // Otherwise, `expr->decl` will simply be bound to the r-value.
+        //
+        // The first case is necessary to make `maybeMoveTemp` operations that
+        // produce l-value results work correctly, but seems slippery.
+        //
+        // TODO: We should probably have two AST node types to cover the two
+        // different use cases of `LetExpr`: the definitely-immutable case that
+        // actually behaves like a `let`, and this other mutable-alias case that
+        // feels kind of messy and gross.
 
         auto initVal = lowerLValueExpr(context, expr->decl->initExpr);
         setGlobalValue(context, expr->decl, initVal);
@@ -3813,17 +3886,66 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
 
     LoweredValInfo visitExtractExistentialValueExpr(ExtractExistentialValueExpr* expr)
     {
+        // We are being asked to extract the value from an existential, which
+        // is itself a single IR op. However, we also need to handle the case
+        // where `expr` might be used as an l-value, in which case we need
+        // additional information to allow any mutations through the extracted
+        // value to be written back.
+
         auto existentialType = lowerType(context, getType(getASTBuilder(), expr->declRef));
-        auto existentialVal = getSimpleVal(context, emitDeclRef(context, expr->declRef, existentialType));
+        auto existentialVal = emitDeclRef(context, expr->declRef, existentialType);
+
+        // Note that we make a *copy* of the existential value that is definitely
+        // a simple r-value. This ensures that all the `extractExistential*()` operations
+        // below work on the same consistent IR value.
+        //
+        auto existentialValCopy = getSimpleVal(context, existentialVal);
 
         auto openedType = lowerType(context, expr->type);
 
-        return LoweredValInfo::simple(getBuilder()->emitExtractExistentialValue(openedType, existentialVal));
+        auto extractedVal = getBuilder()->emitExtractExistentialValue(
+            openedType, existentialValCopy);
+
+        if(!isLValueContext())
+        {
+            // If we are in an r-value context, we can directly use the `extractExistentialValue`
+            // instruction as the result, and life is simple.
+            //
+            return LoweredValInfo::simple(extractedVal);
+        }
+
+        // In an l-value context, we need to track the information necessary so that
+        // if a new/modified value of `openedType` was produced, we could write it
+        // back into the original `existentialVal`'s location.
+        //
+        // The write-back is actually pretty simple: it is just a `makeExisential` op.
+        // In order to be able to emit that op later, we need to track the operands
+        // that it would use. The first operand would be the new concrete value (which
+        // would implicitly encode the concrete type via its IR type) while the second
+        // is the witness table for the conformance to the existential.
+        //
+        // Note: We are assuming/requiring here that any value "written back" must have
+        // the exact same concrete type as `extractedVal`, so taht it can use the same
+        // IR witness table. The front-end should be enforcing that constraint, and we
+        // have no way to check or enforce it at this point.
+
+        auto witnessTable = getBuilder()->emitExtractExistentialWitnessTable(existentialValCopy);
+
+        RefPtr<ExtractedExistentialValInfo> info = new ExtractedExistentialValInfo();
+        info->extractedVal = extractedVal;
+        info->existentialVal = existentialVal;
+        info->existentialType = existentialType;
+        info->witnessTable = witnessTable;
+
+        context->shared->extValues.add(info);
+        return LoweredValInfo::extractedExistential(info);
     }
 };
 
 struct LValueExprLoweringVisitor : ExprLoweringVisitorBase<LValueExprLoweringVisitor>
 {
+    static bool _isLValueContext() { return true; }
+
     // When visiting a swizzle expression in an l-value context,
     // we need to construct a "swizzled l-value."
     LoweredValInfo visitMatrixSwizzleExpr(MatrixSwizzleExpr*)
@@ -3901,6 +4023,8 @@ struct LValueExprLoweringVisitor : ExprLoweringVisitorBase<LValueExprLoweringVis
 
 struct RValueExprLoweringVisitor : ExprLoweringVisitorBase<RValueExprLoweringVisitor>
 {
+    static bool _isLValueContext() { return false; }
+
     // A matrix swizzle in an r-value context can save time by just
     // emitting the matrix swizzle instructions directly.
     LoweredValInfo visitMatrixSwizzleExpr(MatrixSwizzleExpr* expr)
@@ -5289,6 +5413,32 @@ top:
         }
         break;
 
+    case LoweredValInfo::Flavor::ExtractedExistential:
+        {
+            // The `left` value is the result of opening an existential.
+            //
+            auto leftInfo = left.getExtractedExistentialValInfo();
+            auto existentialVal = leftInfo->existentialVal;
+
+            // The actual desitnation we need to store into is the
+            // existential value itself.
+            //
+            left = existentialVal;
+
+            // The `right` value must be of the same concrete type as
+            // the opened value, but the new destination is of the
+            // original existential type, so we need to wrap it up
+            // appropriately.
+            //
+            right = LoweredValInfo::simple(builder->emitMakeExistential(
+                leftInfo->existentialType,
+                getSimpleVal(context, right),
+                leftInfo->witnessTable));
+
+            goto top;
+        }
+        break;
+
     default:
         SLANG_UNIMPLEMENTED_X("assignment");
         break;