[Sema] Correctly re-contextualize if/switch exprs in lazy vars

include/swift/AST/Pattern.h

@@ -566,6 +566,7 @@ class EnumElementPattern : public Pattern {
   void setSubPattern(Pattern *p) { SubPattern = p; }
 
   DeclContext *getDeclContext() const { return DC; }
+  void setDeclContext(DeclContext *newDC) { DC = newDC; }
 
   DeclNameRef getName() const { return Name; }
 
@@ -706,6 +707,12 @@ class ExprPattern : public Pattern {
 
   DeclContext *getDeclContext() const { return DC; }
 
+  void setDeclContext(DeclContext *newDC) {
+    DC = newDC;
+    if (MatchVar)
+      MatchVar->setDeclContext(newDC);
+  }
+
   /// The match expression if it has been computed, \c nullptr otherwise.
   /// Should only be used by the ASTDumper and ASTWalker.
   Expr *getCachedMatchExpr() const { return MatchExpr; }

lib/Sema/MiscDiagnostics.cpp

@@ -3912,7 +3912,8 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
       // Diagnose a SingleValueStmtExpr in a context that we do not currently
-      // support.
+      // support. If we start allowing these in arbitrary places, we'll need
+      // to ensure that autoclosures correctly contextualize them.
       if (!ValidSingleValueStmtExprs.contains(SVE)) {
         Diags.diagnose(SVE->getLoc(), diag::single_value_stmt_out_of_place,
                        SVE->getStmt()->getKind());

lib/Sema/TypeCheckStorage.cpp

@@ -1552,6 +1552,9 @@ synthesizeReadCoroutineGetterBody(AccessorDecl *getter, ASTContext &ctx) {
 namespace {
   /// This ASTWalker explores an expression tree looking for expressions (which
   /// are DeclContext's) and changes their parent DeclContext to NewDC.
+  /// TODO: We ought to consider merging this with
+  /// ContextualizeClosuresAndMacros, or better yet removing it in favor of
+  /// avoiding the recontextualization for lazy vars.
   class RecontextualizeClosures : public ASTWalker {
     DeclContext *NewDC;
   public:
@@ -1567,34 +1570,42 @@ namespace {
         CE->setParent(NewDC);
         return Action::SkipChildren(E);
       }
-
-      if (auto CLE = dyn_cast<CaptureListExpr>(E)) {
-        // Make sure to recontextualize any decls in the capture list as well.
-        for (auto &CLE : CLE->getCaptureList()) {
-          CLE.getVar()->setDeclContext(NewDC);
-          CLE.PBD->setDeclContext(NewDC);
-        }
-      }
-
-      // Unlike a closure, a TapExpr is not a DeclContext, so we need to
-      // recontextualize its variable and then anything else in its body.
-      // FIXME: Might be better to change walkToDeclPre() and walkToStmtPre()
-      // below, but I don't know what other effects that might have.
-      if (auto TE = dyn_cast<TapExpr>(E)) {
-        TE->getVar()->setDeclContext(NewDC);
-        for (auto node : TE->getBody()->getElements())
-          node.walk(RecontextualizeClosures(NewDC));
-      }
 
       return Action::Continue(E);
     }
 
-    /// We don't want to recurse into declarations or statements.
-    PreWalkAction walkToDeclPre(Decl *) override {
-      return Action::SkipChildren();
+    PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
+      if (auto *EP = dyn_cast<ExprPattern>(P))
+        EP->setDeclContext(NewDC);
+      if (auto *EP = dyn_cast<EnumElementPattern>(P))
+        EP->setDeclContext(NewDC);
+
+      return Action::Continue(P);
     }
+
     PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-      return Action::SkipChildren(S);
+      // The ASTWalker doesn't walk the case body variables, contextualize them
+      // ourselves.
+      if (auto *CS = dyn_cast<CaseStmt>(S)) {
+        for (auto *CaseVar : CS->getCaseBodyVariablesOrEmptyArray())
+          CaseVar->setDeclContext(NewDC);
+      }
+      return Action::Continue(S);
+    }
+
+    PreWalkAction walkToDeclPre(Decl *D) override {
+      D->setDeclContext(NewDC);
+
+      // Auxiliary decls need to have their contexts adjusted too.
+      if (auto *VD = dyn_cast<VarDecl>(D)) {
+        VD->visitAuxiliaryDecls([&](VarDecl *D) {
+          D->setDeclContext(NewDC);
+        });
+      }
+
+      // Skip walking the children of any Decls that are also DeclContexts,
+      // they will already have the right parent.
+      return Action::SkipChildrenIf(isa<DeclContext>(D));
     }
   };
 } // end anonymous namespace

test/SILGen/if_expr.swift

@@ -1,5 +1,8 @@
-// RUN: %target-swift-emit-silgen %s | %FileCheck %s
-// RUN: %target-swift-emit-ir %s
+// RUN: %target-swift-emit-silgen -enable-experimental-feature ThenStatements %s | %FileCheck %s
+// RUN: %target-swift-emit-ir -enable-experimental-feature ThenStatements %s
+
+// Needed for experimental features
+// REQUIRES: asserts
 
 func foo() -> Int {
   if .random() { 1 } else { 2 }
@@ -580,3 +583,79 @@ func testConditionalCast<T>(_ x: Any) -> T? {
     nil
   }
 }
+
+@propertyWrapper
+struct Wrapper<T> {
+  var wrappedValue: T
+}
+
+// rdar://119158202 - Make sure we correctly contextualize local bindings.
+func testLazyLocal(_ x: Int?) {
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1aL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+  lazy var a = if let x { x } else { 0 }
+  _ = a
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1bL_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+  lazy var b = if .random() {
+    let x = ""
+    then x
+  } else {
+    ""
+  }
+  _ = b
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1cL_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+  lazy var c = if .random() {
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1cL_SSvg1xL2_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+    lazy var x = ""
+    then x
+  } else {
+    ""
+  }
+  _ = c
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1dL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+  lazy var d = if .random() {
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1dL_Sivg1yL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+    lazy var y = if let x { x } else { 0 }
+    then y
+  } else {
+    0
+  }
+  _ = d
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr13testLazyLocalyySiSgF1eL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }) -> Int
+  lazy var e = if .random() {
+    @Wrapper
+    var x = 0
+    then x
+  } else {
+    0
+  }
+  _ = e
+}
+
+struct LazyProp {
+  var a: Int?
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s7if_expr8LazyPropV1bSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var b = if let a { a } else { 0 }
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s7if_expr8LazyPropV1cSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var c = if .random() {
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr8LazyPropV1cSivg1xL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }) -> Int
+    lazy var x = 0
+    then x
+  } else {
+    0
+  }
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s7if_expr8LazyPropV1dSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var d = if .random() {
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr8LazyPropV1dSivg1xL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, @inout_aliasable LazyProp) -> Int
+    lazy var x = if case let a? = a { a } else { 0 }
+    then x
+  } else {
+    0
+  }
+}

