[Compile Time Constant Extraction] Add extraction of arrays (#62491)

Author: jPaolantonio

include/swift/AST/ConstTypeInfo.h

@@ -27,7 +27,15 @@ class Type;
 /// in a type property initializer expression
 class CompileTimeValue {
 public:
-  enum ValueKind { RawLiteral, InitCall, Builder, Dictionary, Runtime, Tuple };
+  enum ValueKind {
+    RawLiteral,
+    InitCall,
+    Builder,
+    Dictionary,
+    Array,
+    Tuple,
+    Runtime
+  };
 
   ValueKind getKind() const { return Kind; }
 
@@ -126,6 +134,23 @@ class TupleValue : public CompileTimeValue {
   std::vector<TupleElement> Elements;
 };
 
+/// An array literal value representation
+class ArrayValue : public CompileTimeValue {
+public:
+  ArrayValue(std::vector<std::shared_ptr<CompileTimeValue>> Elements)
+      : CompileTimeValue(ValueKind::Array), Elements(Elements) {}
+
+  static bool classof(const CompileTimeValue *T) {
+    return T->getKind() == ValueKind::Array;
+  }
+  std::vector<std::shared_ptr<CompileTimeValue>> getElements() const {
+    return Elements;
+  }
+
+private:
+  std::vector<std::shared_ptr<CompileTimeValue>> Elements;
+};
+
 /// A representation of an arbitrary value that does not fall under
 /// any of the above categories.
 class RuntimeValue : public CompileTimeValue {

lib/ConstExtract/ConstExtract.cpp

@@ -123,7 +123,6 @@ parseProtocolListFromFile(StringRef protocolListFilePath,
 static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
   if (expr) {
     switch (expr->getKind()) {
-    case ExprKind::Array:
     case ExprKind::Dictionary:
 
     case ExprKind::BooleanLiteral:
@@ -140,6 +139,15 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
       break;
     }
 
+    case ExprKind::Array: {
+      auto arrayExpr = cast<ArrayExpr>(expr);
+      std::vector<std::shared_ptr<CompileTimeValue>> elementValues;
+      for (const auto elementExpr : arrayExpr->getElements()) {
+        elementValues.push_back(extractCompileTimeValue(elementExpr));
+      }
+      return std::make_shared<ArrayValue>(elementValues);
+    }
+
     case ExprKind::Tuple: {
       auto tupleExpr = cast<TupleExpr>(expr);
 
@@ -190,6 +198,11 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
       break;
     }
 
+    case ExprKind::Erasure: {
+      auto erasureExpr = cast<ErasureExpr>(expr);
+      return extractCompileTimeValue(erasureExpr->getSubExpr());
+    }
+
     default: {
       break;
     }
@@ -373,6 +386,18 @@ void writeValue(llvm::json::OStream &JSON,
     break;
   }
 
+  case CompileTimeValue::ValueKind::Array: {
+    auto arrayValue = cast<ArrayValue>(value);
+
+    JSON.attribute("valueKind", "Array");
+    JSON.attributeArray("value", [&] {
+      for (auto CTP : arrayValue->getElements()) {
+        JSON.object([&] { writeValue(JSON, CTP); });
+      }
+    });
+    break;
+  }
+
   case CompileTimeValue::ValueKind::Runtime: {
     JSON.attribute("valueKind", "Runtime");
     break;

test/ConstExtraction/ExtractGroups.swift

@@ -14,8 +14,61 @@
 // CHECK-NEXT:        "type": "[Swift.Int]",
 // CHECK-NEXT:        "isStatic": "false",
 // CHECK-NEXT:        "isComputed": "false",
-// CHECK-NEXT:        "valueKind": "RawLiteral",
-// CHECK-NEXT:        "value": "[1, 2, 3]"
+// CHECK-NEXT:        "valueKind": "Array",
+// CHECK-NEXT:        "value": [
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "1"
+// CHECK-NEXT:          },
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "2"
+// CHECK-NEXT:          },
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "3"
+// CHECK-NEXT:          }
+// CHECK-NEXT:        ]
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "label": "array2",
+// CHECK-NEXT:        "type": "[ExtractGroups.Foo]",
+// CHECK-NEXT:        "isStatic": "false",
+// CHECK-NEXT:        "isComputed": "false",
+// CHECK-NEXT:        "valueKind": "Array",
+// CHECK-NEXT:        "value": [
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "InitCall",
+// CHECK-NEXT:            "value": {
+// CHECK-NEXT:              "type": "ExtractGroups.Bar",
+// CHECK-NEXT:              "arguments": []
+// CHECK-NEXT:            }
+// CHECK-NEXT:          },
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "1"
+// CHECK-NEXT:          },
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "\"hi\""
+// CHECK-NEXT:          }
+// CHECK-NEXT:        ]
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "label": "array3",
+// CHECK-NEXT:        "type": "[ExtractGroups.Bar]",
+// CHECK-NEXT:        "isStatic": "false",
+// CHECK-NEXT:        "isComputed": "false",
+// CHECK-NEXT:        "valueKind": "Array",
+// CHECK-NEXT:        "value": [
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "valueKind": "InitCall",
+// CHECK-NEXT:            "value": {
+// CHECK-NEXT:              "type": "ExtractGroups.Bar",
+// CHECK-NEXT:              "arguments": []
+// CHECK-NEXT:            }
+// CHECK-NEXT:          }
+// CHECK-NEXT:        ]
 // CHECK-NEXT:      }
 // CHECK-NEXT:    ]
 // CHECK-NEXT:  },
@@ -96,6 +149,8 @@ protocol MyProto {}
 
 public struct Arrays : MyProto {
     let array1: [Int] = [1, 2, 3]
+    let array2: [Foo] = [Bar(), 1, "hi"]
+    let array3: [Bar] = [Bar()]
 }
 
 public struct Dictionaries : MyProto {
@@ -108,4 +163,7 @@ public struct Tuples : MyProto {
     let tuple3: Void = ()
 }
 
-public struct Bar {}
+public protocol Foo {}
+public struct Bar: Foo {}
+extension Int: Foo {}
+extension String: Foo {}