[Sema] Fix optional chaining behavior with postfix operators

Postfix operators can further be chained within an optional binding
chain, so we need to make sure they're handled in
`getMemberChainSubExpr`. Unresolved member chains still don't allow
them, so we need to add a new `kind` parameter to differentiate the
behavior here.

rdar://147826988

Author: hamishknight

lib/Sema/PreCheckTarget.cpp

@@ -317,54 +317,78 @@ static bool findNonMembers(ArrayRef<LookupResultEntry> lookupResults,
   return AllDeclRefs;
 }
 
+namespace {
+enum class MemberChainKind {
+  OptionalBind,     // A 'x?.y' optional binding chain
+  UnresolvedMember, // A '.foo.bar' chain
+};
+} // end anonymous namespace
+
 /// Find the next element in a chain of members. If this expression is (or
 /// could be) the base of such a chain, this will return \c nullptr.
-static Expr *getMemberChainSubExpr(Expr *expr) {
+static Expr *getMemberChainSubExpr(Expr *expr, MemberChainKind kind) {
   assert(expr && "getMemberChainSubExpr called with null expr!");
-  if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
+  if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr))
     return UDE->getBase();
-  } else if (auto *CE = dyn_cast<CallExpr>(expr)) {
+  if (auto *CE = dyn_cast<CallExpr>(expr))
     return CE->getFn();
-  } else if (auto *BOE = dyn_cast<BindOptionalExpr>(expr)) {
+  if (auto *BOE = dyn_cast<BindOptionalExpr>(expr))
     return BOE->getSubExpr();
-  } else if (auto *FVE = dyn_cast<ForceValueExpr>(expr)) {
+  if (auto *FVE = dyn_cast<ForceValueExpr>(expr))
     return FVE->getSubExpr();
-  } else if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
+  if (auto *SE = dyn_cast<SubscriptExpr>(expr))
     return SE->getBase();
-  } else if (auto *DSE = dyn_cast<DotSelfExpr>(expr)) {
+  if (auto *DSE = dyn_cast<DotSelfExpr>(expr))
     return DSE->getSubExpr();
-  } else if (auto *USE = dyn_cast<UnresolvedSpecializeExpr>(expr)) {
+  if (auto *USE = dyn_cast<UnresolvedSpecializeExpr>(expr))
     return USE->getSubExpr();
-  } else if (auto *CCE = dyn_cast<CodeCompletionExpr>(expr)) {
+  if (auto *CCE = dyn_cast<CodeCompletionExpr>(expr))
     return CCE->getBase();
-  } else {
-    return nullptr;
+
+  if (kind == MemberChainKind::OptionalBind) {
+    // We allow postfix operators to be part of the optional member chain, e.g:
+    //
+    //   for?.bar++
+    //   x.y?^.foo()
+    //
+    // Note this behavior is specific to optional chains, we treat e.g
+    // `.foo^` as `(.foo)^`.
+    if (auto *PO = dyn_cast<PostfixUnaryExpr>(expr))
+      return PO->getOperand();
+
+    // Unresolved member chains can themselves be nested in optional chains
+    // since optional chains can include postfix operators.
+    if (auto *UME = dyn_cast<UnresolvedMemberChainResultExpr>(expr))
+      return UME->getSubExpr();
   }
+
+  return nullptr;
 }
 
 UnresolvedMemberExpr *TypeChecker::getUnresolvedMemberChainBase(Expr *expr) {
-  if (auto *subExpr = getMemberChainSubExpr(expr))
+  if (auto *subExpr =
+          getMemberChainSubExpr(expr, MemberChainKind::UnresolvedMember)) {
     return getUnresolvedMemberChainBase(subExpr);
-  else
-    return dyn_cast<UnresolvedMemberExpr>(expr);
+  }
+  return dyn_cast<UnresolvedMemberExpr>(expr);
 }
 
 static bool isBindOptionalMemberChain(Expr *expr) {
-  if (isa<BindOptionalExpr>(expr)) {
+  if (isa<BindOptionalExpr>(expr))
     return true;
-  } else if (auto *base = getMemberChainSubExpr(expr)) {
+
+  if (auto *base = getMemberChainSubExpr(expr, MemberChainKind::OptionalBind))
     return isBindOptionalMemberChain(base);
-  } else {
-    return false;
-  }
+
+  return false;
 }
 
 /// Whether this expression sits at the end of a chain of member accesses.
