IRGen: Count number of immediate class members in ClassMetadataLayout

Author: slavapestov

lib/IRGen/MetadataLayout.cpp

@@ -203,7 +203,7 @@ Address irgen::emitAddressOfFieldOffsetVector(IRGenFunction &IGF,
 /********************************** CLASSES ***********************************/
 
 ClassMetadataLayout::ClassMetadataLayout(IRGenModule &IGM, ClassDecl *decl)
-    : NominalMetadataLayout(Kind::Class) {
+    : NominalMetadataLayout(Kind::Class), NumImmediateMembers(0) {
 
   struct Scanner : LayoutScanner<Scanner, ClassMetadataScanner> {
     using super = LayoutScanner;
@@ -228,9 +228,26 @@ ClassMetadataLayout::ClassMetadataLayout(IRGenModule &IGM, ClassDecl *decl)
       super::noteStartOfGenericRequirements(forClass);
     }
 
+    void addGenericWitnessTable(CanType argType, ProtocolConformanceRef conf,
+                                ClassDecl *forClass) {
+      if (forClass == Target) {
+        Layout.NumImmediateMembers++;
+      }
+      super::addGenericWitnessTable(argType, conf, forClass);
+    }
+
+    void addGenericArgument(CanType argType, ClassDecl *forClass) {
+      if (forClass == Target) {
+        Layout.NumImmediateMembers++;
+      }
+      super::addGenericArgument(argType, forClass);
+    }
+
     void addMethod(SILDeclRef fn) {
-      if (fn.getDecl()->getDeclContext() == Target)
+      if (fn.getDecl()->getDeclContext() == Target) {
+        Layout.NumImmediateMembers++;
         Layout.MethodInfos.try_emplace(fn, getNextOffset());
+      }
       super::addMethod(fn);
     }
 
@@ -241,11 +258,21 @@ ClassMetadataLayout::ClassMetadataLayout(IRGenModule &IGM, ClassDecl *decl)
     }
 
     void addFieldOffset(VarDecl *field) {
-      if (field->getDeclContext() == Target)
+      if (field->getDeclContext() == Target) {
+        Layout.NumImmediateMembers++;
         Layout.FieldOffsets.try_emplace(field, getNextOffset());
+      }
       super::addFieldOffset(field);
     }
 
+    void addFieldOffsetPlaceholders(MissingMemberDecl *placeholder) {
+      if (placeholder->getDeclContext() == Target) {
+        Layout.NumImmediateMembers +=
+          placeholder->getNumberOfFieldOffsetVectorEntries();
+      }
+      super::addFieldOffsetPlaceholders(placeholder);
+    }
+
     void addVTableEntries(ClassDecl *forClass) {
       if (forClass == Target)
         Layout.VTableOffset = getNextOffset();

lib/IRGen/MetadataLayout.h

@@ -168,6 +168,9 @@ class ClassMetadataLayout : public NominalMetadataLayout {
   /// The start of the field-offset vector.
   StoredOffset FieldOffsetVector;
 
+  /// The number of members to add after superclass metadata.
+  unsigned NumImmediateMembers;
+
   const StoredMethodInfo &getStoredMethodInfo(SILDeclRef method) const {
     auto it = MethodInfos.find(method);
     assert(it != MethodInfos.end());
@@ -222,6 +225,13 @@ class ClassMetadataLayout : public NominalMetadataLayout {
 
   Offset getFieldOffsetVectorOffset(IRGenFunction &IGF) const;
 
+  /// The number of members to add after superclass metadata. The size of
+  /// this metadata is the superclass size plus the number of immediate
+  /// members in the class itself.
+  unsigned getNumImmediateMembers() const {
+    return NumImmediateMembers;
+  }
+
   static bool classof(const MetadataLayout *layout) {
     return layout->getKind() == Kind::Class;
   }