test/SILGen/switch_expr.swift

@@ -1,5 +1,8 @@
-// RUN: %target-swift-emit-silgen %s | %FileCheck %s
-// RUN: %target-swift-emit-ir %s
+// RUN: %target-swift-emit-silgen -enable-experimental-feature ThenStatements %s | %FileCheck %s
+// RUN: %target-swift-emit-ir -enable-experimental-feature ThenStatements %s
+
+// Needed for experimental features
+// REQUIRES: asserts
 
 func foo() -> Int {
   switch Bool.random() {
@@ -775,3 +778,94 @@ func testConditionalCast<T>(_ x: Any, _ y: Int) -> T? {
     x as? T
   }
 }
+
+@propertyWrapper
+struct Wrapper<T> {
+  var wrappedValue: T
+}
+
+// rdar://119158202 - Make sure we correctly contextualize local bindings.
+func testLazyLocal(_ x: Int?) {
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1aL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+  lazy var a = switch x { case let x?: x case nil: 0 }
+  _ = a
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1bL_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+  lazy var b = switch Bool.random() {
+  case true:
+    let x = ""
+    then x
+  case false:
+    ""
+  }
+  _ = b
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1cL_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+  lazy var c = switch Bool.random() {
+  case true:
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1cL_SSvg1xL3_SSvg : $@convention(thin) (@guaranteed { var Optional<String> }) -> @owned String
+    lazy var x = ""
+    then x
+  case false:
+    ""
+  }
+  _ = c
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1dL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+  lazy var d = switch Bool.random() {
+  case true:
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1dL_Sivg1yL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Optional<Int>) -> Int
+    lazy var y = switch x { case let x?: x case nil: 0 }
+    then y
+  case false:
+    0
+  }
+  _ = d
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1eL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }) -> Int
+  lazy var e = switch Bool.random() {
+  case true:
+    @Wrapper
+    var x = 0
+    then x
+  case false:
+    0
+  }
+  _ = e
+
+  // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr13testLazyLocalyySiSgF1fL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }) -> Int
+  lazy var f = switch 0 {
+  case 1:
+    1
+  default:
+    0
+  }
+  _ = f
+}
+
+struct LazyProp {
+  var a: Int?
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s11switch_expr8LazyPropV1bSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var b = switch a { case let a?: a case nil: 0 }
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s11switch_expr8LazyPropV1cSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var c = switch Bool.random() {
+  case true:
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr8LazyPropV1cSivg1xL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }) -> Int
+    lazy var x = 0
+    then x
+  case false:
+    0
+  }
+
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] [ossa] @$s11switch_expr8LazyPropV1dSivg : $@convention(method) (@inout LazyProp) -> Int
+  lazy var d = switch Bool.random() {
+  case true:
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s11switch_expr8LazyPropV1dSivg1xL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, @inout_aliasable LazyProp) -> Int
+    lazy var x = switch a { case let a?: a case nil: 0 }
+    then x
+  case false:
+    0
+  }
+}