Fix some issues around codegen for l-values and assignment (#601)

The problem here arose when a complicated l-value was formed like:

```hlsl
struct Foo { float4 a; }
RWStructuredBuffer<Foo> gBuffer;

gBuffer[index].a.xz += whatever;
```

In this case the `gBuffer[index].a.xz` expression is a complex l-value in multiple ways:

* The `gBuffer[index]` subscript could be routed to either a `get` accessor or a `ref` accessor (and maybe also a `set` accessor if we add one to the stdlib definition), and we defer the choice  of which to call until as late as possible in codegen today.

* The `_.a` part then becomes a "bound member acess" because we can't actually produce a direct pointer until we've resolved how to implement the subscript operation.

* The `_.xz` part becomes a "swizzled l-value" because there is *no* way to materialize it as a pointer to contiguous storage in the orignal object (the `x` and `z` components of a vector aren't contiguous).

Recent changes to support atomic operations on buffer elements introduced the `ref` accessor on `RWStructuredBuffer`, which made it possible to form a pointer to a buffer element in the IR. This interacted with some code for the "bound member" case that was trying to only introduce a temporary when absolutely necessary, and was doing so by assuming anything with an address didn't need to be moved into a temporary.

The first fix is to clean up that logic in the bound-member case for assignment: always create a temporary, rather than do it conditionally.

The second fix here is more systemic: we add logic to try to coerce the representation of an l-value during codegen into being a simple address, and employ that in cases where we know an address is desired. In a case like the above this helps to get things into the form that is required, so that a swizzled store can be issued.

There is still some potential for cleanup in this logic, but I don't want to introduce more changes than seem necessary to fix the original problem.

1 files changed, 99 insertions, 38 deletions

diff --git a/source/slang/lower-to-ir.cpp b/source/slang/lower-to-ir.cpp
index 761fe9c91..27325ee38 100644
--- a/source/slang/lower-to-ir.cpp
+++ b/source/slang/lower-to-ir.cpp
@@ -2948,39 +2948,32 @@ void lowerStmt(
     }
 }
 
-static LoweredValInfo maybeMoveMutableTemp(
+/// Create and return a mutable temporary initialized with `val`
+static LoweredValInfo moveIntoMutableTemp(
     IRGenContext*           context,
     LoweredValInfo const&   val)
 {
-    switch(val.flavor)
-    {
-    case LoweredValInfo::Flavor::Ptr:
-        return val;
+    IRInst* irVal = getSimpleVal(context, val);
+    auto type = irVal->getDataType();
+    auto var = createVar(context, type);
 
-    default:
-        {
-            IRInst* irVal = getSimpleVal(context, val);
-            auto type = irVal->getDataType();
-            auto var = createVar(context, type);
-
-            assign(context, var, LoweredValInfo::simple(irVal));
-            return var;
-        }
-        break;
-    }
+    assign(context, var, LoweredValInfo::simple(irVal));
+    return var;
 }
 
-IRInst* getAddress(
+/// Try to coerce `inVal` into a `LoweredValInfo::ptr()` with a simple address.
+LoweredValInfo tryGetAddress(
     IRGenContext*           context,
-    LoweredValInfo const&   inVal,
-    SourceLoc               diagnosticLocation)
+    LoweredValInfo const&   inVal)
 {
     LoweredValInfo val = inVal;
 
     switch(val.flavor)
     {
     case LoweredValInfo::Flavor::Ptr:
-        return val.val;
+        // The `Ptr` case means that we already have an IR value with
+        // the address of our value. Easy!
+        return val;
 
     case LoweredValInfo::Flavor::BoundSubscript:
         {
@@ -3005,13 +2998,60 @@ IRInst* getAddress(
 
                 // The result from the call should be a pointer, and it
                 // is the address that we wanted in the first place.
-                return getSimpleVal(context, refVal);
+                return LoweredValInfo::ptr(getSimpleVal(context, refVal));
             }
 
             // Otherwise, there was no `ref` accessor, and so it is not possible
             // to materialize this location into a pointer for whatever purpose
             // we have in mind (e.g., passing it to an atomic operation).
         }
+        break;
+
+    case LoweredValInfo::Flavor::BoundMember:
+        {
+            auto boundMemberInfo = val.getBoundMemberInfo();
+
+            // If we hit this case, then it means that we have a reference
+            // to a single field in something, but for whatever reason the
+            // higher-level logic was not able to turn it into a pointer
+            // already (maybe the base value for the field reference is
+            // a `BoundSubscript`, etc.).
+            //
+            // We need to read the entire base value out, modify the field
+            // we care about, and then write it back.
+
+            auto declRef = boundMemberInfo->declRef;
+            if( auto fieldDeclRef = declRef.As<StructField>() )
+            {
+                auto baseVal = boundMemberInfo->base;
+                auto basePtr = tryGetAddress(context, baseVal);
+
+                return extractField(context, boundMemberInfo->type, basePtr, fieldDeclRef);
+            }
+
+        }
+        break;
+
+    case LoweredValInfo::Flavor::SwizzledLValue:
+        {
+            auto originalSwizzleInfo = val.getSwizzledLValueInfo();
+            auto originalBase = originalSwizzleInfo->base;
+
+            UInt elementCount = originalSwizzleInfo->elementCount;
+
+            auto newBase = tryGetAddress(context, originalBase);
+            RefPtr<SwizzledLValueInfo> newSwizzleInfo = new SwizzledLValueInfo();
+            context->shared->extValues.Add(newSwizzleInfo);
+
+            newSwizzleInfo->base = newBase;
+            newSwizzleInfo->type = originalSwizzleInfo->type;
+            newSwizzleInfo->elementCount = elementCount;
+            for(UInt ee = 0; ee < elementCount; ++ee)
+                newSwizzleInfo->elementIndices[ee] = originalSwizzleInfo->elementIndices[ee];
+
+            return LoweredValInfo::swizzledLValue(newSwizzleInfo);
+        }
+        break;
 
     // TODO: are there other cases we need to handled here?
 
@@ -3019,6 +3059,23 @@ IRInst* getAddress(
         break;
     }
 
+    // If none of the special cases above applied, then we werent' able to make
+    // this value into a pointer, and we should just return it as-is.
+    return val;
+}
+
+IRInst* getAddress(
+    IRGenContext*           context,
+    LoweredValInfo const&   inVal,
+    SourceLoc               diagnosticLocation)
+{
+    LoweredValInfo val = tryGetAddress(context, inVal);
+
+    if( val.flavor == LoweredValInfo::Flavor::Ptr )
+    {
+        return val.val;
+    }
+
     context->getSink()->diagnose(diagnosticLocation, Diagnostics::invalidLValueForRefParameter);
     return nullptr;
 }
@@ -3031,29 +3088,26 @@ void assign(
     LoweredValInfo left = inLeft;
     LoweredValInfo right = inRight;
 
+    // Before doing the case analysis on the shape of the `left` value,
+    // we might as well go ahead and see if we can coerce it into
+    // a simple pointer, since that would make our life a lot easier
+    // when handling complex cases.
+    //
+    left = tryGetAddress(context, left);
+
     auto builder = context->irBuilder;
 
 top:
     switch (left.flavor)
     {
     case LoweredValInfo::Flavor::Ptr:
-        switch (right.flavor)
         {
-        case LoweredValInfo::Flavor::Simple:
-        case LoweredValInfo::Flavor::Ptr:
-        case LoweredValInfo::Flavor::SwizzledLValue:
-        case LoweredValInfo::Flavor::BoundSubscript:
-        case LoweredValInfo::Flavor::BoundMember:
-            {
-                builder->emitStore(
-                    left.val,
-                    getSimpleVal(context, right));
-            }
-            break;
-
-        default:
-            SLANG_UNIMPLEMENTED_X("assignment");
-            break;
+            // The `left` value is just a pointer, so we can emit
+            // a store to it directly.
+            //
+            builder->emitStore(
+                left.val,
+                getSimpleVal(context, right));
         }
         break;
 
@@ -3064,6 +3118,11 @@ top:
             auto swizzleInfo = left.getSwizzledLValueInfo();
             auto loweredBase = swizzleInfo->base;
 
+            // Note that the call to `tryGetAddress` at the start should
+            // ensure that `loweredBase` has been simplified as much as
+            // possible (e.g., if it is possible to turn it into a
+            // `LoweredValInfo::ptr()` then that will have been done).
+
             switch( loweredBase.flavor )
             {
             default:
@@ -3209,7 +3268,7 @@ top:
                 // materialize the base value and move it into
                 // a mutable temporary if needed
                 auto baseVal = boundMemberInfo->base;
-                auto tempVal = maybeMoveMutableTemp(context, materialize(context, baseVal));
+                auto tempVal = moveIntoMutableTemp(context, baseVal);
 
                 // extract the field l-value out of the temporary
                 auto tempFieldVal = extractField(context, boundMemberInfo->type, tempVal, fieldDeclRef);
@@ -3219,6 +3278,8 @@ top:
 
                 // write back the modified temporary to the base l-value
                 assign(context, baseVal, tempVal);
+
+                return;
             }
             else
             {