-static bool isMemberChainTail(Expr *expr, Expr *parent) {
+static bool isMemberChainTail(Expr *expr, Expr *parent, MemberChainKind kind) {
   assert(expr && "isMemberChainTail called with null expr!");
   // If this expression's parent is not itself part of a chain (or, this expr
   // has no parent expr), this must be the tail of the chain.
-  return !parent || getMemberChainSubExpr(parent) != expr;
+  return !parent || getMemberChainSubExpr(parent, kind) != expr;
 }
 
 static bool isValidForwardReference(ValueDecl *D, DeclContext *DC,
@@ -2641,7 +2665,6 @@ Expr *PreCheckTarget::simplifyTypeConstructionWithLiteralArg(Expr *E) {
 ///
 ///   foo? = newFoo // LHS of the assignment operator
 ///   foo?.bar += value // LHS of 'assignment: true' precedence group operators.
-///   for?.bar++ // Postfix operator.
 ///
 /// In such cases, the operand is constructed to be an 'OperatorEvaluationExpr'
 /// wrapping the actual operand. This function hoist it and wraps the entire
@@ -2666,12 +2689,6 @@ PreCheckTarget::hoistOptionalEvaluationExprIfNeeded(Expr *expr) {
         }
       }
     }
-  } else if (auto *postfixE = dyn_cast<PostfixUnaryExpr>(expr)) {
-    if (auto *OEE = dyn_cast<OptionalEvaluationExpr>(postfixE->getOperand())) {
-      postfixE->setOperand(OEE->getSubExpr());
-      OEE->setSubExpr(postfixE);
-      return OEE;
-    }
   }
   return nullptr;
 }
@@ -2680,18 +2697,21 @@ Expr *PreCheckTarget::wrapMemberChainIfNeeded(Expr *E) {
   auto *parent = Parent.getAsExpr();
   Expr *wrapped = E;
 
-  if (!isMemberChainTail(E, parent))
+  // If the parent is already wrapped, we've already formed the member chain.
+  if (parent && (isa<OptionalEvaluationExpr>(parent) ||
+                 isa<UnresolvedMemberChainResultExpr>(parent))) {
     return E;
+  }
 
   // If we find an unresolved member chain, wrap it in an
-  // UnresolvedMemberChainResultExpr (unless this has already been done).
-  if (auto *UME = TypeChecker::getUnresolvedMemberChainBase(E)) {
-    if (!parent || !isa<UnresolvedMemberChainResultExpr>(parent))
+  // UnresolvedMemberChainResultExpr.
+  if (isMemberChainTail(E, parent, MemberChainKind::UnresolvedMember)) {
+    if (auto *UME = TypeChecker::getUnresolvedMemberChainBase(E))
       wrapped = new (Ctx) UnresolvedMemberChainResultExpr(E, UME);
   }
   // Wrap optional chain in an OptionalEvaluationExpr.
-  if (isBindOptionalMemberChain(E)) {
-    if (!parent || !isa<OptionalEvaluationExpr>(parent))
+  if (isMemberChainTail(E, parent, MemberChainKind::OptionalBind)) {
+    if (isBindOptionalMemberChain(E))
       wrapped = new (Ctx) OptionalEvaluationExpr(wrapped);
   }
   return wrapped;

test/SILGen/member_chains.swift

@@ -0,0 +1,36 @@
+// RUN: %target-swift-emit-silgen -verify %s
+
+struct ImplicitMembers: Equatable {
+  static var implicit = ImplicitMembers()
+
+  static var optional: ImplicitMembers? = ImplicitMembers()
+  static func createOptional() -> ImplicitMembers? {
+    ImplicitMembers()
+  }
+
+  var another: ImplicitMembers { ImplicitMembers() }
+  var anotherOptional: ImplicitMembers? { ImplicitMembers() }
+}
+
+// Make sure we can SILGen these without issue.
+
+postfix operator ^
+postfix func ^ (_ lhs: ImplicitMembers) -> Int { 0 }
+
+extension Int {
+  func foo() {}
+  var optionalMember: Int? { 0 }
+}
+
+// https://github.com/swiftlang/swift/issues/80265
+// Make sure optional chaining looks through postfix operators.
+var x: ImplicitMembers?
+let _ = x?^.foo()
+let _ = x?^.optionalMember?.foo()
+let _ = x?.another^.optionalMember?.foo()
+
+// Make sure the unresolved member chain extends up to the postfix operator,
+// but the optional chain covers the entire expr.
+let _ = .optional?^.foo()
+let _ = .createOptional()?^.foo()
+let _ = .implicit.anotherOptional?^.foo()