Sema: Allocate PreparedOverload in solver arena

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -67,7 +67,7 @@ class SyntacticElementTarget;
 
 // PreparedOverload.h
 struct DeclReferenceType;
-struct PreparedOverload;
+class PreparedOverload;
 struct PreparedOverloadBuilder;
 
 } // end namespace constraints
@@ -4906,18 +4906,23 @@ class ConstraintSystem {
   void recordResolvedOverload(ConstraintLocator *locator,
                               SelectedOverload choice);
 
+  /// Build and allocate a prepared overload in the solver arena.
+  const PreparedOverload *prepareOverload(ConstraintLocator *locator,
+                                          OverloadChoice choice,
+                                          DeclContext *useDC);
+
   /// Populate the prepared overload with all type variables and constraints
   /// that are to be introduced into the constraint system when this choice
   /// is taken.
   DeclReferenceType
-  prepareOverload(ConstraintLocator *locator,
-                  OverloadChoice choice,
-                  DeclContext *useDC,
-                  PreparedOverloadBuilder *preparedOverload);
+  prepareOverloadImpl(ConstraintLocator *locator,
+                      OverloadChoice choice,
+                      DeclContext *useDC,
+                      PreparedOverloadBuilder *preparedOverload);
 
   void replayChanges(
-      ConstraintLocatorBuilder locator,
-      PreparedOverload preparedOverload);
+      ConstraintLocator *locator,
+      const PreparedOverload *preparedOverload);
 
   /// Resolve the given overload set to the given choice.
   void resolveOverload(ConstraintLocator *locator, Type boundType,

include/swift/Sema/PreparedOverload.h

@@ -13,8 +13,10 @@
 #define SWIFT_SEMA_PREPAREDOVERLOAD_H
 
 #include "swift/AST/PropertyWrappers.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/TrailingObjects.h"
 
 namespace swift {
 
@@ -59,77 +61,118 @@ struct DeclReferenceType {
 /// variable.
 using OpenedType = std::pair<GenericTypeParamType *, TypeVariableType *>;
 
-/// A "pre-cooked" representation of all type variables and constraints
-/// that are generated as part of an overload choice.
-struct PreparedOverload {
-  /// A change to be introduced into the constraint system when this
-  /// overload choice is chosen.
-  struct Change {
-    enum ChangeKind : unsigned {
-      /// A generic parameter was opened to a type variable.
-      AddedTypeVariable,
+/// A change to be introduced into the constraint system when this
+/// overload choice is chosen.
+struct PreparedOverloadChange {
+  enum ChangeKind : unsigned {
+    /// A generic parameter was opened to a type variable.
+    AddedTypeVariable,
 
-      /// A generic requirement was opened to a constraint.
-      AddedConstraint,
+    /// A generic requirement was opened to a constraint.
+    AddedConstraint,
 
-      /// A mapping of generic parameter types to type variables
-      /// was recorded.
-      OpenedTypes,
+    /// A mapping of generic parameter types to type variables
+    /// was recorded.
+    OpenedTypes,
 
-      /// An existential type was opened.
-      OpenedExistentialType,
+    /// An existential type was opened.
+    OpenedExistentialType,
 
-      /// A pack expansion type was opened.
-      OpenedPackExpansionType,
+    /// A pack expansion type was opened.
+    OpenedPackExpansionType,
 
-      /// A property wrapper was applied to a parameter.
-      AppliedPropertyWrapper,
+    /// A property wrapper was applied to a parameter.
+    AppliedPropertyWrapper,
 
-      /// A fix was recorded because a property wrapper application failed.
-      AddedFix
-    };
+    /// A fix was recorded because a property wrapper application failed.
+    AddedFix
+  };
 
-    /// The kind of change.
-    ChangeKind Kind;
+  /// The kind of change.
+  ChangeKind Kind;
 
-    union {
-      /// For ChangeKind::AddedTypeVariable.
-      TypeVariableType *TypeVar;
+  union {
+    /// For ChangeKind::AddedTypeVariable.
+    TypeVariableType *TypeVar;
 
-      /// For ChangeKind::AddedConstraint.
-      Constraint *TheConstraint;
-
-      /// For ChangeKind::OpenedTypes.
-      struct {
-        const OpenedType *Data;
-        size_t Count;
-      } Replacements;
-
-      /// For ChangeKind::OpenedExistentialType.
-      ExistentialArchetypeType *TheExistential;
-
-      /// For ChangeKind::OpenedPackExpansionType.
-      struct {
-        PackExpansionType *TheExpansion;
-        TypeVariableType *TypeVar;
-      } PackExpansion;
-
-      /// For ChangeKind::AppliedPropertyWrapper.
-      struct {
-        TypeBase *WrapperType;
-        PropertyWrapperInitKind InitKind;
-      } PropertyWrapper;
-
-      /// For ChangeKind::Fix.
-      struct {
-        ConstraintFix *TheFix;
-        unsigned Impact;
-      } Fix;
-    };
+    /// For ChangeKind::AddedConstraint.
+    Constraint *TheConstraint;
+
+    /// For ChangeKind::OpenedTypes.
+    struct {
+      const OpenedType *Data;
+      size_t Count;
+    } Replacements;
+
+    /// For ChangeKind::OpenedExistentialType.
+    ExistentialArchetypeType *TheExistential;
+
+    /// For ChangeKind::OpenedPackExpansionType.
+    struct {
+      PackExpansionType *TheExpansion;
+      TypeVariableType *TypeVar;
+    } PackExpansion;
+
+    /// For ChangeKind::AppliedPropertyWrapper.
+    struct {
+      TypeBase *WrapperType;
+      PropertyWrapperInitKind InitKind;
+    } PropertyWrapper;
+
+    /// For ChangeKind::Fix.
+    struct {
+      ConstraintFix *TheFix;
+      unsigned Impact;
+    } Fix;
   };
+};
+
+/// A "pre-cooked" representation of all type variables and constraints
+/// that are generated as part of an overload choice.
+class PreparedOverload final :
+  public llvm::TrailingObjects<PreparedOverload, PreparedOverloadChange> {
 
-  ArrayRef<Change> Changes;
+public:
+  using Change = PreparedOverloadChange;
+
+private:
+  size_t Count;
   DeclReferenceType DeclType;
+
+  size_t numTrailingObjects(OverloadToken<Change>) const {
+    return Count;
+  }
+
+public:
+  PreparedOverload(const DeclReferenceType &declType, ArrayRef<Change> changes)
+    : Count(changes.size()), DeclType(declType) {
+    std::uninitialized_copy(changes.begin(), changes.end(),
+                            getTrailingObjects<Change>());
+  }
+
+  Type getOpenedType() const {
+    return DeclType.openedType;
+  }
+
+  Type getAdjustedOpenedType() const {
+    return DeclType.adjustedOpenedType;
+  }
+
+  Type getReferenceType() const {
+    return DeclType.referenceType;
+  }
+
+  Type getAdjustedReferenceType() const {
+    return DeclType.adjustedReferenceType;
+  }
+
+  Type getThrownErrorTypeOnAccess() const {
+    return DeclType.thrownErrorTypeOnAccess;
+  }
+
+  ArrayRef<Change> getChanges() const {
+    return ArrayRef<Change>(getTrailingObjects<Change>(), Count);
+  }
 };
 
 struct PreparedOverloadBuilder {

lib/Sema/TypeOfReference.cpp

@@ -2498,9 +2498,9 @@ void ConstraintSystem::recordResolvedOverload(ConstraintLocator *locator,
 }
 
 void ConstraintSystem::replayChanges(
-    ConstraintLocatorBuilder locator,
-    PreparedOverload preparedOverload) {
-  for (auto change : preparedOverload.Changes) {
+    ConstraintLocator *locator,
+    const PreparedOverload *preparedOverload) {
+  for (auto change : preparedOverload->getChanges()) {
     switch (change.Kind) {
     case PreparedOverload::Change::AddedTypeVariable:
       addTypeVariable(change.TypeVar);
@@ -2511,20 +2511,18 @@ void ConstraintSystem::replayChanges(
       break;
 
     case PreparedOverload::Change::OpenedTypes: {
-      auto *locatorPtr = getConstraintLocator(locator);
       ArrayRef<OpenedType> replacements(
           change.Replacements.Data,
           change.Replacements.Count);
 
       // FIXME: Get rid of this conditional.
-      if (getOpenedTypes(locatorPtr).empty())
-        recordOpenedType(locatorPtr, replacements);
+      if (getOpenedTypes(locator).empty())
+        recordOpenedType(locator, replacements);
       break;
     }
 
     case PreparedOverload::Change::OpenedExistentialType: {
-      auto *locatorPtr = getConstraintLocator(locator);
-      recordOpenedExistentialType(locatorPtr,
+      recordOpenedExistentialType(locator,
                                   change.TheExistential);
       break;
     }
@@ -2536,8 +2534,7 @@ void ConstraintSystem::replayChanges(
       break;
 
     case PreparedOverload::Change::AppliedPropertyWrapper: {
-      auto *locatorPtr = getConstraintLocator(locator);
-      Expr *anchor = getAsExpr(locatorPtr->getAnchor());
+      Expr *anchor = getAsExpr(locator->getAnchor());
       applyPropertyWrapper(anchor,
           { Type(change.PropertyWrapper.WrapperType),
             change.PropertyWrapper.InitKind });
@@ -2558,10 +2555,10 @@ void ConstraintSystem::replayChanges(
 /// FIXME: As a transitional mechanism, if preparedOverload is nullptr, this
 /// immediately performs all operations.
 DeclReferenceType
-ConstraintSystem::prepareOverload(ConstraintLocator *locator,
-                                  OverloadChoice choice,
-                                  DeclContext *useDC,
-                                  PreparedOverloadBuilder *preparedOverload) {
+ConstraintSystem::prepareOverloadImpl(ConstraintLocator *locator,
+                                      OverloadChoice choice,
+                                      DeclContext *useDC,
+                                      PreparedOverloadBuilder *preparedOverload) {
   // If we refer to a top-level decl with special type-checking semantics,
   // handle it now.
   auto semantics =
@@ -2584,6 +2581,25 @@ ConstraintSystem::prepareOverload(ConstraintLocator *locator,
   }
 }
 
+const PreparedOverload *
+ConstraintSystem::prepareOverload(ConstraintLocator *locator,
+                                  OverloadChoice choice,
+                                  DeclContext *useDC) {
+  ASSERT(!PreparingOverload);
+  PreparingOverload = true;
+
+  PreparedOverloadBuilder builder;
+  auto declRefType = prepareOverloadImpl(locator, choice, useDC, &builder);
+
+  PreparingOverload = false;
+
+  size_t count = builder.Changes.size();
+  auto size = PreparedOverload::totalSizeToAlloc<PreparedOverload::Change>(count);
+  auto mem = Allocator.Allocate(size, alignof(PreparedOverload));
+
+  return new (mem) PreparedOverload(declRefType, builder.Changes);
+}
+
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType, OverloadChoice choice,
                                        DeclContext *useDC) {
@@ -2602,30 +2618,27 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   case OverloadChoiceKind::DynamicMemberLookup:
   case OverloadChoiceKind::KeyPathDynamicMemberLookup: {
     // FIXME: Transitional hack
-    bool enablePreparedOverloads = false;
+    bool enablePreparedOverloads = true;
 
-    PreparedOverloadBuilder preparedOverload;
     if (enablePreparedOverloads) {
-      ASSERT(!PreparingOverload);
-      PreparingOverload = true;
-    }
-
-    auto declRefType = prepareOverload(locator, choice, useDC,
-                                       (enablePreparedOverloads
-                                          ? &preparedOverload
-                                          : nullptr));
+      auto *preparedOverload = prepareOverload(locator, choice, useDC);
+      replayChanges(locator, preparedOverload);
+
+      openedType = preparedOverload->getOpenedType();
+      adjustedOpenedType = preparedOverload->getAdjustedOpenedType();
+      refType = preparedOverload->getReferenceType();
+      adjustedRefType = preparedOverload->getAdjustedReferenceType();
+      thrownErrorTypeOnAccess = preparedOverload->getThrownErrorTypeOnAccess();
+    } else {
+      auto declRefType = prepareOverloadImpl(locator, choice, useDC, nullptr);
 
-    if (enablePreparedOverloads) {
-      PreparingOverload = false;
-      PreparedOverload result{preparedOverload.Changes, declRefType};
-      replayChanges(locator, result);
+      openedType = declRefType.openedType;
+      adjustedOpenedType = declRefType.adjustedOpenedType;
+      refType = declRefType.referenceType;
+      adjustedRefType = declRefType.adjustedReferenceType;
+      thrownErrorTypeOnAccess = declRefType.thrownErrorTypeOnAccess;
     }
 
-    openedType = declRefType.openedType;
-    adjustedOpenedType = declRefType.adjustedOpenedType;
-    refType = declRefType.referenceType;
-    adjustedRefType = declRefType.adjustedReferenceType;
-    thrownErrorTypeOnAccess = declRefType.thrownErrorTypeOnAccess;
     break;
   }