Merge pull request #75219 from venkatesh5789/result-builder-extraction

[Compile Time Constant Extraction] Extract Result Builder Expressions

include/swift/AST/ConstTypeInfo.h

@@ -96,17 +96,155 @@ private:
   std::vector<FunctionParameter> Parameters;
 };
 
-/// A representation of a Builder pattern initialization expression
+/// A representation of a Builder pattern initialization expression. For
+/// example:
+///
+/// @FooBuilder
+/// public static var foos: [Foo] {
+///    Foo(name: "foos.1")
+///    Foo(name: "foos.2")
+/// }
+///
+/// In this example, the result builder type is FooBuilder
+/// The members are Foo(name: "foos.1") and Foo(name: "foos.2")
+///
 class BuilderValue : public CompileTimeValue {
 public:
-  BuilderValue() : CompileTimeValue(ValueKind::Builder) {}
+  enum MemberKind {
+    Expression,
+    Either,
+    Optional,
+    LimitedAvailability,
+    Array,
+    Unknown
+  };
+
+  /// A base class for individual members being declared inside the result
+  /// builder
+  class BuilderMember {
+  public:
+    MemberKind getKind() const { return Kind; }
+
+  protected:
+    BuilderMember(MemberKind MemberKind) : Kind(MemberKind) {}
+
+  private:
+    MemberKind Kind;
+  };
+
+  /// A basic expression that is defined inside the result builder. For example:
+  ///  {
+  ///     Foo(name: "1")
+  ///  }
+  ///
+  class SingleMember : public BuilderMember {
+  public:
+    SingleMember(std::shared_ptr<CompileTimeValue> Element)
+        : BuilderMember(MemberKind::Expression), Element(Element) {}
+
+    static bool classof(const BuilderMember *T) {
+      return T->getKind() == MemberKind::Expression;
+    }
+
+    std::shared_ptr<CompileTimeValue> getElement() const { return Element; }
+
+  private:
+    std::shared_ptr<CompileTimeValue> Element;
+  };
+
+  /// A member that represents when the individual values are defined by
+  /// iterating over an array. For example:
+  ///     for i in 1...3 {
+  ///         Foo(name: "MyFooProviderInferred.foos.Array.\(i)")
+  ///     }
+  ///
+  class ArrayMember : public BuilderMember {
+  public:
+    ArrayMember(std::vector<std::shared_ptr<BuilderMember>> Elements)
+        : BuilderMember(MemberKind::Array), Elements(Elements) {}
+
+    static bool classof(const BuilderMember *T) {
+      return T->getKind() == MemberKind::Array;
+    }
+
+    std::vector<std::shared_ptr<BuilderMember>> getElements() const {
+      return Elements;
+    }
+
+  private:
+    std::vector<std::shared_ptr<BuilderMember>> Elements;
+  };
+
+  /// A member that is defined conditionally. It can be of the following types:
+  ///
+  /// 1. A regular if-else condition
+  ///      if condition {
+  ///         Foo(name: "1")
+  ///      } else {
+  ///         Foo(name: "2")
+  ///      }
+  ///
+  /// 2. An optional
+  ///     if condition {
+  ///         Foo(name: "1")
+  ///     }
+  ///
+  /// 3. Limited availability
+  ///     if #available(macOS 99, *) {
+  ///         Foo(name: "1")
+  ///         Foo(name: "2")
+  ///     }
+  ///
+  class ConditionalMember : public BuilderMember {
+  public:
+    ConditionalMember(MemberKind MemberKind,
+                      std::vector<std::shared_ptr<BuilderMember>> IfElements,
+                      std::vector<std::shared_ptr<BuilderMember>> ElseElements)
+        : BuilderMember(MemberKind), IfElements(IfElements),
+          ElseElements(ElseElements) {}
+
+    static bool classof(const BuilderMember *T) {
+      auto Kind = T->getKind();
+      return (Kind == MemberKind::Either) ||
+             (Kind == MemberKind::LimitedAvailability) ||
+             (Kind == MemberKind::Optional);
+    }
+
+    std::vector<std::shared_ptr<BuilderMember>> getIfElements() const {
+      return IfElements;
+    }
+    std::vector<std::shared_ptr<BuilderMember>> getElseElements() const {
+      return ElseElements;
+    }
+
+  private:
+    std::vector<std::shared_ptr<BuilderMember>> IfElements;
+    std::vector<std::shared_ptr<BuilderMember>> ElseElements;
+  };
+
+  BuilderValue(std::vector<std::shared_ptr<BuilderMember>> Members)
+      : CompileTimeValue(ValueKind::Builder), ResultBuilderType(std::nullopt),
+        Members(Members) {}
+
+  BuilderValue(CustomAttr *ResultBuilderType,
+               std::vector<std::shared_ptr<BuilderMember>> Members)
+      : CompileTimeValue(ValueKind::Builder),
+        ResultBuilderType(ResultBuilderType), Members(Members) {}
+
+  std::optional<CustomAttr *> getResultBuilderType() const {
+    return ResultBuilderType;
+  }
+  std::vector<std::shared_ptr<BuilderMember>> getMembers() const {
+    return Members;
+  }
 
   static bool classof(const CompileTimeValue *T) {
     return T->getKind() == ValueKind::Builder;
   }
 
 private:
-  std::vector<CompileTimeValue> Members;
+  std::optional<CustomAttr *> ResultBuilderType;
+  std::vector<std::shared_ptr<BuilderMember>> Members;
 };
 
 struct TupleElement {