[ConstraintSystem] Fix missing arguments

While trying to match function types, detect and fix any missing
arguments (by introducing type variables), such arguments would
get type information from corresponding parameters and aid in
producing solutions which are much easier to diagnose.

Author: xedin

include/swift/AST/Types.h

@@ -2762,6 +2762,8 @@ class AnyFunctionType : public TypeBase {
     bool operator!=(Param const &b) const { return !(*this == b); }
 
     Param getWithoutLabel() const { return Param(Ty, Identifier(), Flags); }
+
+    Param withType(Type newType) const { return Param(newType, Label, Flags); }
   };
 
   class CanParam : public Param {

lib/Sema/CSFix.cpp

@@ -342,3 +342,16 @@ AllowInvalidInitRef::create(RefKind kind, ConstraintSystem &cs, Type baseTy,
   return new (cs.getAllocator()) AllowInvalidInitRef(
       cs, kind, baseTy, init, isStaticallyDerived, baseRange, locator);
 }
+
+bool AddMissingArguments::diagnose(Expr *root, bool asNote) const {
+  return false;
+}
+
+AddMissingArguments *
+AddMissingArguments::create(ConstraintSystem &cs, FunctionType *funcType,
+                            llvm::ArrayRef<Param> synthesizedArgs,
+                            ConstraintLocator *locator) {
+  unsigned size = totalSizeToAlloc<Param>(synthesizedArgs.size());
+  void *mem = cs.getAllocator().Allocate(size, alignof(AddMissingArguments));
+  return new (mem) AddMissingArguments(cs, funcType, synthesizedArgs, locator);
+}

lib/Sema/CSFix.h

@@ -122,6 +122,10 @@ enum class FixKind : uint8_t {
   /// derived (rather than an arbitrary value of metatype type) or the
   /// referenced constructor must be required.
   AllowInvalidInitRef,
+
+  /// If there are fewer arguments than parameters, let's fix that up
+  /// by adding new arguments to the list represented as type variables.
+  AddMissingArguments,
 };
 
 class ConstraintFix {
@@ -614,6 +618,45 @@ class AllowInvalidInitRef final : public ConstraintFix {
                                      ConstraintLocator *locator);
 };
 
+class AddMissingArguments final
+    : public ConstraintFix,
+      private llvm::TrailingObjects<AddMissingArguments,
+                                    AnyFunctionType::Param> {
+  friend TrailingObjects;
+
+  using Param = AnyFunctionType::Param;
+
+  FunctionType *Fn;
+  unsigned NumSynthesized;
+
+  AddMissingArguments(ConstraintSystem &cs, FunctionType *funcType,
+                      llvm::ArrayRef<AnyFunctionType::Param> synthesizedArgs,
+                      ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::AddMissingArguments, locator), Fn(funcType),
+        NumSynthesized(synthesizedArgs.size()) {
+    std::uninitialized_copy(synthesizedArgs.begin(), synthesizedArgs.end(),
+                            getSynthesizedArgumentsBuf().begin());
+  }
+
+public:
+  std::string getName() const override { return "synthesize missing argument(s)"; }
+
+  ArrayRef<Param> getSynthesizedArguments() const {
+    return {getTrailingObjects<Param>(), NumSynthesized};
+  }
+
+  bool diagnose(Expr *root, bool asNote = false) const override;
+
+  static AddMissingArguments *create(ConstraintSystem &cs, FunctionType *fnType,
+                                     llvm::ArrayRef<Param> synthesizedArgs,
+                                     ConstraintLocator *locator);
+
+private:
+  MutableArrayRef<Param> getSynthesizedArgumentsBuf() {
+    return {getTrailingObjects<Param>(), NumSynthesized};
+  }
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSSimplify.cpp

@@ -1250,8 +1250,39 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
   auto argumentLocator = locator.withPathElement(
       ConstraintLocator::FunctionArgument);
 
-  if (func1Params.size() != func2Params.size())
-    return getTypeMatchFailure(argumentLocator);
+  int diff = func1Params.size() - func2Params.size();
+  if (diff != 0) {
+    if (!shouldAttemptFixes())
+      return getTypeMatchFailure(argumentLocator);
+
+    auto *anchor = locator.trySimplifyToExpr();
+    if (!anchor)
+      return getTypeMatchFailure(argumentLocator);
+
+    // If there are missing arguments, let's add them
+    // using parameter as a template.
+    if (diff < 0) {
+      for (unsigned i = func1Params.size(),
+                    n = func2Params.size(); i != n; ++i) {
+        auto *argLoc =
+            getConstraintLocator(anchor, LocatorPathElt::getTupleElement(i));
+
+        auto arg = func2Params[i].withType(createTypeVariable(argLoc));
+        func1Params.push_back(arg);
+      }
+
+      ArrayRef<AnyFunctionType::Param> argsRef(func1Params);
+      auto *fix = AddMissingArguments::create(*this, func2,
+                                              argsRef.take_back(abs(diff)),
+                                              getConstraintLocator(locator));
+
+      if (recordFix(fix))
+        return getTypeMatchFailure(argumentLocator);
+    } else {
+      // TODO(diagnostics): Add handling of extraneous arguments.
+      return getTypeMatchFailure(argumentLocator);
+    }
+  }
 
   bool hasLabelingFailures = false;
   for (unsigned i : indices(func1Params)) {
@@ -5628,6 +5659,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AllowTypeOrInstanceMember:
   case FixKind::AllowInvalidPartialApplication:
   case FixKind::AllowInvalidInitRef:
+  case FixKind::AddMissingArguments:
     llvm_unreachable("handled elsewhere");
   }