Sema: Introduce TypeChecker::availabilityAtLocation().

Implement a query that returns the `AvailabilityContext` for a given
`SourceLoc` and `DeclContext`. Re-implement the existing type checker query
that just returns an `AvailabilityRange` on top of the new query.

Author: tshortli

include/swift/AST/AvailabilityContext.h

@@ -71,6 +71,12 @@ class AvailabilityContext {
   /// Returns true if this context is deprecated on the current platform.
   bool isDeprecated() const;
 
+  /// Constrain with another `AvailabilityContext`.
+  void constrainWithContext(const AvailabilityContext &other, ASTContext &ctx);
+
+  /// Constrain with the availability attributes of `decl`.
+  void constrainWithDecl(const Decl *decl);
+
   /// Constrain the platform availability range with `platformRange`.
   void constrainWithPlatformRange(const AvailabilityRange &platformRange,
                                   ASTContext &ctx);

include/swift/AST/AvailabilityContextStorage.h

@@ -40,10 +40,13 @@ struct AvailabilityContext::PlatformInfo {
   /// platform.
   unsigned IsDeprecated : 1;
 
-  /// Sets `Range` to `other` if `other` is more restrictive. Returns true if
-  /// any property changed as a result of adding this constraint.
+  /// Sets each field to the value of the corresponding field in `other` if the
+  /// other is more restrictive. Returns true if any field changed as a result
+  /// of adding this constraint.
+  bool constrainWith(const PlatformInfo &other);
+
   /// Updates each field to reflect the availability of `decl`, if that
-  /// availability is more restrictive. Return true if any field was updated.
+  /// availability is more restrictive. Returns true if any field was updated.
   bool constrainWith(const Decl *decl);
 
   bool constrainRange(const AvailabilityRange &other) {

lib/AST/AvailabilityContext.cpp

@@ -18,6 +18,18 @@
 
 using namespace swift;
 
+bool AvailabilityContext::PlatformInfo::constrainWith(
+    const PlatformInfo &other) {
+  bool isConstrained = false;
+  isConstrained |= constrainRange(other.Range);
+  if (other.IsUnavailable) {
+    isConstrained |= constrainUnavailability(other.UnavailablePlatform);
+  }
+  isConstrained |= constrainDeprecated(other.IsDeprecated);
+
+  return isConstrained;
+}
+
 bool AvailabilityContext::PlatformInfo::constrainWith(const Decl *decl) {
   bool isConstrained = false;
   auto &ctx = decl->getASTContext();
@@ -129,6 +141,18 @@ bool AvailabilityContext::isDeprecated() const {
   return Info->Platform.IsDeprecated;
 }
 
+void AvailabilityContext::constrainWithContext(const AvailabilityContext &other,
+                                               ASTContext &ctx) {
+  PlatformInfo platformAvailability{Info->Platform};
+  if (platformAvailability.constrainWith(other.Info->Platform)) {
+    Info = Storage::get(platformAvailability, ctx);
+  }
+}
+
+void AvailabilityContext::constrainWithDecl(const Decl *decl) {
+  constrainWithDeclAndPlatformRange(decl, AvailabilityRange::alwaysAvailable());
+}
+
 void AvailabilityContext::constrainWithPlatformRange(
     const AvailabilityRange &platformRange, ASTContext &ctx) {
   PlatformInfo platformAvailability{Info->Platform};

lib/Sema/TypeCheckAvailability.cpp

@@ -1411,9 +1411,9 @@ evaluator::SideEffect ExpandChildTypeRefinementContextsRequest::evaluate(
   return evaluator::SideEffect();
 }
 
-AvailabilityRange TypeChecker::overApproximateAvailabilityAtLocation(
-    SourceLoc loc, const DeclContext *DC,
-    const TypeRefinementContext **MostRefined) {
+AvailabilityContext
+TypeChecker::availabilityAtLocation(SourceLoc loc, const DeclContext *DC,
+                                    const TypeRefinementContext **MostRefined) {
   SourceFile *SF;
   if (loc.isValid())
     SF = DC->getParentModule()->getSourceFileContainingLocation(loc);
@@ -1436,7 +1436,7 @@ AvailabilityRange TypeChecker::overApproximateAvailabilityAtLocation(
   // this will be a real problem.
 
   // We can assume we are running on at least the minimum inlining target.
-  auto OverApproximateContext = AvailabilityRange::forInliningTarget(Context);
+  auto baseAvailability = AvailabilityContext::getDefault(Context);
   auto isInvalidLoc = [SF](SourceLoc loc) {
     return SF ? loc.isInvalid() : true;
   };
@@ -1445,34 +1445,35 @@ AvailabilityRange TypeChecker::overApproximateAvailabilityAtLocation(
     if (!D)
       break;
 
+    baseAvailability.constrainWithDecl(D);
     loc = D->getLoc();
+    DC = D->getDeclContext();
+  }
 
-    std::optional<AvailabilityRange> Info =
-        AvailabilityInference::annotatedAvailableRange(D);
+  if (!SF || loc.isInvalid())
+    return baseAvailability;
 
-    if (Info.has_value()) {
-      OverApproximateContext.constrainWith(Info.value());
-    }
+  TypeRefinementContext *rootTRC = getOrBuildTypeRefinementContext(SF);
+  if (!rootTRC)
+    return baseAvailability;
 
-    DC = D->getDeclContext();
-  }
+  TypeRefinementContext *TRC = rootTRC->findMostRefinedSubContext(loc, Context);
+  if (!TRC)
+    return baseAvailability;
 
-  if (SF && loc.isValid()) {
-    TypeRefinementContext *rootTRC = getOrBuildTypeRefinementContext(SF);
-    if (rootTRC) {
-      TypeRefinementContext *TRC =
-          rootTRC->findMostRefinedSubContext(loc, Context);
-      if (TRC) {
-        OverApproximateContext.constrainWith(
-            TRC->getPlatformAvailabilityRange());
-        if (MostRefined) {
-          *MostRefined = TRC;
-        }
-      }
-    }
+  if (MostRefined) {
+    *MostRefined = TRC;
   }
 
-  return OverApproximateContext;
+  auto availability = TRC->getAvailabilityContext();
+  availability.constrainWithContext(baseAvailability, Context);
+  return availability;
+}
+
+AvailabilityRange TypeChecker::overApproximateAvailabilityAtLocation(
+    SourceLoc loc, const DeclContext *DC,
+    const TypeRefinementContext **MostRefined) {
+  return availabilityAtLocation(loc, DC, MostRefined).getPlatformRange();
 }
 
 bool TypeChecker::isDeclarationUnavailable(

lib/Sema/TypeChecker.h

@@ -998,6 +998,13 @@ bool isAvailabilitySafeForConformance(
     const ValueDecl *witness, const DeclContext *dc,
     AvailabilityRange &requiredAvailability);
 
+/// Returns the most refined `AvailabilityContext` for the given location.
+/// If `MostRefined` is not `nullptr`, it will be set to the most refined TRC
+/// that contains the given location.
+AvailabilityContext
+availabilityAtLocation(SourceLoc loc, const DeclContext *DC,
+                       const TypeRefinementContext **MostRefined = nullptr);
+
 /// Returns an over-approximation of the range of operating system versions
 /// that could the passed-in location could be executing upon for
 /// the target platform. If MostRefined != nullptr, set to the most-refined