From 48eb4b771472e122f776a982d24232be3e867c29 Mon Sep 17 00:00:00 2001
From: Holly Borla <hborla@apple.com>
Date: Fri, 2 Feb 2024 17:26:46 -0800
Subject: [PATCH] [Availability] Lazily expand type refinement contexts for
 extensions.

---
 lib/Sema/TypeCheckAvailability.cpp            |  4 ++++
 test/Macros/macro_expand_extensions.swift     | 10 ++++++++
 .../availability_refinement_contexts.swift    | 24 ++++++++++---------
 test/stdlib/Observation/Observable.swift      |  6 +++++
 4 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/lib/Sema/TypeCheckAvailability.cpp b/lib/Sema/TypeCheckAvailability.cpp
index 4815090a814c5..c4b191d02e989 100644
--- a/lib/Sema/TypeCheckAvailability.cpp
+++ b/lib/Sema/TypeCheckAvailability.cpp
@@ -551,6 +551,10 @@ class TypeRefinementContextBuilder : private ASTWalker {
       return true;
     }
 
+    if (isa<ExtensionDecl>(D)) {
+      return true;
+    }
+
     return false;
   }
 
diff --git a/test/Macros/macro_expand_extensions.swift b/test/Macros/macro_expand_extensions.swift
index e88628cdee978..edbe672821c44 100644
--- a/test/Macros/macro_expand_extensions.swift
+++ b/test/Macros/macro_expand_extensions.swift
@@ -256,3 +256,13 @@ final class SendableClass {
 @AddSendable
 class InvalidSendableClass {
 }
+
+@AddSendable
+struct HasNestedType {
+  struct Inner {}
+}
+
+// Make sure no circularity error is produced when resolving
+// extensions of nested types when the outer type has an
+// attached macro that can add other nested types.
+extension HasNestedType.Inner {}
diff --git a/test/Sema/availability_refinement_contexts.swift b/test/Sema/availability_refinement_contexts.swift
index 837ae7847b3c5..448feff32dd01 100644
--- a/test/Sema/availability_refinement_contexts.swift
+++ b/test/Sema/availability_refinement_contexts.swift
@@ -73,8 +73,9 @@ protocol SomeProtocol {
   var protoProperty: Int { get }
 }
 
-// CHECK-NEXT: {{^}}  (decl versions=[10.51,+Inf) decl=extension.SomeClass
-// CHECK-NEXT: {{^}}    (decl versions=[10.52,+Inf) decl=someExtensionFunction()
+// CHECK-NEXT: {{^}}  (decl_implicit versions=[10.13,+Inf) decl=extension.SomeClass
+// CHECK-NEXT: {{^}}    (decl versions=[10.51,+Inf) decl=extension.SomeClass
+// CHECK-NEXT: {{^}}      (decl versions=[10.52,+Inf) decl=someExtensionFunction()
 @available(OSX 10.51, *)
 extension SomeClass {
   @available(OSX 10.52, *)
@@ -207,15 +208,16 @@ func functionWithWhile() {
   }
 }
 
-// CHECK-NEXT: {{^}}  (decl versions=[10.51,+Inf) decl=extension.SomeClass
-// CHECK-NEXT: {{^}}    (decl_implicit versions=[10.51,+Inf) decl=someStaticPropertyWithClosureInit
-// CHECK-NEXT: {{^}}      (decl versions=[10.52,+Inf) decl=someStaticPropertyWithClosureInit
-// CHECK-NEXT: {{^}}        (condition_following_availability versions=[10.54,+Inf)
-// CHECK-NEXT: {{^}}        (if_then versions=[10.54,+Inf)
-// CHECK-NEXT: {{^}}    (decl_implicit versions=[10.51,+Inf) decl=someStaticPropertyWithClosureInitInferred
-// CHECK-NEXT: {{^}}      (decl versions=[10.52,+Inf) decl=someStaticPropertyWithClosureInitInferred
-// CHECK-NEXT: {{^}}        (condition_following_availability versions=[10.54,+Inf)
-// CHECK-NEXT: {{^}}        (if_then versions=[10.54,+Inf)
+// CHECK-NEXT: {{^}}  (decl_implicit versions=[10.13,+Inf) decl=extension.SomeClass
+// CHECK-NEXT: {{^}}    (decl versions=[10.51,+Inf) decl=extension.SomeClass
+// CHECK-NEXT: {{^}}      (decl_implicit versions=[10.51,+Inf) decl=someStaticPropertyWithClosureInit
+// CHECK-NEXT: {{^}}        (decl versions=[10.52,+Inf) decl=someStaticPropertyWithClosureInit
+// CHECK-NEXT: {{^}}          (condition_following_availability versions=[10.54,+Inf)
+// CHECK-NEXT: {{^}}          (if_then versions=[10.54,+Inf)
+// CHECK-NEXT: {{^}}      (decl_implicit versions=[10.51,+Inf) decl=someStaticPropertyWithClosureInitInferred
+// CHECK-NEXT: {{^}}        (decl versions=[10.52,+Inf) decl=someStaticPropertyWithClosureInitInferred
+// CHECK-NEXT: {{^}}          (condition_following_availability versions=[10.54,+Inf)
+// CHECK-NEXT: {{^}}          (if_then versions=[10.54,+Inf)
 @available(OSX 10.51, *)
 extension SomeClass {
   @available(OSX 10.52, *)
diff --git a/test/stdlib/Observation/Observable.swift b/test/stdlib/Observation/Observable.swift
index f4280e33036a3..06dbd24137eaf 100644
--- a/test/stdlib/Observation/Observable.swift
+++ b/test/stdlib/Observation/Observable.swift
@@ -225,6 +225,12 @@ class GuardedAvailability {
   }()
 }
 
+@Observable class Parent {
+  class Nested {}
+}
+
+extension Parent.Nested {}
+
 struct CowContainer {
   final class Contents { }