Added a '_typeList' property to 'TensorArrayProtocol'.

Author: eaplatanios

lib/Sema/DerivedConformanceTensorArrayProtocol.cpp

@@ -363,6 +363,89 @@ static ValueDecl *deriveTensorArrayProtocol_tensorHandleCount(
   return tensorHandleCountDecl;
 }
 
+
+/// Derive the body for the '_typeList' getter.
+static void
+deriveBodyTensorArrayProtocol_typeList(AbstractFunctionDecl *funcDecl) {
+  auto *parentDC = funcDecl->getParent();
+  auto *nominal = funcDecl->getDeclContext()->getSelfNominalTypeDecl();
+  auto &C = nominal->getASTContext();
+
+  auto *tensorGroupProto = C.getProtocol(KnownProtocolKind::TensorGroup);
+  auto *typeListReq = getProtocolRequirement(tensorGroupProto, C.Id_typeList);
+
+  // Concatenate all member `_typeList` arrays.
+  Type arrayType = BoundGenericType::get(
+    C.getArrayDecl(), Type(), 
+    {C.getTensorDataTypeDecl()->getDeclaredInterfaceType()});
+  auto *arrayTypeExpr = TypeExpr::createImplicit(arrayType, C);
+  auto plusOpLookup = C.getArrayDecl()->lookupDirect(C.getIdentifier("+"));
+  assert(plusOpLookup.size() == 1 && "Ambiguous 'Array.+' operator.");
+  ValueDecl *plusOpDecl = plusOpLookup.front();
+  auto plusOpDRE = new (C) 
+      DeclRefExpr(plusOpDecl, DeclNameLoc(), /*Implicit*/ true);
+  auto plusOpExpr = new (C)
+      DotSyntaxCallExpr(plusOpDRE, SourceLoc(), arrayTypeExpr);
+  Expr *typeListExpr = ArrayExpr::create(C, SourceLoc(), {}, {}, SourceLoc());
+  for (auto member : nominal->getStoredProperties()) {
+    auto memberType =
+        parentDC->mapTypeIntoContext(member->getValueInterfaceType());
+    auto *memberTypeExpr = TypeExpr::createImplicit(memberType, C);
+    auto *memberTypeListExpr = new (C) 
+        MemberRefExpr(memberTypeExpr, SourceLoc(), typeListReq,
+                      DeclNameLoc(), /*Implicit*/ true);
+    // Create expression `lhsArg + rhsArg`.
+    auto *plusOpArgs =
+        TupleExpr::create(C, SourceLoc(), {typeListExpr, memberTypeListExpr}, 
+                          {}, {}, SourceLoc(), /*HasTrailingClosure*/ false,
+                          /*Implicit*/ true);
+    typeListExpr = new (C) BinaryExpr(plusOpExpr, plusOpArgs, 
+                                      /*Implicit*/ true);
+  }
+
+  // Return the resulting data types array.
+  auto *returnStmt = new (C) ReturnStmt(SourceLoc(), typeListExpr);
+  auto *body = BraceStmt::create(C, SourceLoc(), {returnStmt}, SourceLoc(),
+                                 /*Implicit*/ true);
+  funcDecl->setBody(BraceStmt::create(C, SourceLoc(), {body}, SourceLoc(),
+                                      /*Implicit*/ true));
+}
+
+/// Derive a '_typeList' implementation.
+static ValueDecl *deriveTensorArrayProtocol_typeList(
+    DerivedConformance &derived) {
+  auto nominal = derived.Nominal;
+  auto &TC = derived.TC;
+  ASTContext &C = TC.Context;
+
+  auto parentDC = derived.getConformanceContext();
+  Type dataTypeArrayType = BoundGenericType::get(
+    C.getArrayDecl(), Type(), 
+    {C.getTensorDataTypeDecl()->getDeclaredInterfaceType()});
+  auto returnType = parentDC->mapTypeIntoContext(dataTypeArrayType);
+
+  // Create `_typeList` property declaration.
+  VarDecl *typeListDecl;
+  PatternBindingDecl *patDecl;
+  std::tie(typeListDecl, patDecl) = derived.declareDerivedProperty(
+      C.Id_typeList, returnType, returnType, /*isStatic*/ false,
+      /*isFinal*/ false);
+
+  // Add `@inlinable` to the `_typeList` declaration.
+  if (nominal->getEffectiveAccess() > AccessLevel::Internal)
+    typeListDecl->getAttrs().add(new (C) InlinableAttr(/*implicit*/ true));
+
+  // Create `_typeList` getter.
+  auto *getterDecl = derived.declareDerivedPropertyGetter(
+      TC, typeListDecl, returnType);
+  getterDecl->setBodySynthesizer(deriveBodyTensorArrayProtocol_typeList);
+  typeListDecl->setAccessors(StorageImplInfo::getImmutableComputed(),
+                             SourceLoc(), {getterDecl}, SourceLoc());
+  derived.addMembersToConformanceContext({getterDecl, typeListDecl, patDecl});
+
+  return typeListDecl;
+}
+
 // Synthesize body for `init(_owning:count:)`.
 static void 
 deriveBodyTensorArrayProtocol_init(AbstractFunctionDecl *funcDecl) {
@@ -584,6 +667,8 @@ ValueDecl *DerivedConformance::deriveTensorArrayProtocol(
     return deriveTensorArrayProtocol_unpackTensorHandles(*this);
   if (requirement->getBaseName() == TC.Context.Id_tensorHandleCount)
     return deriveTensorArrayProtocol_tensorHandleCount(*this);
+  if (requirement->getBaseName() == TC.Context.Id_typeList)
+    return deriveTensorArrayProtocol_typeList(*this);
   if (requirement->getBaseName() == DeclBaseName::createConstructor())
     return deriveTensorArrayProtocol_init(*this);
   TC.diagnose(requirement->getLoc(), 

lib/Sema/DerivedConformances.cpp

@@ -230,6 +230,11 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,
     // TensorArrayProtocol._tensorHandleCount
     if (name.isSimpleName(ctx.Id_tensorHandleCount))
       return getRequirement(KnownProtocolKind::TensorArrayProtocol);
+    
+    // SWIFT_ENABLE_TENSORFLOW
+    // TensorArrayProtocol._typeList
+    if (name.isSimpleName(ctx.Id_typeList) && !requirement->isStatic())
+      return getRequirement(KnownProtocolKind::TensorArrayProtocol);
 
     // SWIFT_ENABLE_TENSORFLOW
     // TensorGroup._typeList

stdlib/public/TensorFlow/TensorGroup.swift

@@ -36,6 +36,7 @@ public protocol TensorArrayProtocol {
   func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?)
 
   var _tensorHandleCount: Int32 { get }
+  var _typeList: [TensorDataType] { get }
 
   init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int)
 }
@@ -69,13 +70,16 @@ public protocol TensorGroup : TensorArrayProtocol {
 public extension TensorGroup {
   /// The number of tensor fields in this type.
   static var _tensorHandleCount: Int32 { return Int32(Self._typeList.count) }
-  var _tensorHandleCount: Int32 { return Int32(Self._typeList.count) }
 
   /// An array of `nil`s with the same number of elements as `_outputTypeList`.
   /// The `nil` represents unknown shape.
   static var _unknownShapeList: [TensorShape?] {
     return Array(repeating: nil, count: _typeList.count)
   }
+  
+  // The following instance properties are from `TensorArrayProtocol`.
+  var _tensorHandleCount: Int32 { return Int32(Self._typeList.count) }
+  var _typeList: [TensorDataType] { return Self._typeList }
 
   init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int) {
     precondition(count == Self._typeList.count)
@@ -223,9 +227,13 @@ extension Array : TensorArrayProtocol where Element : TensorGroup {
   }
 
   public var _tensorHandleCount: Int32 {
-    var count: Int32 = 0
-    for elem in self { count += elem._tensorHandleCount }
-    return count
+    return Element._tensorHandleCount * Int32(count)
+  }
+
+  public var _typeList: [TensorDataType] {
+    return Array<TensorDataType>([[TensorDataType]](
+      repeating: Element._typeList,
+      count: Int(Element._tensorHandleCount)).joined())
   }
 
   public init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int) {

test/TensorFlowRuntime/tracer.swift

@@ -191,6 +191,10 @@ TracerTests.testAllBackends("Advanced") {
       return model._tensorHandleCount + optimizer._tensorHandleCount
     }
 
+    public var _typeList: [TensorDataType] {
+      return model._typeList + optimizer._typeList
+    }
+
     func _makeInstance<C: Collection>(owning inputs: C) -> State
       where C.Element == CTensorHandle {
       assert(inputs.count == 4)