WIP full types throws inference

DougGregor · DougGregor · commit 8b1f3680ae3f · 2024-02-16T20:18:03.000-08:00
diff --git a/include/swift/Sema/ConstraintLocator.h b/include/swift/Sema/ConstraintLocator.h
@@ -337,6 +337,13 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     return false;
   }
 
+  /// Determine whether this locator points directly to a given statement.
+  template <typename E> bool directlyAtStmt() const {
+    if (auto *stmt = getAnchor().dyn_cast<Stmt *>())
+      return isa<E>(stmt) && getPath().empty();
+    return false;
+  }
+
   /// Check whether the first element in the path of this locator (if any)
   /// is a given \c LocatorPathElt subclass.
   template <class T>
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -3434,17 +3434,13 @@ class ConstraintSystem {
       PotentialThrowSite::Kind kind, Type type,
       ConstraintLocatorBuilder locator);
 
-  /// Determine the caught error type for the given catch node.
+  /// Retrieve the caught error type for the given catch node, which could be
+  /// a type variable if the caught error type is going to be inferred.
   Type getCaughtErrorType(CatchNode node);
 
-  /// Infer the caught error type for this catch node, or introduce an
-  /// appropriate type variable to describe that caught error type if
-  /// it cannot be computed yet.
-  Type inferCaughtErrorType(CatchNode node);
-
-  /// Return the type variable that represents the inferred thrown error
-  /// type for this closure, or NULL if the thrown error type is not inferred.
-  TypeVariableType *getInferredThrownError(ClosureExpr *closure);
+  /// Finalize the caught error type type once all of the potential throw
+  /// sites are known.
+  void finalizeCaughtErrorType(CatchNode node);
 
   /// Retrieve the constraint locator for the given anchor and
   /// path, uniqued.
@@ -6474,6 +6470,9 @@ class TypeVarRefCollector : public ASTWalker {
   /// Infer the referenced type variables from a given decl.
   void inferTypeVars(Decl *D);
 
+  /// Infer the referenced type variables from a type.
+  void inferTypeVars(Type type);
+
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Arguments;
   }
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -865,8 +865,12 @@ void TypeVarRefCollector::inferTypeVars(Decl *D) {
   if (!ty)
     return;
 
+  inferTypeVars(ty);
+}
+
+void TypeVarRefCollector::inferTypeVars(Type type) {
   SmallPtrSet<TypeVariableType *, 4> typeVars;
-  ty->getTypeVariables(typeVars);
+  type->getTypeVariables(typeVars);
   TypeVars.insert(typeVars.begin(), typeVars.end());
 }
 
@@ -915,10 +919,12 @@ TypeVarRefCollector::walkToStmtPre(Stmt *stmt) {
     if (isa<ReturnStmt>(stmt) && DCDepth == 0 &&
         !Locator->directlyAt<ClosureExpr>()) {
       SmallPtrSet<TypeVariableType *, 4> typeVars;
-      CS.getClosureType(CE)->getResult()->getTypeVariables(typeVars);
+      auto closureType = CS.getClosureType(CE);
+      closureType->getResult()->getTypeVariables(typeVars);
+
+      if (auto thrownErrorType = closureType->getThrownError())
+        thrownErrorType->getTypeVariables(typeVars);
 
-      FIXME if we're doing full typed throws, also look at the thrown
-        error type here?
       TypeVars.insert(typeVars.begin(), typeVars.end());
     }
   }
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -13840,11 +13840,11 @@ ConstraintSystem::simplifyCaughtErrorConstraint(
     CatchNode catchNode,
     TypeMatchOptions flags,
     ConstraintLocatorBuilder locator) {
-  Type caughtErrorType = inferCaughtErrorType(catchNode);
-  if (caughtErrorType->isEqual(type))
+  // Keep the constraint around until it simplifies beyond a type variable.
+  Type simplified = simplifyType(type);
+  if (simplified->is<TypeVariableType>())
     return SolutionKind::Unsolved;
 
-  addConstraint(ConstraintKind::Bind, type, caughtErrorType, locator);
   return SolutionKind::Solved;
 }
 
