lib/Sema/BuilderTransform.cpp - third_party/swift - Git at Google

 //===--- BuilderTransform.cpp - Function-builder transformation -----------===//
 //
 // This source file is part of the Swift.org open source project
 //
 // Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements routines associated with the function-builder
 // transformation.
 //
 //===----------------------------------------------------------------------===//

 #include "ConstraintSystem.h"
 #include "TypeChecker.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include <iterator>
 #include <map>
 #include <memory>
 #include <utility>
 #include <tuple>

 using namespace swift;
 using namespace constraints;

 namespace {

 /// Visitor to classify the contents of the given closure.
 class BuilderClosureVisitor
     : public StmtVisitor<BuilderClosureVisitor, Expr *> {
   ConstraintSystem *cs;
   ASTContext &ctx;
   bool wantExpr;
   Type builderType;
   NominalTypeDecl *builder = nullptr;
   llvm::SmallDenseMap<Identifier, bool> supportedOps;

 public:
   SkipUnhandledConstructInFunctionBuilder::UnhandledNode unhandledNode;

 private:
   /// Produce a builder call to the given named function with the given arguments.
   Expr *buildCallIfWanted(SourceLoc loc,
                           Identifier fnName, ArrayRef<Expr *> args,
                           ArrayRef<Identifier> argLabels,
                           bool allowOneWay) {
     if (!wantExpr)
       return nullptr;

     // FIXME: Setting a TypeLoc on this expression is necessary in order
     // to get diagnostics if something about this builder call fails,
     // e.g. if there isn't a matching overload for `buildBlock`.
     // But we can only do this if there isn't a type variable in the type.
     TypeLoc typeLoc;
     if (!builderType->hasTypeVariable()) {
       typeLoc = TypeLoc(new (ctx) FixedTypeRepr(builderType, loc), builderType);
     }

     auto typeExpr = new (ctx) TypeExpr(typeLoc);
     if (cs) {
       cs->setType(typeExpr, MetatypeType::get(builderType));
       cs->setType(&typeExpr->getTypeLoc(), builderType);
     }

     SmallVector<SourceLoc, 4> argLabelLocs;
     for (auto i : indices(argLabels)) {
       argLabelLocs.push_back(args[i]->getStartLoc());
     }

     typeExpr->setImplicit();
     auto memberRef = new (ctx) UnresolvedDotExpr(
         typeExpr, loc, fnName, DeclNameLoc(loc), /*implicit=*/true);
     SourceLoc openLoc = args.empty() ? loc : args.front()->getStartLoc();
     SourceLoc closeLoc = args.empty() ? loc : args.back()->getEndLoc();
     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.
   bool builderSupports(Identifier fnName,
                        ArrayRef<Identifier> argLabels = {}) {
     auto known = supportedOps.find(fnName);
     if (known != supportedOps.end()) {
       return known->second;
     }

     bool found = false;
     for (auto decl : builder->lookupDirect(fnName)) {
       if (auto func = dyn_cast<FuncDecl>(decl)) {
         // Function must be static.
         if (!func->isStatic())
           continue;

         // Function must have the right argument labels, if provided.
         if (!argLabels.empty()) {
           auto funcLabels = func->getFullName().getArgumentNames();
           if (argLabels.size() > funcLabels.size() ||
               funcLabels.slice(0, argLabels.size()) != argLabels)
             continue;
         }

         // Okay, it's a good-enough match.
         found = true;
         break;
       }
     }

     return supportedOps[fnName] = found;
   }

 public:
   BuilderClosureVisitor(ASTContext &ctx, ConstraintSystem *cs,
                         bool wantExpr, Type builderType)
       : cs(cs), ctx(ctx), wantExpr(wantExpr), builderType(builderType) {
     assert((cs || !builderType->hasTypeVariable()) &&
            "cannot handle builder type with type variables without "
            "constraint system");
     builder = builderType->getAnyNominal();
   }

 #define CONTROL_FLOW_STMT(StmtClass)                      \
   Expr *visit##StmtClass##Stmt(StmtClass##Stmt *stmt) { \
     if (!unhandledNode)                                 \
       unhandledNode = stmt;                             \
                                                         \
     return nullptr;                                     \
   }

   Expr *visitBraceStmt(BraceStmt *braceStmt) {
     SmallVector<Expr *, 4> expressions;
     for (const auto &node : braceStmt->getElements()) {
       if (auto stmt = node.dyn_cast<Stmt *>()) {
         auto expr = visit(stmt);
         if (expr)
           expressions.push_back(expr);
         continue;
       }

       if (auto decl = node.dyn_cast<Decl *>()) {
         // Just ignore #if; the chosen children should appear in the
         // surrounding context.  This isn't good for source tools but it
         // at least works.
         if (isa<IfConfigDecl>(decl))
           continue;

         if (!unhandledNode)
           unhandledNode = decl;

         continue;
       }

       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,
                              /*argLabels=*/{ },
                              /*allowOneWay=*/true);
   }

   Expr *visitReturnStmt(ReturnStmt *stmt) {
     // Allow implicit returns due to 'return' elision.
     if (!stmt->isImplicit() || !stmt->hasResult()) {
       if (!unhandledNode)
         unhandledNode = stmt;
       return nullptr;
     }

     return stmt->getResult();
   }

   Expr *visitDoStmt(DoStmt *doStmt) {
     if (!builderSupports(ctx.Id_buildDo)) {
       if (!unhandledNode)
         unhandledNode = doStmt;
       return nullptr;
     }

     auto arg = visit(doStmt->getBody());
     if (!arg)
       return nullptr;

     return buildCallIfWanted(doStmt->getStartLoc(), ctx.Id_buildDo, arg,
                              /*argLabels=*/{ }, /*allowOneWay=*/true);
   }

   CONTROL_FLOW_STMT(Yield)
   CONTROL_FLOW_STMT(Defer)

   static Expr *getTrivialBooleanCondition(StmtCondition condition) {
     if (condition.size() != 1)
       return nullptr;

     return condition.front().getBooleanOrNull();
   }

   static bool isBuildableIfChainRecursive(IfStmt *ifStmt,
                                           unsigned &numPayloads,
                                           bool &isOptional) {
     // The conditional must be trivial.
     if (!getTrivialBooleanCondition(ifStmt->getCond()))
       return false;

     // The 'then' clause contributes a payload.
     numPayloads++;

     // If there's an 'else' clause, it contributes payloads:
     if (auto elseStmt = ifStmt->getElseStmt()) {
       // If it's 'else if', it contributes payloads recursively.
       if (auto elseIfStmt = dyn_cast<IfStmt>(elseStmt)) {
         return isBuildableIfChainRecursive(elseIfStmt, numPayloads,
                                            isOptional);
       // Otherwise it's just the one.
       } else {
         numPayloads++;
       }

     // If not, the chain result is at least optional.
     } else {
       isOptional = true;
     }

     return true;
   }

   bool isBuildableIfChain(IfStmt *ifStmt, unsigned &numPayloads,
                           bool &isOptional) {
     if (!isBuildableIfChainRecursive(ifStmt, numPayloads, isOptional))
       return false;

     // If there's a missing 'else', we need 'buildIf' to exist.
     if (isOptional && !builderSupports(ctx.Id_buildIf))
       return false;

     // If there are multiple clauses, we need 'buildEither(first:)' and
     // 'buildEither(second:)' to both exist.
     if (numPayloads > 1) {
       if (!builderSupports(ctx.Id_buildEither, {ctx.Id_first}) ||
           !builderSupports(ctx.Id_buildEither, {ctx.Id_second}))
         return false;
     }

     return true;
   }

   Expr *visitIfStmt(IfStmt *ifStmt) {
     // Check whether the chain is buildable and whether it terminates
     // without an `else`.
     bool isOptional = false;
     unsigned numPayloads = 0;
     if (!isBuildableIfChain(ifStmt, numPayloads, isOptional)) {
       if (!unhandledNode)
         unhandledNode = ifStmt;
       return nullptr;
     }

     // Attempt to build the chain, propagating short-circuits, which
     // might arise either do to error or not wanting an expression.
     auto chainExpr =
       buildIfChainRecursive(ifStmt, 0, numPayloads, isOptional);
     if (!chainExpr)
       return nullptr;
     assert(wantExpr);

     // The operand should have optional type if we had optional results,
     // 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,
                                     /*argLabels=*/{ },
                                     /*allowOneWay=*/true);
     } else if (ctx.LangOpts.FunctionBuilderOneWayConstraints) {
       // Form a one-way constraint to prevent backward propagation.
       chainExpr = new (ctx) OneWayExpr(chainExpr);
     }

     return chainExpr;
   }

   /// Recursively build an if-chain: build an expression which will have
   /// a value of the chain result type before any call to `buildIf`.
   /// The expression will perform any necessary calls to `buildEither`,
   /// and the result will have optional type if `isOptional` is true.
   Expr *buildIfChainRecursive(IfStmt *ifStmt, unsigned payloadIndex,
                               unsigned numPayloads, bool isOptional) {
     assert(payloadIndex < numPayloads);
     // Make sure we recursively visit both sides even if we're not
     // building expressions.

     // Build the then clause.  This will have the corresponding payload
     // type (i.e. not wrapped in any way).
     Expr *thenArg = visit(ifStmt->getThenStmt());

     // Build the else clause, if present.  If this is from an else-if,
     // this will be fully wrapped; otherwise it will have the corresponding
     // payload type (at index `payloadIndex + 1`).
     assert(ifStmt->getElseStmt() || isOptional);
     bool isElseIf = false;
     Optional<Expr *> elseChain;
     if (auto elseStmt = ifStmt->getElseStmt()) {
       if (auto elseIfStmt = dyn_cast<IfStmt>(elseStmt)) {
         isElseIf = true;
         elseChain = buildIfChainRecursive(elseIfStmt, payloadIndex + 1,
                                           numPayloads, isOptional);
       } else {
         elseChain = visit(elseStmt);
       }
     }

     // Short-circuit if appropriate.
     if (!wantExpr || !thenArg || (elseChain && !*elseChain))
       return nullptr;

     // Okay, build the conditional expression.

     // Prepare the `then` operand by wrapping it to produce a chain result.
     SourceLoc thenLoc = ifStmt->getThenStmt()->getStartLoc();
     Expr *thenExpr = buildWrappedChainPayload(thenArg, payloadIndex,
                                               numPayloads, isOptional);

     // Prepare the `else operand:
     Expr *elseExpr;
     SourceLoc elseLoc;

     // - If there's no `else` clause, use `Optional.none`.
     if (!elseChain) {
       assert(isOptional);
       elseLoc = ifStmt->getEndLoc();
       elseExpr = buildNoneExpr(elseLoc);

     // - If there's an `else if`, the chain expression from that
     //   should already be producing a chain result.
     } else if (isElseIf) {
       elseExpr = *elseChain;
       elseLoc = ifStmt->getElseLoc();

     // - Otherwise, wrap it to produce a chain result.
     } else {
       elseLoc = ifStmt->getElseLoc();
       elseExpr = buildWrappedChainPayload(*elseChain,
                                           payloadIndex + 1, numPayloads,
                                           isOptional);
     }

     Expr *condition = getTrivialBooleanCondition(ifStmt->getCond());
     assert(condition && "checked by isBuildableIfChain");

     auto ifExpr = new (ctx) IfExpr(condition, thenLoc, thenExpr,
                                    elseLoc, elseExpr);
     ifExpr->setImplicit();
     return ifExpr;
   }

   /// Wrap a payload value in an expression which will produce a chain
   /// result (without `buildIf`).
   Expr *buildWrappedChainPayload(Expr *operand, unsigned payloadIndex,
                                  unsigned numPayloads, bool isOptional) {
     assert(payloadIndex < numPayloads);

     // Inject into the appropriate chain position.
     //
     // We produce a (left-biased) balanced binary tree of Eithers in order
     // to prevent requiring a linear number of injections in the worst case.
     // That is, if we have 13 clauses, we want to produce:
     //
     //                      /------------------Either------------\
     //           /-------Either-------\                     /--Either--\
     //     /--Either--\          /--Either--\          /--Either--\     \
     //   /-E-\      /-E-\      /-E-\      /-E-\      /-E-\      /-E-\    \
     // 0000 0001  0010 0011  0100 0101  0110 0111  1000 1001  1010 1011 1100
     //
     // Note that a prefix of length D of the payload index acts as a path
     // through the tree to the node at depth D.  On the rightmost path
     // through the tree (when this prefix is equal to the corresponding
     // prefix of the maximum payload index), the bits of the index mark
     // where Eithers are required.
     //
     // Since we naturally want to build from the innermost Either out, and
     // therefore work with progressively shorter prefixes, we can do it all
     // with right-shifts.
     for (auto path = payloadIndex, maxPath = numPayloads - 1;
          maxPath != 0; path >>= 1, maxPath >>= 1) {
       // Skip making Eithers on the rightmost path where they aren't required.
       // This isn't just an optimization: adding spurious Eithers could
       // leave us with unresolvable type variables if `buildEither` has
       // a signature like:
       //    static func buildEither<T,U>(first value: T) -> Either<T,U>
       // which relies on unification to work.
       if (path == maxPath && !(maxPath & 1)) continue;

       bool isSecond = (path & 1);
       operand = buildCallIfWanted(operand->getStartLoc(),
                                   ctx.Id_buildEither, operand,
                                   {isSecond ? ctx.Id_second : ctx.Id_first},
                                   /*allowOneWay=*/false);
     }

     // Inject into Optional if required.  We'll be adding the call to
     // `buildIf` after all the recursive calls are complete.
     if (isOptional) {
       operand = buildSomeExpr(operand);
     }

     return operand;
   }

   Expr *buildSomeExpr(Expr *arg) {
     auto optionalDecl = ctx.getOptionalDecl();
     auto optionalType = optionalDecl->getDeclaredType();

     auto loc = arg->getStartLoc();
     auto optionalTypeExpr =
       TypeExpr::createImplicitHack(loc, optionalType, ctx);
     auto someRef = new (ctx) UnresolvedDotExpr(
         optionalTypeExpr, loc, ctx.getIdentifier("some"),
         DeclNameLoc(loc), /*implicit=*/true);
     return CallExpr::createImplicit(ctx, someRef, arg, { });
   }

   Expr *buildNoneExpr(SourceLoc endLoc) {
     auto optionalDecl = ctx.getOptionalDecl();
     auto optionalType = optionalDecl->getDeclaredType();

     auto optionalTypeExpr =
       TypeExpr::createImplicitHack(endLoc, optionalType, ctx);
     return new (ctx) UnresolvedDotExpr(
         optionalTypeExpr, endLoc, ctx.getIdentifier("none"),
         DeclNameLoc(endLoc), /*implicit=*/true);
   }

   CONTROL_FLOW_STMT(Guard)
   CONTROL_FLOW_STMT(While)
   CONTROL_FLOW_STMT(DoCatch)
   CONTROL_FLOW_STMT(RepeatWhile)
   CONTROL_FLOW_STMT(ForEach)
   CONTROL_FLOW_STMT(Switch)
   CONTROL_FLOW_STMT(Case)
   CONTROL_FLOW_STMT(Catch)
   CONTROL_FLOW_STMT(Break)
   CONTROL_FLOW_STMT(Continue)
   CONTROL_FLOW_STMT(Fallthrough)
   CONTROL_FLOW_STMT(Fail)
   CONTROL_FLOW_STMT(Throw)
   CONTROL_FLOW_STMT(PoundAssert)

 #undef CONTROL_FLOW_STMT
 };

 } // end anonymous namespace

 BraceStmt *
 TypeChecker::applyFunctionBuilderBodyTransform(FuncDecl *FD,
                                                BraceStmt *body,
                                                Type builderType) {
   // Try to build a single result expression.
   BuilderClosureVisitor visitor(Context, nullptr,
                                 /*wantExpr=*/true, builderType);
   Expr *returnExpr = visitor.visit(body);
   if (!returnExpr)
     return nullptr;

   // Make sure we have a usable result type for the body.
   Type returnType = AnyFunctionRef(FD).getBodyResultType();
   if (!returnType || returnType->hasError())
     return nullptr;

   auto loc = returnExpr->getStartLoc();
   auto returnStmt =
     new (Context) ReturnStmt(loc, returnExpr, /*implicit*/ true);
   return BraceStmt::create(Context, body->getLBraceLoc(), { returnStmt },
                            body->getRBraceLoc());
 }

 ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
     ClosureExpr *closure, Type builderType, ConstraintLocator *calleeLocator,
     ConstraintLocatorBuilder locator) {
   auto builder = builderType->getAnyNominal();
   assert(builder && "Bad function builder type");
   assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());

   // FIXME: Right now, single-expression closures suppress the function
   // builder translation.
   if (closure->hasSingleExpressionBody())
     return getTypeMatchSuccess();

   // Pre-check the closure body: pre-check any expressions in it and look
   // for return statements.
   switch (TC.preCheckFunctionBuilderClosureBody(closure)) {
   case FunctionBuilderClosurePreCheck::Okay:
     // If the pre-check was okay, apply the function-builder transform.
     break;

   case FunctionBuilderClosurePreCheck::Error:
     // If the pre-check had an error, flag that.
     return getTypeMatchFailure(locator);

   case FunctionBuilderClosurePreCheck::HasReturnStmt:
     // If the closure has a return statement, suppress the transform but
     // continue solving the constraint system.
     return getTypeMatchSuccess();
   }

   // Check the form of this closure to see if we can apply the
   // function-builder translation at all.
   {
     // Check whether we can apply this specific function builder.
     BuilderClosureVisitor visitor(getASTContext(), this,
                                   /*wantExpr=*/false, builderType);
     (void)visitor.visit(closure->getBody());

     // If we saw a control-flow statement or declaration that the builder
     // cannot handle, we don't have a well-formed function builder application.
     if (visitor.unhandledNode) {
       // If we aren't supposed to attempt fixes, fail.
       if (!shouldAttemptFixes()) {
         return getTypeMatchFailure(locator);
       }

       // Record the first unhandled construct as a fix.
       if (recordFix(
               SkipUnhandledConstructInFunctionBuilder::create(
                 *this, visitor.unhandledNode, builder,
                 getConstraintLocator(locator)))) {
         return getTypeMatchFailure(locator);
       }
     }
   }

   // If the builder type has a type parameter, substitute in the type
   // variables.
   if (builderType->hasTypeParameter()) {
     // Find the opened type for this callee and substitute in the type
     // parametes.
     for (const auto &opened : OpenedTypes) {
       if (opened.first == calleeLocator) {
         OpenedTypeMap replacements(opened.second.begin(),
                                    opened.second.end());
         builderType = openType(builderType, replacements);
         break;
       }
     }
     assert(!builderType->hasTypeParameter());
   }

   BuilderClosureVisitor visitor(getASTContext(), this,
                                 /*wantExpr=*/true, builderType);
   Expr *singleExpr = visitor.visit(closure->getBody());

   // We've already pre-checked all the original expressions, but do the
   // pre-check to the generated expression just to set up any preconditions
   // that CSGen might have.
   //
   // TODO: just build the AST the way we want it in the first place.
   if (TC.preCheckExpression(singleExpr, closure))
     return getTypeMatchFailure(locator);

   singleExpr = generateConstraints(singleExpr, closure);
   if (!singleExpr)
     return getTypeMatchFailure(locator);

   Type transformedType = getType(singleExpr);
   assert(transformedType && "Missing type");

   // Record the transformation.
   assert(std::find_if(
       builderTransformedClosures.begin(),
       builderTransformedClosures.end(),
       [&](const std::pair<ClosureExpr *, AppliedBuilderTransform> &elt) {
         return elt.first == closure;
       }) == builderTransformedClosures.end() &&
          "already transformed this closure along this path!?!");
   builderTransformedClosures.push_back(
       std::make_pair(closure,
                      AppliedBuilderTransform{builderType, singleExpr}));

   // Bind the result type of the closure to the type of the transformed
   // expression.
   Type closureType = getType(closure);
   auto fnType = closureType->castTo<FunctionType>();
   addConstraint(ConstraintKind::Equal, fnType->getResult(), transformedType,
                 locator);
   return getTypeMatchSuccess();
 }

 namespace {

 /// Pre-check all the expressions in the closure body.
 class PreCheckFunctionBuilderClosure : public ASTWalker {
   TypeChecker &TC;
   ClosureExpr *Closure;
   bool HasReturnStmt = false;
   bool HasError = false;
 public:
   PreCheckFunctionBuilderClosure(TypeChecker &tc, ClosureExpr *closure)
     : TC(tc), Closure(closure) {}

   FunctionBuilderClosurePreCheck run() {
     Stmt *oldBody = Closure->getBody();

     Stmt *newBody = oldBody->walk(*this);

     // If the walk was aborted, it was because we had a problem of some kind.
     assert((newBody == nullptr) == (HasError || HasReturnStmt) &&
            "unexpected short-circuit while walking closure body");
     if (!newBody) {
       if (HasError)
         return FunctionBuilderClosurePreCheck::Error;

       return FunctionBuilderClosurePreCheck::HasReturnStmt;
     }

     assert(oldBody == newBody && "pre-check walk wasn't in-place?");

     return FunctionBuilderClosurePreCheck::Okay;
   }

   std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
     // Pre-check the expression.  If this fails, abort the walk immediately.
     // Otherwise, replace the expression with the result of pre-checking.
     // In either case, don't recurse into the expression.
     if (TC.preCheckExpression(E, /*DC*/ Closure)) {
       HasError = true;
       return std::make_pair(false, nullptr);
     }

     return std::make_pair(false, E);
   }

   std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
     // If we see a return statement, abort the walk immediately.
     if (isa<ReturnStmt>(S)) {
       HasReturnStmt = true;
       return std::make_pair(false, nullptr);
     }

     // Otherwise, recurse into the statement normally.
     return std::make_pair(true, S);
   }
 };

 }

 FunctionBuilderClosurePreCheck
 TypeChecker::preCheckFunctionBuilderClosureBody(ClosureExpr *closure) {
   // Single-expression closures should already have been pre-checked.
   if (closure->hasSingleExpressionBody())
     return FunctionBuilderClosurePreCheck::Okay;

   // Check whether we've already done this analysis.
   auto it = precheckedFunctionBuilderClosures.find(closure);
   if (it != precheckedFunctionBuilderClosures.end())
     return it->second;

   auto result = PreCheckFunctionBuilderClosure(*this, closure).run();

   // Cache the result.
   precheckedFunctionBuilderClosures.insert(std::make_pair(closure, result));

   return result;
 }
	//===--- BuilderTransform.cpp - Function-builder transformation -----------===//
	//
	// This source file is part of the Swift.org open source project
	//
	// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
	// Licensed under Apache License v2.0 with Runtime Library Exception
	//
	// See https://swift.org/LICENSE.txt for license information
	// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
	//
	//===----------------------------------------------------------------------===//
	//
	// This file implements routines associated with the function-builder
	// transformation.
	//
	//===----------------------------------------------------------------------===//

	#include "ConstraintSystem.h"
	#include "TypeChecker.h"
	#include "swift/AST/ASTVisitor.h"
	#include "swift/AST/ASTWalker.h"
	#include "swift/AST/NameLookup.h"
	#include "swift/AST/NameLookupRequests.h"
	#include "swift/AST/ParameterList.h"
	#include "llvm/ADT/DenseMap.h"
	#include "llvm/ADT/SmallVector.h"
	#include <iterator>
	#include <map>
	#include <memory>
	#include <utility>
	#include <tuple>

	using namespace swift;
	using namespace constraints;

	namespace {

	/// Visitor to classify the contents of the given closure.
	class BuilderClosureVisitor
	: public StmtVisitor<BuilderClosureVisitor, Expr *> {
	ConstraintSystem *cs;
	ASTContext &ctx;
	bool wantExpr;
	Type builderType;
	NominalTypeDecl *builder = nullptr;
	llvm::SmallDenseMap<Identifier, bool> supportedOps;

	public:
	SkipUnhandledConstructInFunctionBuilder::UnhandledNode unhandledNode;

	private:
	/// Produce a builder call to the given named function with the given arguments.
	Expr *buildCallIfWanted(SourceLoc loc,
	Identifier fnName, ArrayRef<Expr *> args,
	ArrayRef<Identifier> argLabels,
	bool allowOneWay) {
	if (!wantExpr)
	return nullptr;

	// FIXME: Setting a TypeLoc on this expression is necessary in order
	// to get diagnostics if something about this builder call fails,
	// e.g. if there isn't a matching overload for `buildBlock`.
	// But we can only do this if there isn't a type variable in the type.
	TypeLoc typeLoc;
	if (!builderType->hasTypeVariable()) {
	typeLoc = TypeLoc(new (ctx) FixedTypeRepr(builderType, loc), builderType);
	}

	auto typeExpr = new (ctx) TypeExpr(typeLoc);
	if (cs) {
	cs->setType(typeExpr, MetatypeType::get(builderType));
	cs->setType(&typeExpr->getTypeLoc(), builderType);
	}

	SmallVector<SourceLoc, 4> argLabelLocs;
	for (auto i : indices(argLabels)) {
	argLabelLocs.push_back(args[i]->getStartLoc());
	}

	typeExpr->setImplicit();
	auto memberRef = new (ctx) UnresolvedDotExpr(
	typeExpr, loc, fnName, DeclNameLoc(loc), /implicit=/true);
	SourceLoc openLoc = args.empty() ? loc : args.front()->getStartLoc();
	SourceLoc closeLoc = args.empty() ? loc : args.back()->getEndLoc();
	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.
	bool builderSupports(Identifier fnName,
	ArrayRef<Identifier> argLabels = {}) {
	auto known = supportedOps.find(fnName);
	if (known != supportedOps.end()) {
	return known->second;
	}

	bool found = false;
	for (auto decl : builder->lookupDirect(fnName)) {
	if (auto func = dyn_cast<FuncDecl>(decl)) {
	// Function must be static.
	if (!func->isStatic())
	continue;

	// Function must have the right argument labels, if provided.
	if (!argLabels.empty()) {
	auto funcLabels = func->getFullName().getArgumentNames();
	if (argLabels.size() > funcLabels.size() \|\|
	funcLabels.slice(0, argLabels.size()) != argLabels)
	continue;
	}

	// Okay, it's a good-enough match.
	found = true;
	break;
	}
	}

	return supportedOps[fnName] = found;
	}

	public:
	BuilderClosureVisitor(ASTContext &ctx, ConstraintSystem *cs,
	bool wantExpr, Type builderType)
	: cs(cs), ctx(ctx), wantExpr(wantExpr), builderType(builderType) {
	assert((cs \|\| !builderType->hasTypeVariable()) &&
	"cannot handle builder type with type variables without "
	"constraint system");
	builder = builderType->getAnyNominal();
	}

	#define CONTROL_FLOW_STMT(StmtClass) \
	Expr visit##StmtClass##Stmt(StmtClass##Stmt stmt) { \
	if (!unhandledNode) \
	unhandledNode = stmt; \
	\
	return nullptr; \
	}

	Expr visitBraceStmt(BraceStmt braceStmt) {
	SmallVector<Expr *, 4> expressions;
	for (const auto &node : braceStmt->getElements()) {
	if (auto stmt = node.dyn_cast<Stmt *>()) {
	auto expr = visit(stmt);
	if (expr)
	expressions.push_back(expr);
	continue;
	}

	if (auto decl = node.dyn_cast<Decl *>()) {
	// Just ignore #if; the chosen children should appear in the
	// surrounding context. This isn't good for source tools but it
	// at least works.
	if (isa<IfConfigDecl>(decl))
	continue;

	if (!unhandledNode)
	unhandledNode = decl;

	continue;
	}

	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,
	/argLabels=/{ },
	/allowOneWay=/true);
	}

	Expr visitReturnStmt(ReturnStmt stmt) {
	// Allow implicit returns due to 'return' elision.
	if (!stmt->isImplicit() \|\| !stmt->hasResult()) {
	if (!unhandledNode)
	unhandledNode = stmt;
	return nullptr;
	}

	return stmt->getResult();
	}

	Expr visitDoStmt(DoStmt doStmt) {
	if (!builderSupports(ctx.Id_buildDo)) {
	if (!unhandledNode)
	unhandledNode = doStmt;
	return nullptr;
	}

	auto arg = visit(doStmt->getBody());
	if (!arg)
	return nullptr;

	return buildCallIfWanted(doStmt->getStartLoc(), ctx.Id_buildDo, arg,
	/argLabels=/{ }, /allowOneWay=/true);
	}

	CONTROL_FLOW_STMT(Yield)
	CONTROL_FLOW_STMT(Defer)

	static Expr *getTrivialBooleanCondition(StmtCondition condition) {
	if (condition.size() != 1)
	return nullptr;

	return condition.front().getBooleanOrNull();
	}

	static bool isBuildableIfChainRecursive(IfStmt *ifStmt,
	unsigned &numPayloads,
	bool &isOptional) {
	// The conditional must be trivial.
	if (!getTrivialBooleanCondition(ifStmt->getCond()))
	return false;

	// The 'then' clause contributes a payload.
	numPayloads++;

	// If there's an 'else' clause, it contributes payloads:
	if (auto elseStmt = ifStmt->getElseStmt()) {
	// If it's 'else if', it contributes payloads recursively.
	if (auto elseIfStmt = dyn_cast<IfStmt>(elseStmt)) {
	return isBuildableIfChainRecursive(elseIfStmt, numPayloads,
	isOptional);
	// Otherwise it's just the one.
	} else {
	numPayloads++;
	}

	// If not, the chain result is at least optional.
	} else {
	isOptional = true;
	}

	return true;
	}

	bool isBuildableIfChain(IfStmt *ifStmt, unsigned &numPayloads,
	bool &isOptional) {
	if (!isBuildableIfChainRecursive(ifStmt, numPayloads, isOptional))
	return false;

	// If there's a missing 'else', we need 'buildIf' to exist.
	if (isOptional && !builderSupports(ctx.Id_buildIf))
	return false;

	// If there are multiple clauses, we need 'buildEither(first:)' and
	// 'buildEither(second:)' to both exist.
	if (numPayloads > 1) {
	if (!builderSupports(ctx.Id_buildEither, {ctx.Id_first}) \|\|
	!builderSupports(ctx.Id_buildEither, {ctx.Id_second}))
	return false;
	}

	return true;
	}

	Expr visitIfStmt(IfStmt ifStmt) {
	// Check whether the chain is buildable and whether it terminates
	// without an `else`.
	bool isOptional = false;
	unsigned numPayloads = 0;
	if (!isBuildableIfChain(ifStmt, numPayloads, isOptional)) {
	if (!unhandledNode)
	unhandledNode = ifStmt;
	return nullptr;
	}

	// Attempt to build the chain, propagating short-circuits, which
	// might arise either do to error or not wanting an expression.
	auto chainExpr =
	buildIfChainRecursive(ifStmt, 0, numPayloads, isOptional);
	if (!chainExpr)
	return nullptr;
	assert(wantExpr);

	// The operand should have optional type if we had optional results,
	// 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,
	/argLabels=/{ },
	/allowOneWay=/true);
	} else if (ctx.LangOpts.FunctionBuilderOneWayConstraints) {
	// Form a one-way constraint to prevent backward propagation.
	chainExpr = new (ctx) OneWayExpr(chainExpr);
	}

	return chainExpr;
	}

	/// Recursively build an if-chain: build an expression which will have
	/// a value of the chain result type before any call to `buildIf`.
	/// The expression will perform any necessary calls to `buildEither`,
	/// and the result will have optional type if `isOptional` is true.
	Expr buildIfChainRecursive(IfStmt ifStmt, unsigned payloadIndex,
	unsigned numPayloads, bool isOptional) {
	assert(payloadIndex < numPayloads);
	// Make sure we recursively visit both sides even if we're not
	// building expressions.

	// Build the then clause. This will have the corresponding payload
	// type (i.e. not wrapped in any way).
	Expr *thenArg = visit(ifStmt->getThenStmt());

	// Build the else clause, if present. If this is from an else-if,
	// this will be fully wrapped; otherwise it will have the corresponding
	// payload type (at index `payloadIndex + 1`).
	assert(ifStmt->getElseStmt() \|\| isOptional);
	bool isElseIf = false;
	Optional<Expr *> elseChain;
	if (auto elseStmt = ifStmt->getElseStmt()) {
	if (auto elseIfStmt = dyn_cast<IfStmt>(elseStmt)) {
	isElseIf = true;
	elseChain = buildIfChainRecursive(elseIfStmt, payloadIndex + 1,
	numPayloads, isOptional);
	} else {
	elseChain = visit(elseStmt);
	}
	}

	// Short-circuit if appropriate.
	if (!wantExpr \|\| !thenArg \|\| (elseChain && !*elseChain))
	return nullptr;

	// Okay, build the conditional expression.

	// Prepare the `then` operand by wrapping it to produce a chain result.
	SourceLoc thenLoc = ifStmt->getThenStmt()->getStartLoc();
	Expr *thenExpr = buildWrappedChainPayload(thenArg, payloadIndex,
	numPayloads, isOptional);

	// Prepare the `else operand:
	Expr *elseExpr;
	SourceLoc elseLoc;

	// - If there's no `else` clause, use `Optional.none`.
	if (!elseChain) {
	assert(isOptional);
	elseLoc = ifStmt->getEndLoc();
	elseExpr = buildNoneExpr(elseLoc);

	// - If there's an `else if`, the chain expression from that
	// should already be producing a chain result.
	} else if (isElseIf) {
	elseExpr = *elseChain;
	elseLoc = ifStmt->getElseLoc();

	// - Otherwise, wrap it to produce a chain result.
	} else {
	elseLoc = ifStmt->getElseLoc();
	elseExpr = buildWrappedChainPayload(*elseChain,
	payloadIndex + 1, numPayloads,
	isOptional);
	}

	Expr *condition = getTrivialBooleanCondition(ifStmt->getCond());
	assert(condition && "checked by isBuildableIfChain");

	auto ifExpr = new (ctx) IfExpr(condition, thenLoc, thenExpr,
	elseLoc, elseExpr);
	ifExpr->setImplicit();
	return ifExpr;
	}

	/// Wrap a payload value in an expression which will produce a chain
	/// result (without `buildIf`).
	Expr buildWrappedChainPayload(Expr operand, unsigned payloadIndex,
	unsigned numPayloads, bool isOptional) {
	assert(payloadIndex < numPayloads);

	// Inject into the appropriate chain position.
	//
	// We produce a (left-biased) balanced binary tree of Eithers in order
	// to prevent requiring a linear number of injections in the worst case.
	// That is, if we have 13 clauses, we want to produce:
	//
	// /------------------Either------------\
	// /-------Either-------\ /--Either--\
	// /--Either--\ /--Either--\ /--Either--\ \
	// /-E-\ /-E-\ /-E-\ /-E-\ /-E-\ /-E-\ \
	// 0000 0001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100
	//
	// Note that a prefix of length D of the payload index acts as a path
	// through the tree to the node at depth D. On the rightmost path
	// through the tree (when this prefix is equal to the corresponding
	// prefix of the maximum payload index), the bits of the index mark
	// where Eithers are required.
	//
	// Since we naturally want to build from the innermost Either out, and
	// therefore work with progressively shorter prefixes, we can do it all
	// with right-shifts.
	for (auto path = payloadIndex, maxPath = numPayloads - 1;
	maxPath != 0; path >>= 1, maxPath >>= 1) {
	// Skip making Eithers on the rightmost path where they aren't required.
	// This isn't just an optimization: adding spurious Eithers could
	// leave us with unresolvable type variables if `buildEither` has
	// a signature like:
	// static func buildEither<T,U>(first value: T) -> Either<T,U>
	// which relies on unification to work.
	if (path == maxPath && !(maxPath & 1)) continue;

	bool isSecond = (path & 1);
	operand = buildCallIfWanted(operand->getStartLoc(),
	ctx.Id_buildEither, operand,
	{isSecond ? ctx.Id_second : ctx.Id_first},
	/allowOneWay=/false);
	}

	// Inject into Optional if required. We'll be adding the call to
	// `buildIf` after all the recursive calls are complete.
	if (isOptional) {
	operand = buildSomeExpr(operand);
	}

	return operand;
	}

	Expr buildSomeExpr(Expr arg) {
	auto optionalDecl = ctx.getOptionalDecl();
	auto optionalType = optionalDecl->getDeclaredType();

	auto loc = arg->getStartLoc();
	auto optionalTypeExpr =
	TypeExpr::createImplicitHack(loc, optionalType, ctx);
	auto someRef = new (ctx) UnresolvedDotExpr(
	optionalTypeExpr, loc, ctx.getIdentifier("some"),
	DeclNameLoc(loc), /implicit=/true);
	return CallExpr::createImplicit(ctx, someRef, arg, { });
	}

	Expr *buildNoneExpr(SourceLoc endLoc) {
	auto optionalDecl = ctx.getOptionalDecl();
	auto optionalType = optionalDecl->getDeclaredType();

	auto optionalTypeExpr =
	TypeExpr::createImplicitHack(endLoc, optionalType, ctx);
	return new (ctx) UnresolvedDotExpr(
	optionalTypeExpr, endLoc, ctx.getIdentifier("none"),
	DeclNameLoc(endLoc), /implicit=/true);
	}

	CONTROL_FLOW_STMT(Guard)
	CONTROL_FLOW_STMT(While)
	CONTROL_FLOW_STMT(DoCatch)
	CONTROL_FLOW_STMT(RepeatWhile)
	CONTROL_FLOW_STMT(ForEach)
	CONTROL_FLOW_STMT(Switch)
	CONTROL_FLOW_STMT(Case)
	CONTROL_FLOW_STMT(Catch)
	CONTROL_FLOW_STMT(Break)
	CONTROL_FLOW_STMT(Continue)
	CONTROL_FLOW_STMT(Fallthrough)
	CONTROL_FLOW_STMT(Fail)
	CONTROL_FLOW_STMT(Throw)
	CONTROL_FLOW_STMT(PoundAssert)

	#undef CONTROL_FLOW_STMT
	};

	} // end anonymous namespace

	BraceStmt *
	TypeChecker::applyFunctionBuilderBodyTransform(FuncDecl *FD,
	BraceStmt *body,
	Type builderType) {
	// Try to build a single result expression.
	BuilderClosureVisitor visitor(Context, nullptr,
	/wantExpr=/true, builderType);
	Expr *returnExpr = visitor.visit(body);
	if (!returnExpr)
	return nullptr;

	// Make sure we have a usable result type for the body.
	Type returnType = AnyFunctionRef(FD).getBodyResultType();
	if (!returnType \|\| returnType->hasError())
	return nullptr;

	auto loc = returnExpr->getStartLoc();
	auto returnStmt =
	new (Context) ReturnStmt(loc, returnExpr, /implicit/ true);
	return BraceStmt::create(Context, body->getLBraceLoc(), { returnStmt },
	body->getRBraceLoc());
	}

	ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
	ClosureExpr closure, Type builderType, ConstraintLocator calleeLocator,
	ConstraintLocatorBuilder locator) {
	auto builder = builderType->getAnyNominal();
	assert(builder && "Bad function builder type");
	assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());

	// FIXME: Right now, single-expression closures suppress the function
	// builder translation.
	if (closure->hasSingleExpressionBody())
	return getTypeMatchSuccess();

	// Pre-check the closure body: pre-check any expressions in it and look
	// for return statements.
	switch (TC.preCheckFunctionBuilderClosureBody(closure)) {
	case FunctionBuilderClosurePreCheck::Okay:
	// If the pre-check was okay, apply the function-builder transform.
	break;

	case FunctionBuilderClosurePreCheck::Error:
	// If the pre-check had an error, flag that.
	return getTypeMatchFailure(locator);

	case FunctionBuilderClosurePreCheck::HasReturnStmt:
	// If the closure has a return statement, suppress the transform but
	// continue solving the constraint system.
	return getTypeMatchSuccess();
	}

	// Check the form of this closure to see if we can apply the
	// function-builder translation at all.
	{
	// Check whether we can apply this specific function builder.
	BuilderClosureVisitor visitor(getASTContext(), this,
	/wantExpr=/false, builderType);
	(void)visitor.visit(closure->getBody());

	// If we saw a control-flow statement or declaration that the builder
	// cannot handle, we don't have a well-formed function builder application.
	if (visitor.unhandledNode) {
	// If we aren't supposed to attempt fixes, fail.
	if (!shouldAttemptFixes()) {
	return getTypeMatchFailure(locator);
	}

	// Record the first unhandled construct as a fix.
	if (recordFix(
	SkipUnhandledConstructInFunctionBuilder::create(
	*this, visitor.unhandledNode, builder,
	getConstraintLocator(locator)))) {
	return getTypeMatchFailure(locator);
	}
	}
	}

	// If the builder type has a type parameter, substitute in the type
	// variables.
	if (builderType->hasTypeParameter()) {
	// Find the opened type for this callee and substitute in the type
	// parametes.
	for (const auto &opened : OpenedTypes) {
	if (opened.first == calleeLocator) {
	OpenedTypeMap replacements(opened.second.begin(),
	opened.second.end());
	builderType = openType(builderType, replacements);
	break;
	}
	}
	assert(!builderType->hasTypeParameter());
	}

	BuilderClosureVisitor visitor(getASTContext(), this,
	/wantExpr=/true, builderType);
	Expr *singleExpr = visitor.visit(closure->getBody());

	// We've already pre-checked all the original expressions, but do the
	// pre-check to the generated expression just to set up any preconditions
	// that CSGen might have.
	//
	// TODO: just build the AST the way we want it in the first place.
	if (TC.preCheckExpression(singleExpr, closure))
	return getTypeMatchFailure(locator);

	singleExpr = generateConstraints(singleExpr, closure);
	if (!singleExpr)
	return getTypeMatchFailure(locator);

	Type transformedType = getType(singleExpr);
	assert(transformedType && "Missing type");

	// Record the transformation.
	assert(std::find_if(
	builderTransformedClosures.begin(),
	builderTransformedClosures.end(),
	[&](const std::pair<ClosureExpr *, AppliedBuilderTransform> &elt) {
	return elt.first == closure;
	}) == builderTransformedClosures.end() &&
	"already transformed this closure along this path!?!");
	builderTransformedClosures.push_back(
	std::make_pair(closure,
	AppliedBuilderTransform{builderType, singleExpr}));

	// Bind the result type of the closure to the type of the transformed
	// expression.
	Type closureType = getType(closure);
	auto fnType = closureType->castTo<FunctionType>();
	addConstraint(ConstraintKind::Equal, fnType->getResult(), transformedType,
	locator);
	return getTypeMatchSuccess();
	}

	namespace {

	/// Pre-check all the expressions in the closure body.
	class PreCheckFunctionBuilderClosure : public ASTWalker {
	TypeChecker &TC;
	ClosureExpr *Closure;
	bool HasReturnStmt = false;
	bool HasError = false;
	public:
	PreCheckFunctionBuilderClosure(TypeChecker &tc, ClosureExpr *closure)
	: TC(tc), Closure(closure) {}

	FunctionBuilderClosurePreCheck run() {
	Stmt *oldBody = Closure->getBody();

	Stmt newBody = oldBody->walk(this);

	// If the walk was aborted, it was because we had a problem of some kind.
	assert((newBody == nullptr) == (HasError \|\| HasReturnStmt) &&
	"unexpected short-circuit while walking closure body");
	if (!newBody) {
	if (HasError)
	return FunctionBuilderClosurePreCheck::Error;

	return FunctionBuilderClosurePreCheck::HasReturnStmt;
	}

	assert(oldBody == newBody && "pre-check walk wasn't in-place?");

	return FunctionBuilderClosurePreCheck::Okay;
	}

	std::pair<bool, Expr > walkToExprPre(Expr E) override {
	// Pre-check the expression. If this fails, abort the walk immediately.
	// Otherwise, replace the expression with the result of pre-checking.
	// In either case, don't recurse into the expression.
	if (TC.preCheckExpression(E, /DC/ Closure)) {
	HasError = true;
	return std::make_pair(false, nullptr);
	}

	return std::make_pair(false, E);
	}

	std::pair<bool, Stmt > walkToStmtPre(Stmt S) override {
	// If we see a return statement, abort the walk immediately.
	if (isa<ReturnStmt>(S)) {
	HasReturnStmt = true;
	return std::make_pair(false, nullptr);
	}

	// Otherwise, recurse into the statement normally.
	return std::make_pair(true, S);
	}
	};

	}

	FunctionBuilderClosurePreCheck
	TypeChecker::preCheckFunctionBuilderClosureBody(ClosureExpr *closure) {
	// Single-expression closures should already have been pre-checked.
	if (closure->hasSingleExpressionBody())
	return FunctionBuilderClosurePreCheck::Okay;

	// Check whether we've already done this analysis.
	auto it = precheckedFunctionBuilderClosures.find(closure);
	if (it != precheckedFunctionBuilderClosures.end())
	return it->second;

	auto result = PreCheckFunctionBuilderClosure(*this, closure).run();

	// Cache the result.
	precheckedFunctionBuilderClosures.insert(std::make_pair(closure, result));

	return result;
	}