[autodiff] Change getTangentStoredProperty() to use a Projection instead of FieldIndexCacheBase.

This is NFCI.

THis is in preparation for making FieldIndexCacheBase a templated subclass.

Author: gottesmm

include/swift/SILOptimizer/Differentiation/Common.h

@@ -22,6 +22,7 @@
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/SIL/SILDifferentiabilityWitness.h"
 #include "swift/SIL/SILFunction.h"
+#include "swift/SIL/Projection.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/Analysis/ArraySemantic.h"
@@ -166,8 +167,11 @@ VarDecl *getTangentStoredProperty(ADContext &context, VarDecl *originalField,
 /// Returns the tangent stored property of the original stored property
 /// referenced by the given projection instruction with the given base type.
 /// On error, emits diagnostic and returns nullptr.
+///
+/// NOTE: Asserts if \p projectionInst is not one of: struct_extract,
+/// struct_element_addr, or ref_element_addr.
 VarDecl *getTangentStoredProperty(ADContext &context,
-                                  FieldIndexCacheBase *projectionInst,
+                                  SingleValueInstruction *projectionInst,
                                   CanType baseType,
                                   DifferentiationInvoker invoker);
 

lib/SILOptimizer/Differentiation/Common.cpp

@@ -329,14 +329,16 @@ VarDecl *getTangentStoredProperty(ADContext &context, VarDecl *originalField,
 }
 
 VarDecl *getTangentStoredProperty(ADContext &context,
-                                  FieldIndexCacheBase *projectionInst,
+                                  SingleValueInstruction *projectionInst,
                                   CanType baseType,
                                   DifferentiationInvoker invoker) {
   assert(isa<StructExtractInst>(projectionInst) ||
          isa<StructElementAddrInst>(projectionInst) ||
          isa<RefElementAddrInst>(projectionInst));
+  Projection proj(projectionInst);
   auto loc = getValidLocation(projectionInst);
-  return getTangentStoredProperty(context, projectionInst->getField(), baseType,
+  auto *field = proj.getVarDecl(projectionInst->getOperand(0)->getType());
+  return getTangentStoredProperty(context, field, baseType,
                                   loc, invoker);
 }
 

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -1043,7 +1043,7 @@ class JVPCloner::Implementation final
     auto structType =
         remapSILTypeInDifferential(sei->getOperand()->getType()).getASTType();
     auto *tanField =
-        getTangentStoredProperty(context, sei, structType, invoker);
+      getTangentStoredProperty(context, sei, structType, invoker);
     if (!tanField) {
       errorOccurred = true;
       return;
@@ -1074,7 +1074,7 @@ class JVPCloner::Implementation final
     auto structType =
         remapSILTypeInDifferential(seai->getOperand()->getType()).getASTType();
     auto *tanField =
-        getTangentStoredProperty(context, seai, structType, invoker);
+      getTangentStoredProperty(context, seai, structType, invoker);
     if (!tanField) {
       errorOccurred = true;
       return;

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -1832,7 +1832,9 @@ bool PullbackCloner::Implementation::run() {
         }
       }
       // Diagnose unsupported stored property projections.
-      if (auto *inst = dyn_cast<FieldIndexCacheBase>(v)) {
+      if (isa<StructExtractInst>(v) || isa<RefElementAddrInst>(v) ||
+          isa<StructElementAddrInst>(v)) {
+        auto *inst = cast<SingleValueInstruction>(v);
         assert(inst->getNumOperands() == 1);
         auto baseType = remapType(inst->getOperand(0)->getType()).getASTType();
         if (!getTangentStoredProperty(getContext(), inst, baseType,