diff --git a/lib/Sema/CSStep.cpp b/lib/Sema/CSStep.cpp
@@ -970,6 +970,18 @@ StepResult ConjunctionStep::resume(bool prevFailed) {
     if (HadFailure)
       return done(/*isSuccess=*/false);
 
+    auto locator = Conjunction->getLocator();
+    if (auto syntacticElement =
+          locator->getLastElementAs<LocatorPathElt::SyntacticElement>()) {
+      // If we are at a do..catch, and we're inferring a throw error type for
+      // the do..catch, we're now in a position to finalize the thrown type.
+      if (auto doCatch = dyn_cast_or_null<DoCatchStmt>(
+                           syntacticElement->getElement().dyn_cast<Stmt *>())) {
+        if (CS.getCaughtErrorType(doCatch)->is<TypeVariableType>())
+          CS.finalizeCaughtErrorType(doCatch);
+      }
+    }
+
     // If this was an isolated conjunction solver needs to do
     // the following:
     //
@@ -1084,6 +1096,20 @@ void ConjunctionStep::restoreOuterState(const Score &solutionScore) const {
 void ConjunctionStep::SolverSnapshot::applySolution(const Solution &solution) {
   CS.applySolution(solution);
 
+  // If we are at a closure, and the closure type has a type variable for
+  // its thrown error type, we're now in a position to finalize that type.
+  auto locator = Conjunction->getLocator();
+  if (locator->directlyAt<ClosureExpr>()) {
+    auto closureTy =
+        CS.getClosureType(castToExpr<ClosureExpr>(locator->getAnchor()));
+    if (auto thrownError = closureTy->getEffectiveThrownErrorTypeOrNever()) {
+      if (thrownError->is<TypeVariableType>()) {
+        auto *closure = castToExpr<ClosureExpr>(locator->getAnchor());
+        CS.finalizeCaughtErrorType(closure);
+      }
+    }
+  }
+
   if (!CS.shouldAttemptFixes())
     return;
 
@@ -1096,7 +1122,6 @@ void ConjunctionStep::SolverSnapshot::applySolution(const Solution &solution) {
   // has failed, let's bind all of unresolved type variables
   // in its interface type to holes to avoid extraneous
   // fixes produced by outer context.
-  auto locator = Conjunction->getLocator();
   if (locator->directlyAt<ClosureExpr>()) {
     auto closureTy =
         CS.getClosureType(castToExpr<ClosureExpr>(locator->getAnchor()));
diff --git a/lib/Sema/CSSyntacticElement.cpp b/lib/Sema/CSSyntacticElement.cpp
@@ -287,23 +287,6 @@ static bool isViableElement(ASTNode element,
 using ElementInfo = std::tuple<ASTNode, ContextualTypeInfo,
                                /*isDiscarded=*/bool, ConstraintLocator *>;
 
-static TypeVariableType *assignClosureThrownErrorType(
-    ConstraintSystem &cs, ClosureExpr *closure) {
-  // FIXME: Remove this once the inference is working in general.
-  if (!cs.getASTContext().LangOpts.hasFeature(Feature::FullTypedThrows))
-    return nullptr;
-
-  auto closureType = cs.getClosureType(closure);
-  auto thrownType = closureType->getEffectiveThrownErrorTypeOrNever();
-  auto computedThrownType = cs.inferCaughtErrorType(closure);
-  cs.addConstraint(
-      ConstraintKind::Conversion, computedThrownType, thrownType, 
-      cs.getConstraintLocator(closure,
-                              ConstraintLocator::ClosureThrownError));
-
-  return computedThrownType->getAs<TypeVariableType>();
-}
-
 static void createConjunction(ConstraintSystem &cs, DeclContext *dc,
                               ArrayRef<ElementInfo> elements,
                               ConstraintLocator *locator, bool isIsolated,
@@ -341,6 +324,17 @@ static void createConjunction(ConstraintSystem &cs, DeclContext *dc,
     isIsolated = true;
   }
 
+  if (auto syntacticElement =
+        locator->getLastElementAs<LocatorPathElt::SyntacticElement>()) {
+    // If we are at a do..catch, and we're inferring a throw error type for
+    // the do..catch, we're now in a position to finalize the thrown type.
+    if (auto doCatch = dyn_cast_or_null<DoCatchStmt>(
+                         syntacticElement->getElement().dyn_cast<Stmt *>())) {
+      if (cs.getCaughtErrorType(doCatch)->is<TypeVariableType>())
+        isIsolated = true;
+    }
+  }
+
   if (locator->isForSingleValueStmtConjunction()) {
     auto *SVE = castToExpr<SingleValueStmtExpr>(locator->getAnchor());
     referencedVars.push_back(cs.getType(SVE)->castTo<TypeVariableType>());
@@ -384,19 +378,17 @@ static void createConjunction(ConstraintSystem &cs, DeclContext *dc,
         cs, element, context, elementLoc, isDiscarded));
   }
 
-  for (auto *externalVar : paramCollector.getTypeVars())
-    referencedVars.push_back(externalVar);
-
-#if false
   // If the body of the closure is being used to infer the thrown error type
   // of that closure, introduce a constraint to do so.
   if (locator->directlyAt<ClosureExpr>()) {
     auto *closure = castToExpr<ClosureExpr>(locator->getAnchor());
-    if (auto thrownErrorTypeVar = assignClosureThrownErrorType(cs, closure))
-      referencedVars.push_back(thrownErrorTypeVar);
+    if (auto thrownError = cs.getCaughtErrorType(closure))
+      paramCollector.inferTypeVars(thrownError);
   }
-#endif
-  
+
+  for (auto *externalVar : paramCollector.getTypeVars())
+    referencedVars.push_back(externalVar);
+
   // It's possible that there are no viable elements in the body,
   // because e.g. whole body is an `#if` statement or it only has
   // declarations that are checked during solution application.
@@ -1055,20 +1047,26 @@ class SyntacticElementConstraintGenerator
   void visitDoCatchStmt(DoCatchStmt *doStmt) {
     SmallVector<ElementInfo, 4> elements;
 
+    auto myLocator = locator;
+    if (cs.getASTContext().LangOpts.hasFeature(Feature::FullTypedThrows)) {
+      myLocator = cs.getConstraintLocator(
+          locator, LocatorPathElt::SyntacticElement(doStmt));
+    }
+
     // First, let's record a body of `do` statement. Note we need to add a
     // SyntaticElement locator path element here to avoid treating the inner
     // brace conjunction as being isolated if 'doLoc' is for an isolated
     // conjunction (as is the case with 'do' expressions).
     auto *doBodyLoc = cs.getConstraintLocator(
-        locator, LocatorPathElt::SyntacticElement(doStmt->getBody()));
+        myLocator, LocatorPathElt::SyntacticElement(doStmt->getBody()));
     elements.push_back(makeElement(doStmt->getBody(), doBodyLoc));
 
     // After that has been type-checked, let's switch to
     // individual `catch` statements.
     for (auto *catchStmt : doStmt->getCatches())
       elements.push_back(makeElement(catchStmt, locator));
 
-    createConjunction(elements, locator);
+    createConjunction(elements, myLocator);
   }
 
   void visitCaseStmt(CaseStmt *caseStmt) {
@@ -1156,7 +1154,7 @@ class SyntacticElementConstraintGenerator
       }
 
       if (closure) {
-        assignClosureThrownErrorType(cs, closure);
+        cs.finalizeCaughtErrorType(closure);
       }
 
       return;
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
diff --git a/test/expr/closure/typed_throws_full.swift b/test/expr/closure/typed_throws_full.swift
