[Sema] NFC: Refactor AllowIUO policy into individual flags

Whether implicitly unwrapped optionals are allowed or not it is separable
from the contexts in which they are allowed.

Author: davezarzycki

lib/Sema/TypeCheckConstraints.cpp

@@ -1248,7 +1248,6 @@ bool PreCheckExpression::walkToClosureExprPre(ClosureExpr *closure) {
   options |= TypeResolutionFlags::AllowUnspecifiedTypes;
   options |= TypeResolutionFlags::AllowUnboundGenerics;
   options |= TypeResolutionFlags::InExpression;
-  options |= TypeResolutionFlags::AllowIUO;
   bool hadParameterError = false;
 
   GenericTypeToArchetypeResolver resolver(closure);

lib/Sema/TypeCheckDecl.cpp

@@ -1045,9 +1045,10 @@ static void validatePatternBindingEntry(TypeChecker &tc,
   // In particular, it's /not/ correct to check the PBD's DeclContext because
   // top-level variables in a script file are accessible from other files,
   // even though the PBD is inside a TopLevelCodeDecl.
-  TypeResolutionOptions options = TypeResolutionFlags::InExpression;
+  TypeResolutionOptions options = TypeResolutionFlags::PatternBindingEntry;
+  options |= TypeResolutionFlags::InExpression;
+  options |= TypeResolutionFlags::Direct;
 
-  options |= TypeResolutionFlags::AllowIUO;
   if (binding->getInit(entryNumber)) {
     // If we have an initializer, we can also have unknown types.
     options |= TypeResolutionFlags::AllowUnspecifiedTypes;

lib/Sema/TypeCheckGeneric.cpp

@@ -531,7 +531,8 @@ static bool checkGenericFuncSignature(TypeChecker &tc,
   if (auto fn = dyn_cast<FuncDecl>(func)) {
     if (!fn->getBodyResultTypeLoc().isNull()) {
       // Check the result type of the function.
-      TypeResolutionOptions options = TypeResolutionFlags::AllowIUO;
+      TypeResolutionOptions options = TypeResolutionFlags::FunctionResult;
+      options |= TypeResolutionFlags::Direct;
       if (fn->hasDynamicSelf())
         options |= TypeResolutionFlags::DynamicSelfResult;
 
@@ -899,8 +900,10 @@ static bool checkGenericSubscriptSignature(TypeChecker &tc,
                     &resolver);
 
   // Check the element type.
+  TypeResolutionOptions elementOptions = TypeResolutionFlags::Direct;
+  elementOptions |= TypeResolutionFlags::FunctionResult;
   badType |= tc.validateType(subscript->getElementTypeLoc(), subscript,
-                             TypeResolutionFlags::AllowIUO, &resolver);
+                             elementOptions, &resolver);
 
   // Infer requirements from it.
   if (genericParams && builder) {

lib/Sema/TypeCheckPattern.cpp

@@ -749,11 +749,12 @@ static bool validateParameterType(ParamDecl *decl, DeclContext *DC,
   // If the element is a variadic parameter, resolve the parameter type as if
   // it were in non-parameter position, since we want functions to be
   // @escaping in this case.
-  auto elementOptions = (options | (decl->isVariadic()
-                                    ? TypeResolutionFlags::VariadicFunctionInput
-                                    : TypeResolutionFlags::FunctionInput));
-  if (!decl->isVariadic())
-    elementOptions |= TypeResolutionFlags::AllowIUO;
+  options |= TypeResolutionFlags::Direct;
+  if (decl->isVariadic()) {
+    options |= TypeResolutionFlags::VariadicFunctionInput;
+  } else {
+    options |= TypeResolutionFlags::FunctionInput;
+  }
 
   bool hadError = false;
 
@@ -762,8 +763,7 @@ static bool validateParameterType(ParamDecl *decl, DeclContext *DC,
   // We might have a null typeLoc if this is a closure parameter list,
   // where parameters are allowed to elide their types.
   if (!TL.isNull()) {
-    hadError |= TC.validateType(TL, DC,
-                                elementOptions, &resolver);
+    hadError |= TC.validateType(TL, DC, options, &resolver);
   }
 
   auto *TR = TL.getTypeRepr();
@@ -773,7 +773,7 @@ static bool validateParameterType(ParamDecl *decl, DeclContext *DC,
   // If this is declared with '!' indicating that it is an Optional
   // that we should implicitly unwrap if doing so is required to type
   // check, then add an attribute to the decl.
-  if (elementOptions.contains(TypeResolutionFlags::AllowIUO) && TR &&
+  if (!decl->isVariadic() && TR &&
       TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional) {
     auto &C = DC->getASTContext();
     decl->getAttrs().add(

lib/Sema/TypeCheckType.cpp

@@ -319,7 +319,6 @@ adjustOptionsForGenericArgs(TypeResolutionOptions options) {
   options -= TypeResolutionFlags::FunctionInput;
   options -= TypeResolutionFlags::TypeAliasUnderlyingType;
   options -= TypeResolutionFlags::AllowUnavailableProtocol;
-  options -= TypeResolutionFlags::AllowIUO;
 
   return options;
 }
@@ -1515,7 +1514,8 @@ namespace {
     Type resolveOptionalType(OptionalTypeRepr *repr,
                              TypeResolutionOptions options);
     Type resolveImplicitlyUnwrappedOptionalType(ImplicitlyUnwrappedOptionalTypeRepr *repr,
-                                      TypeResolutionOptions options);
+                                                TypeResolutionOptions options,
+                                                bool isDirect);
     Type resolveTupleType(TupleTypeRepr *repr,
                           TypeResolutionOptions options);
     Type resolveCompositionType(CompositionTypeRepr *repr,
@@ -1567,14 +1567,21 @@ Type TypeResolver::resolveType(TypeRepr *repr, TypeResolutionOptions options) {
   // them.
   if (!isa<SpecifierTypeRepr>(repr) && !isa<TupleTypeRepr>(repr) &&
       !isa<AttributedTypeRepr>(repr) && !isa<FunctionTypeRepr>(repr) &&
-      !isa<IdentTypeRepr>(repr)) {
+      !isa<IdentTypeRepr>(repr) &&
+      !isa<ImplicitlyUnwrappedOptionalTypeRepr>(repr)) {
     options -= TypeResolutionFlags::FunctionInput;
     options -= TypeResolutionFlags::TypeAliasUnderlyingType;
   }
 
   if (Context.LangOpts.DisableAvailabilityChecking)
     options |= TypeResolutionFlags::AllowUnavailable;
 
+  bool isDirect = false;
+  if ((options & TypeResolutionFlags::Direct) && !isa<SpecifierTypeRepr>(repr)){
+    isDirect = true;
+    options -= TypeResolutionFlags::Direct;
+  }
+
   switch (repr->getKind()) {
   case TypeReprKind::Error:
     return ErrorType::get(Context);
@@ -1616,10 +1623,10 @@ Type TypeResolver::resolveType(TypeRepr *repr, TypeResolutionOptions options) {
   case TypeReprKind::Optional:
     return resolveOptionalType(cast<OptionalTypeRepr>(repr), options);
 
-  case TypeReprKind::ImplicitlyUnwrappedOptional:
-    return resolveImplicitlyUnwrappedOptionalType(
-             cast<ImplicitlyUnwrappedOptionalTypeRepr>(repr),
-             options);
+  case TypeReprKind::ImplicitlyUnwrappedOptional: {
+    auto iuoRepr = cast<ImplicitlyUnwrappedOptionalTypeRepr>(repr);
+    return resolveImplicitlyUnwrappedOptionalType(iuoRepr, options, isDirect);
+  }
 
   case TypeReprKind::Tuple:
     return resolveTupleType(cast<TupleTypeRepr>(repr), options);
@@ -2100,9 +2107,10 @@ bool TypeResolver::resolveASTFunctionTypeParams(
 Type TypeResolver::resolveASTFunctionType(FunctionTypeRepr *repr,
                                           TypeResolutionOptions options,
                                           FunctionType::ExtInfo extInfo) {
+  options -= TypeResolutionFlags::Direct;
   options -= TypeResolutionFlags::FunctionInput;
+  options -= TypeResolutionFlags::FunctionResult;
   options -= TypeResolutionFlags::TypeAliasUnderlyingType;
-  options -= TypeResolutionFlags::AllowIUO;
 
   SmallVector<AnyFunctionType::Param, 8> params;
   if (resolveASTFunctionTypeParams(repr->getArgsTypeRepr(), options,
@@ -2597,9 +2605,11 @@ Type TypeResolver::resolveSpecifierTypeRepr(SpecifierTypeRepr *repr,
     return ErrorType::get(Context);
   }
 
-  // Anything within the inout isn't a parameter anymore.
-  options -= TypeResolutionFlags::FunctionInput;
-  options -= TypeResolutionFlags::TypeAliasUnderlyingType;
+  if (!isa<ImplicitlyUnwrappedOptionalTypeRepr>(repr->getBase())) {
+    // Anything within the inout isn't a parameter anymore.
+    options -= TypeResolutionFlags::FunctionInput;
+    options -= TypeResolutionFlags::TypeAliasUnderlyingType;
+  }
 
   Type ty = resolveType(repr->getBase(), options);
   if (!ty || ty->hasError()) return ty;
@@ -2682,9 +2692,15 @@ Type TypeResolver::resolveOptionalType(OptionalTypeRepr *repr,
 }
 
 Type TypeResolver::resolveImplicitlyUnwrappedOptionalType(
-       ImplicitlyUnwrappedOptionalTypeRepr *repr,
-       TypeResolutionOptions options) {
-  if (!options.contains(TypeResolutionFlags::AllowIUO)) {
+      ImplicitlyUnwrappedOptionalTypeRepr *repr,
+      TypeResolutionOptions options,
+      bool isDirect) {
+  TypeResolutionOptions allowIUO = TypeResolutionFlags::SILType;
+  allowIUO |= TypeResolutionFlags::FunctionInput;
+  allowIUO |= TypeResolutionFlags::FunctionResult;
+  allowIUO |= TypeResolutionFlags::PatternBindingEntry;
+
+  if (!isDirect || !(options & allowIUO)) {
     // Prior to Swift 5, we allow 'as T!' and turn it into a disjunction.
     if (TC.Context.isSwiftVersionAtLeast(5)) {
       TC.diagnose(repr->getStartLoc(),
@@ -2724,11 +2740,9 @@ Type TypeResolver::resolveTupleType(TupleTypeRepr *repr,
   SmallVector<TupleTypeElt, 8> elements;
   elements.reserve(repr->getNumElements());
 
-
   auto elementOptions = options;
   if (repr->isParenType()) {
-    // We also want to disallow IUO within even a paren.
-    elementOptions -= TypeResolutionFlags::AllowIUO;
+    elementOptions -= TypeResolutionFlags::Direct;
   } else {
     elementOptions = withoutContext(elementOptions, true);
   }

lib/Sema/TypeChecker.cpp

@@ -839,7 +839,6 @@ bool swift::performTypeLocChecking(ASTContext &Ctx, TypeLoc &T,
   options |= TypeResolutionFlags::AllowUnboundGenerics;
   if (isSILMode) {
     options |= TypeResolutionFlags::SILMode;
-    options |= TypeResolutionFlags::AllowIUO;
   }
   if (isSILType)
     options |= TypeResolutionFlags::SILType;

lib/Sema/TypeChecker.h

@@ -479,6 +479,7 @@ enum class TypeResolutionFlags : unsigned {
 
   /// Whether we are in the input type of a function, or under one level of
   /// tuple type.  This is not set for multi-level tuple arguments.
+  /// See also: TypeResolutionFlags::Direct
   FunctionInput = 1 << 5,
 
   /// Whether this is a variadic function input.
@@ -543,8 +544,8 @@ enum class TypeResolutionFlags : unsigned {
   /// Whether we are checking the parameter list of a subscript.
   SubscriptParameters = 1 << 23,
 
-  /// Is it okay to resolve an IUO sigil ("!") here?
-  AllowIUO = 1 << 24,
+  /// Whether we are at the direct base of a type expression.
+  Direct = 1 << 24,
 
   /// Whether this type is being used in a cast or coercion expression.
   InCastOrCoercionExpression = 1 << 25,
@@ -557,6 +558,13 @@ enum class TypeResolutionFlags : unsigned {
 
   /// Whether we should not produce diagnostics if the type is invalid.
   SilenceErrors = 1 << 28,
+
+  /// Whether we are in the result type of a function, including multi-level
+  /// tuple return values. See also: TypeResolutionFlags::Direct
+  FunctionResult = 1 << 29,
+
+  /// Whether this is a pattern binding entry.
+  PatternBindingEntry = 1 << 30,
 };
 
 /// Option set describing how type resolution should work.
@@ -565,12 +573,14 @@ using TypeResolutionOptions = OptionSet<TypeResolutionFlags>;
 /// Strip the contextual options from the given type resolution options.
 static inline TypeResolutionOptions
 withoutContext(TypeResolutionOptions options, bool preserveSIL = false) {
+  options -= TypeResolutionFlags::Direct;
   options -= TypeResolutionFlags::FunctionInput;
   options -= TypeResolutionFlags::VariadicFunctionInput;
+  options -= TypeResolutionFlags::FunctionResult;
   options -= TypeResolutionFlags::EnumCase;
   options -= TypeResolutionFlags::SubscriptParameters;
   options -= TypeResolutionFlags::ImmediateOptionalTypeArgument;
-  options -= TypeResolutionFlags::AllowIUO;
+  options -= TypeResolutionFlags::PatternBindingEntry;
   if (!preserveSIL) options -= TypeResolutionFlags::SILType;
   return options;
 }