[CS] NFC: Store ContextualTypeInfo in SyntacticElementTarget

Move the contextual type locator onto
ContextualTypeInfo, and consolidate the separate
fields in SyntacticElementTarget into storing a
ContextualTypeInfo. This then lets us plumb down
the locator for the branch contextual type of an
if/switch expression from the initial constraint
generation, rather than introducing it later. This
should be NFC.

Author: hamishknight

include/swift/Sema/Constraint.h

@@ -25,6 +25,7 @@
 #include "swift/AST/TypeLoc.h"
 #include "swift/Basic/Debug.h"
 #include "swift/Sema/ConstraintLocator.h"
+#include "swift/Sema/ContextualTypeInfo.h"
 #include "swift/Sema/OverloadChoice.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ilist.h"
@@ -50,23 +51,6 @@ class ConstraintLocator;
 class ConstraintSystem;
 enum class TrailingClosureMatching;
 
-/// Describes contextual type information about a particular element
-/// (expression, statement etc.) within a constraint system.
-struct ContextualTypeInfo {
-  TypeLoc typeLoc;
-  ContextualTypePurpose purpose;
-
-  ContextualTypeInfo() : typeLoc(TypeLoc()), purpose(CTP_Unused) {}
-
-  ContextualTypeInfo(Type contextualTy, ContextualTypePurpose purpose)
-      : typeLoc(TypeLoc::withoutLoc(contextualTy)), purpose(purpose) {}
-
-  ContextualTypeInfo(TypeLoc typeLoc, ContextualTypePurpose purpose)
-      : typeLoc(typeLoc), purpose(purpose) {}
-
-  Type getType() const { return typeLoc.getType(); }
-};
-
 /// Describes the kind of constraint placed on one or more types.
 enum class ConstraintKind : char {
   /// The two types must be bound to the same type. This is the only

include/swift/Sema/ConstraintSystem.h

@@ -3114,12 +3114,11 @@ class ConstraintSystem {
     return E;
   }
 
-  void setContextualType(ASTNode node, TypeLoc T,
-                         ContextualTypePurpose purpose) {
+  void setContextualInfo(ASTNode node, ContextualTypeInfo info) {
     assert(bool(node) && "Expected non-null expression!");
     assert(contextualTypes.count(node) == 0 &&
            "Already set this contextual type");
-    contextualTypes[node] = {{T, purpose}, Type()};
+    contextualTypes[node] = {info, Type()};
   }
 
   llvm::Optional<ContextualTypeInfo> getContextualTypeInfo(ASTNode node) const {

include/swift/Sema/ContextualTypeInfo.h

@@ -0,0 +1,55 @@
+//===--- ContextualTypeInfo.h - Contextual Type Info ------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides the \c ContextualTypeInfo class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SEMA_CONTEXTUAL_TYPE_INFO_H
+#define SWIFT_SEMA_CONTEXTUAL_TYPE_INFO_H
+
+#include "swift/AST/TypeLoc.h"
+#include "swift/Sema/ConstraintLocator.h"
+
+namespace swift {
+namespace constraints {
+
+/// Describes contextual type information about a particular element
+/// (expression, statement etc.) within a constraint system.
+struct ContextualTypeInfo {
+  TypeLoc typeLoc;
+  ContextualTypePurpose purpose;
+
+  /// The locator for the contextual type conversion constraint, or
+  /// \c nullptr to use the default locator which is anchored directly on
+  /// the expression.
+  ConstraintLocator *locator;
+
+  ContextualTypeInfo()
+      : typeLoc(TypeLoc()), purpose(CTP_Unused), locator(nullptr) {}
+
+  ContextualTypeInfo(Type contextualTy, ContextualTypePurpose purpose,
+                     ConstraintLocator *locator = nullptr)
+      : typeLoc(TypeLoc::withoutLoc(contextualTy)), purpose(purpose),
+        locator(locator) {}
+
+  ContextualTypeInfo(TypeLoc typeLoc, ContextualTypePurpose purpose,
+                     ConstraintLocator *locator = nullptr)
+      : typeLoc(typeLoc), purpose(purpose), locator(locator) {}
+
+  Type getType() const { return typeLoc.getType(); }
+};
+
+} // end namespace constraints
+} // end namespace swift
+
+#endif // SWIFT_SEMA_CONTEXTUAL_TYPE_INFO_H

include/swift/Sema/SyntacticElementTarget.h

@@ -23,6 +23,7 @@
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeLoc.h"
 #include "swift/Sema/ConstraintLocator.h"
+#include "swift/Sema/ContextualTypeInfo.h"
 
 namespace swift {
 
@@ -71,18 +72,8 @@ class SyntacticElementTarget {
       /// type-checked.
       DeclContext *dc;
 
-      // TODO: Fold the 3 below fields into ContextualTypeInfo
-
-      /// The purpose of the contextual type.
-      ContextualTypePurpose contextualPurpose;
-
-      /// The type to which the expression should be converted.
-      TypeLoc convertType;
-
-      /// The locator for the contextual type conversion constraint, or
-      /// \c nullptr to use the default locator which is anchored directly on
-      /// the expression.
-      ConstraintLocator *convertTypeLocator;
+      /// The contextual type info for the expression.
+      ContextualTypeInfo contextualInfo;
 
       /// When initializing a pattern from the expression, this is the
       /// pattern.
@@ -169,26 +160,15 @@ class SyntacticElementTarget {
   void maybeApplyPropertyWrapper();
 
 public:
-  SyntacticElementTarget(Expr *expr, DeclContext *dc,
-                         ContextualTypePurpose contextualPurpose,
-                         Type convertType,
-                         ConstraintLocator *convertTypeLocator,
-                         bool isDiscarded)
-      : SyntacticElementTarget(expr, dc, contextualPurpose,
-                               TypeLoc::withoutLoc(convertType),
-                               convertTypeLocator, isDiscarded) {}
-
   SyntacticElementTarget(Expr *expr, DeclContext *dc,
                          ContextualTypePurpose contextualPurpose,
                          Type convertType, bool isDiscarded)
-      : SyntacticElementTarget(expr, dc, contextualPurpose, convertType,
-                               /*convertTypeLocator*/ nullptr, isDiscarded) {}
+      : SyntacticElementTarget(
+            expr, dc, ContextualTypeInfo(convertType, contextualPurpose),
+            isDiscarded) {}
 
   SyntacticElementTarget(Expr *expr, DeclContext *dc,
-                         ContextualTypePurpose contextualPurpose,
-                         TypeLoc convertType,
-                         ConstraintLocator *convertTypeLocator,
-                         bool isDiscarded);
+                         ContextualTypeInfo contextualInfo, bool isDiscarded);
 
   SyntacticElementTarget(ClosureExpr *closure, Type convertType) {
     kind = Kind::closure;
@@ -375,26 +355,31 @@ class SyntacticElementTarget {
     llvm_unreachable("invalid decl context type");
   }
 
-  ContextualTypePurpose getExprContextualTypePurpose() const {
+  /// Get the contextual type info for an expression target.
+  ContextualTypeInfo getExprContextualTypeInfo() const {
     assert(kind == Kind::expression);
-    return expression.contextualPurpose;
+    return expression.contextualInfo;
+  }
+
+  /// Get the contextual type purpose for an expression target.
+  ContextualTypePurpose getExprContextualTypePurpose() const {
+    return getExprContextualTypeInfo().purpose;
   }
 
+  /// Get the contextual type for an expression target.
   Type getExprContextualType() const {
     return getExprContextualTypeLoc().getType();
   }
 
+  /// Get the contextual type for an expression target.
   TypeLoc getExprContextualTypeLoc() const {
-    assert(kind == Kind::expression);
-
     // For an @autoclosure parameter, the conversion type is
     // the result of the function type.
-    if (FunctionType *autoclosureParamType = getAsAutoclosureParamType()) {
-      return TypeLoc(expression.convertType.getTypeRepr(),
-                     autoclosureParamType->getResult());
-    }
+    auto typeLoc = getExprContextualTypeInfo().typeLoc;
+    if (FunctionType *autoclosureParamType = getAsAutoclosureParamType())
+      return TypeLoc(typeLoc.getTypeRepr(), autoclosureParamType->getResult());
 
-    return expression.convertType;
+    return typeLoc;
   }
 
   /// Retrieve the type to which an expression should be converted, or
@@ -408,33 +393,32 @@ class SyntacticElementTarget {
   /// Retrieve the conversion type locator for the expression, or \c nullptr
   /// if it has not been set.
   ConstraintLocator *getExprConvertTypeLocator() const {
-    assert(kind == Kind::expression);
-    return expression.convertTypeLocator;
+    return getExprContextualTypeInfo().locator;
   }
 
   /// Returns the autoclosure parameter type, or \c nullptr if the
   /// expression has a different kind of context.
   FunctionType *getAsAutoclosureParamType() const {
-    assert(kind == Kind::expression);
-    if (expression.contextualPurpose == CTP_AutoclosureDefaultParameter)
-      return expression.convertType.getType()->castTo<FunctionType>();
+    if (getExprContextualTypePurpose() == CTP_AutoclosureDefaultParameter)
+      return getExprContextualTypeInfo().getType()->castTo<FunctionType>();
+
     return nullptr;
   }
 
   void setExprConversionType(Type type) {
     assert(kind == Kind::expression);
-    expression.convertType = TypeLoc::withoutLoc(type);
+    expression.contextualInfo.typeLoc = TypeLoc::withoutLoc(type);
   }
 
   void setExprConversionTypeLoc(TypeLoc type) {
     assert(kind == Kind::expression);
-    expression.convertType = type;
+    expression.contextualInfo.typeLoc = type;
   }
 
   /// Whether this target is for an initialization expression and pattern.
   bool isForInitialization() const {
     return kind == Kind::expression &&
-           expression.contextualPurpose == CTP_Initialization;
+           getExprContextualTypePurpose() == CTP_Initialization;
   }
 
   /// For a pattern initialization target, retrieve the pattern.
@@ -448,7 +432,7 @@ class SyntacticElementTarget {
 
   ExprPattern *getExprPattern() const {
     assert(kind == Kind::expression);
-    assert(expression.contextualPurpose == CTP_ExprPattern);
+    assert(getExprContextualTypePurpose() == CTP_ExprPattern);
     return cast<ExprPattern>(expression.pattern);
   }
 
@@ -575,11 +559,16 @@ class SyntacticElementTarget {
       return;
     }
 
-    assert(kind == Kind::expression);
-    assert(expression.contextualPurpose == CTP_Initialization ||
-           expression.contextualPurpose == CTP_ForEachStmt ||
-           expression.contextualPurpose == CTP_ForEachSequence ||
-           expression.contextualPurpose == CTP_ExprPattern);
+    switch (getExprContextualTypePurpose()) {
+    case CTP_Initialization:
+    case CTP_ForEachStmt:
+    case CTP_ForEachSequence:
+    case CTP_ExprPattern:
+      break;
+    default:
+      assert(false && "Unexpected contextual type purpose");
+      break;
+    }
     expression.pattern = pattern;
   }
 

lib/Sema/CSGen.cpp

@@ -4334,8 +4334,9 @@ bool ConstraintSystem::generateWrappedPropertyTypeConstraints(
       ConstraintKind::Equal, propertyType, wrappedValueType,
       getConstraintLocator(
           wrappedVar, LocatorPathElt::ContextualType(CTP_WrappedProperty)));
-  setContextualType(wrappedVar, TypeLoc::withoutLoc(wrappedValueType),
-                    CTP_WrappedProperty);
+
+  ContextualTypeInfo contextInfo(wrappedValueType, CTP_WrappedProperty);
+  setContextualInfo(wrappedVar, contextInfo);
   return false;
 }
 
@@ -4419,10 +4420,9 @@ generateForEachStmtConstraints(ConstraintSystem &cs,
         makeIteratorCall, dc, /*patternType=*/Type(), PB, /*index=*/0,
         /*shouldBindPatternsOneWay=*/false);
 
-    cs.setContextualType(
-        sequenceExpr,
-        TypeLoc::withoutLoc(sequenceProto->getDeclaredInterfaceType()),
-        CTP_ForEachSequence);
+    ContextualTypeInfo contextInfo(sequenceProto->getDeclaredInterfaceType(),
+                                   CTP_ForEachSequence);
+    cs.setContextualInfo(sequenceExpr, contextInfo);
 
     if (cs.generateConstraints(makeIteratorTarget))
       return llvm::None;
@@ -4480,10 +4480,9 @@ generateForEachStmtConstraints(ConstraintSystem &cs,
       if (!iteratorProto)
         return llvm::None;
 
-      cs.setContextualType(
-          nextRef->getBase(),
-          TypeLoc::withoutLoc(iteratorProto->getDeclaredInterfaceType()),
-          CTP_ForEachSequence);
+      ContextualTypeInfo contextInfo(iteratorProto->getDeclaredInterfaceType(),
+                                     CTP_ForEachSequence);
+      cs.setContextualInfo(nextRef->getBase(), contextInfo);
     }
 
     SyntacticElementTarget nextTarget(nextCall, dc, CTP_Unused,
@@ -4549,8 +4548,9 @@ generateForEachStmtConstraints(ConstraintSystem &cs,
       return llvm::None;
 
     cs.setTargetFor(whereExpr, whereTarget);
-    cs.setContextualType(whereExpr, TypeLoc::withoutLoc(boolType),
-                         CTP_Condition);
+
+    ContextualTypeInfo contextInfo(boolType, CTP_Condition);
+    cs.setContextualInfo(whereExpr, contextInfo);
   }
 
   // Populate all of the information for a for-each loop.

lib/Sema/CSSolver.cpp

@@ -323,10 +323,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Add the contextual types.
   for (const auto &contextualType : solution.contextualTypes) {
-    if (!getContextualTypeInfo(contextualType.first)) {
-      setContextualType(contextualType.first, contextualType.second.typeLoc,
-                        contextualType.second.purpose);
-    }
+    if (!getContextualTypeInfo(contextualType.first))
+      setContextualInfo(contextualType.first, contextualType.second);
   }
 
   // Register the statement condition targets.
@@ -1731,8 +1729,7 @@ bool ConstraintSystem::solveForCodeCompletion(
     SyntacticElementTarget &target, SmallVectorImpl<Solution> &solutions) {
   if (auto *expr = target.getAsExpr()) {
     // Tell the constraint system what the contextual type is.
-    setContextualType(expr, target.getExprContextualTypeLoc(),
-                      target.getExprContextualTypePurpose());
+    setContextualInfo(expr, target.getExprContextualTypeInfo());
 
     // Set up the expression type checker timer.
     Timer.emplace(expr, *this);