Merge pull request #26661 from DougGregor/one-way-constraints-5.1

[5.1] One-way constraints for function builders

Author: DougGregor

include/swift/AST/Expr.h

@@ -5353,6 +5353,33 @@ class KeyPathDotExpr : public Expr {
   }
 };
 
+/// Expression node that effects a "one-way" constraint in
+/// the constraint system, allowing type information to flow from the
+/// subexpression outward but not the other way.
+///
+/// One-way expressions are generally implicit and synthetic, introduced by
+/// the type checker. However, there is a built-in expression of the
+/// form \c Builtin.one_way(x) that forms a one-way constraint coming out
+/// of expression `x` that can be used for testing purposes.
+class OneWayExpr : public Expr {
+  Expr *SubExpr;
+
+public:
+  /// Construct an implicit one-way expression from the given subexpression.
+  OneWayExpr(Expr *subExpr)
+     : Expr(ExprKind::OneWay, /*isImplicit=*/true), SubExpr(subExpr) { }
+
+  SourceLoc getLoc() const { return SubExpr->getLoc(); }
+  SourceRange getSourceRange() const { return SubExpr->getSourceRange(); }
+
+  Expr *getSubExpr() const { return SubExpr; }
+  void setSubExpr(Expr *subExpr) { SubExpr = subExpr; }
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::OneWay;
+  }
+};
+
 inline bool Expr::isInfixOperator() const {
   return isa<BinaryExpr>(this) || isa<IfExpr>(this) ||
          isa<AssignExpr>(this) || isa<ExplicitCastExpr>(this);

include/swift/AST/ExprNodes.def

@@ -190,6 +190,7 @@ EXPR(EditorPlaceholder, Expr)
 EXPR(ObjCSelector, Expr)
 EXPR(KeyPath, Expr)
 UNCHECKED_EXPR(KeyPathDot, Expr)
+UNCHECKED_EXPR(OneWay, Expr)
 EXPR(Tap, Expr)
 LAST_EXPR(Tap)
 

include/swift/Basic/LangOptions.h

@@ -207,6 +207,9 @@ namespace swift {
     /// before termination of the shrink phrase of the constraint solver.
     unsigned SolverShrinkUnsolvedThreshold = 10;
 
+    /// Enable one-way constraints in function builders.
+    bool FunctionBuilderOneWayConstraints = true;
+
     /// Disable the shrink phase of the expression type checker.
     bool SolverDisableShrink = false;
 

include/swift/Basic/Statistics.def

@@ -175,6 +175,11 @@ FRONTEND_STATISTIC(Sema, NumLeafScopes)
 /// This is a measure of complexity of the contraction algorithm.
 FRONTEND_STATISTIC(Sema, NumConstraintsConsideredForEdgeContraction)
 
+/// Number of constraint-solving scopes created in the typechecker, while
+/// solving expression type constraints. A rough proxy for "how much work the
+/// expression typechecker did".
+FRONTEND_STATISTIC(Sema, NumCyclicOneWayComponentsCollapsed)
+
 /// Number of declarations that were deserialized. A rough proxy for the amount
 /// of material loaded from other modules.
 FRONTEND_STATISTIC(Sema, NumDeclsDeserialized)

include/swift/Option/FrontendOptions.td

@@ -394,6 +394,14 @@ def Rmodule_interface_rebuild : Flag<["-"], "Rmodule-interface-rebuild">,
 
 def solver_expression_time_threshold_EQ : Joined<["-"], "solver-expression-time-threshold=">;
 
+def enable_function_builder_one_way_constraints : Flag<["-"],
+  "enable-function-builder-one-way-constraints">,
+  HelpText<"Enable one-way constraints in the function builder transformation">;
+
+def disable_function_builder_one_way_constraints : Flag<["-"],
+  "disable-function-builder-one-way-constraints">,
+  HelpText<"Disable one-way constraints in the function builder transformation">;
+
 def solver_disable_shrink :
   Flag<["-"], "solver-disable-shrink">,
   HelpText<"Disable the shrink phase of expression type checking">;

lib/AST/ASTDumper.cpp

@@ -2745,6 +2745,13 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
 
+  void visitOneWayExpr(OneWayExpr *E) {
+    printCommon(E, "one_way_expr");
+    OS << '\n';
+    printRec(E->getSubExpr());
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
+
   void visitTapExpr(TapExpr *E) {
     printCommon(E, "tap_expr");
     PrintWithColorRAII(OS, DeclColor) << " var=";

lib/AST/ASTWalker.cpp

@@ -1087,6 +1087,18 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
   Expr *visitKeyPathDotExpr(KeyPathDotExpr *E) { return E; }
 
+  Expr *visitOneWayExpr(OneWayExpr *E) {
+    if (auto oldSubExpr = E->getSubExpr()) {
+      if (auto subExpr = doIt(oldSubExpr)) {
+        E->setSubExpr(subExpr);
+      } else {
+        return nullptr;
+      }
+    }
+
+    return E;
+  }
+
   Expr *visitTapExpr(TapExpr *E) {
     if (auto oldSubExpr = E->getSubExpr()) {
       if (auto subExpr = doIt(oldSubExpr)) {

lib/AST/Expr.cpp

@@ -366,6 +366,7 @@ ConcreteDeclRef Expr::getReferencedDecl() const {
   NO_REFERENCE(ObjCSelector);
   NO_REFERENCE(KeyPath);
   NO_REFERENCE(KeyPathDot);
+  PASS_THROUGH_REFERENCE(OneWay, getSubExpr);
   NO_REFERENCE(Tap);
 
 #undef SIMPLE_REFERENCE
@@ -537,6 +538,7 @@ bool Expr::canAppendPostfixExpression(bool appendingPostfixOperator) const {
   case ExprKind::Error:
   case ExprKind::CodeCompletion:
   case ExprKind::LazyInitializer:
+  case ExprKind::OneWay:
     return false;
 
   case ExprKind::NilLiteral:

lib/Frontend/CompilerInvocation.cpp

@@ -419,6 +419,10 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
 
   if (Args.getLastArg(OPT_solver_disable_shrink))
     Opts.SolverDisableShrink = true;
+  Opts.FunctionBuilderOneWayConstraints =
+    Args.hasFlag(OPT_enable_function_builder_one_way_constraints,
+                 OPT_disable_function_builder_one_way_constraints,
+                 /*Default=*/true);
 
   if (const Arg *A = Args.getLastArg(OPT_value_recursion_threshold)) {
     unsigned threshold;

lib/Sema/BuilderTransform.cpp

@@ -50,9 +50,10 @@ class BuilderClosureVisitor
 
 private:
   /// Produce a builder call to the given named function with the given arguments.
-  CallExpr *buildCallIfWanted(SourceLoc loc,
-                              Identifier fnName, ArrayRef<Expr *> args,
-                              ArrayRef<Identifier> argLabels = {}) {
+  Expr *buildCallIfWanted(SourceLoc loc,
+                          Identifier fnName, ArrayRef<Expr *> args,
+                          ArrayRef<Identifier> argLabels,
+                          bool allowOneWay) {
     if (!wantExpr)
       return nullptr;
 
@@ -81,9 +82,17 @@ class BuilderClosureVisitor
         typeExpr, loc, fnName, DeclNameLoc(loc), /*implicit=*/true);
     SourceLoc openLoc = args.empty() ? loc : args.front()->getStartLoc();
     SourceLoc closeLoc = args.empty() ? loc : args.back()->getEndLoc();
-    return CallExpr::create(ctx, memberRef, openLoc, args,
-                            argLabels, argLabelLocs, closeLoc,
-                            /*trailing closure*/ nullptr, /*implicit*/true);
+    Expr *result = CallExpr::create(ctx, memberRef, openLoc, args,
+                                    argLabels, argLabelLocs, closeLoc,
+                                    /*trailing closure*/ nullptr,
+                                    /*implicit*/true);
+
+    if (ctx.LangOpts.FunctionBuilderOneWayConstraints && allowOneWay) {
+      // Form a one-way constraint to prevent backward propagation.
+      result = new (ctx) OneWayExpr(result);
+    }
+
+    return result;
   }
 
   /// Check whether the builder supports the given operation.
@@ -160,12 +169,17 @@ class BuilderClosureVisitor
       }
 
       auto expr = node.get<Expr *>();
+      if (wantExpr && ctx.LangOpts.FunctionBuilderOneWayConstraints)
+        expr = new (ctx) OneWayExpr(expr);
+
       expressions.push_back(expr);
     }
 
     // Call Builder.buildBlock(... args ...)
     return buildCallIfWanted(braceStmt->getStartLoc(),
-                             ctx.Id_buildBlock, expressions);
+                             ctx.Id_buildBlock, expressions,
+                             /*argLabels=*/{ },
+                             /*allowOneWay=*/true);
   }
 
   Expr *visitReturnStmt(ReturnStmt *stmt) {
@@ -190,7 +204,8 @@ class BuilderClosureVisitor
     if (!arg)
       return nullptr;
 
-    return buildCallIfWanted(doStmt->getStartLoc(), ctx.Id_buildDo, arg);
+    return buildCallIfWanted(doStmt->getStartLoc(), ctx.Id_buildDo, arg,
+                             /*argLabels=*/{ }, /*allowOneWay=*/true);
   }
 
   CONTROL_FLOW_STMT(Yield)
@@ -275,7 +290,12 @@ class BuilderClosureVisitor
     // so we just need to call `buildIf` now, since we're at the top level.
     if (isOptional) {
       chainExpr = buildCallIfWanted(ifStmt->getStartLoc(),
-                                    ctx.Id_buildIf, chainExpr);
+                                    ctx.Id_buildIf, chainExpr,
+                                    /*argLabels=*/{ },
+                                    /*allowOneWay=*/true);
+    } else if (ctx.LangOpts.FunctionBuilderOneWayConstraints) {
+      // Form a one-way constraint to prevent backward propagation.
+      chainExpr = new (ctx) OneWayExpr(chainExpr);
     }
 
     return chainExpr;
@@ -395,7 +415,8 @@ class BuilderClosureVisitor
       bool isSecond = (path & 1);
       operand = buildCallIfWanted(operand->getStartLoc(),
                                   ctx.Id_buildEither, operand,
-                                  {isSecond ? ctx.Id_second : ctx.Id_first});
+                                  {isSecond ? ctx.Id_second : ctx.Id_first},
+                                  /*allowOneWay=*/false);
     }
 
     // Inject into Optional if required.  We'll be adding the call to