C++ Interop: import subscript operators

lib/ClangImporter/ImportDecl.cpp

@@ -3642,9 +3642,21 @@ namespace {
 
       result->setHasUnreferenceableStorage(hasUnreferenceableStorage);
 
-      if (cxxRecordDecl)
+      if (cxxRecordDecl) {
         result->setIsCxxNonTrivial(!cxxRecordDecl->isTriviallyCopyable());
 
+        for (auto &subscriptInfo : Impl.cxxSubscripts) {
+          auto declAndParameterType = subscriptInfo.first;
+          if (declAndParameterType.first != result)
+            continue;
+
+          auto getterAndSetter = subscriptInfo.second;
+          auto subscript = makeSubscript(getterAndSetter.first,
+                                         getterAndSetter.second);
+          result->addMember(subscript);
+        }
+      }
+
       return result;
     }
 
@@ -3917,12 +3929,10 @@ namespace {
       AbstractStorageDecl *owningStorage;
       switch (importedName.getAccessorKind()) {
       case ImportedAccessorKind::None:
-        owningStorage = nullptr;
-        break;
-
       case ImportedAccessorKind::SubscriptGetter:
       case ImportedAccessorKind::SubscriptSetter:
-        llvm_unreachable("Not possible for a function");
+        owningStorage = nullptr;
+        break;
 
       case ImportedAccessorKind::PropertyGetter: {
         auto property = getImplicitProperty(importedName, decl);
@@ -4122,6 +4132,30 @@ namespace {
             func->setImportAsStaticMember();
           }
         }
+
+        if (importedName.isSubscriptAccessor()) {
+          assert(func->getParameters()->size() == 1);
+          auto typeDecl = dc->getSelfNominalTypeDecl();
+          auto parameterType = func->getParameters()->get(0)->getType();
+          if (!typeDecl || !parameterType)
+            return nullptr;
+
+          auto &getterAndSetter = Impl.cxxSubscripts[{ typeDecl,
+                                                       parameterType }];
+
+          switch (importedName.getAccessorKind()) {
+          case ImportedAccessorKind::SubscriptGetter:
+            getterAndSetter.first = func;
+            break;
+          case ImportedAccessorKind::SubscriptSetter:
+            getterAndSetter.second = func;
+            break;
+          default:
+            llvm_unreachable("invalid subscript kind");
+          }
+
+          Impl.markUnavailable(func, "use subscript");
+        }
         // Someday, maybe this will need to be 'open' for C++ virtual methods.
         func->setAccess(AccessLevel::Public);
       }
@@ -4950,6 +4984,15 @@ namespace {
     SubscriptDecl *importSubscript(Decl *decl,
                                    const clang::ObjCMethodDecl *objcMethod);
 
+    /// Given either the getter, the setter, or both getter & setter
+    /// for a subscript operation, create the Swift subscript declaration.
+    ///
+    /// \param getter function returning `UnsafePointer<T>`
+    /// \param setter function returning `UnsafeMutablePointer<T>`
+    /// \return subscript declaration
+    SubscriptDecl *makeSubscript(FuncDecl *getter,
+                                 FuncDecl *setter);
+
     /// Import the accessor and its attributes.
     AccessorDecl *importAccessor(clang::ObjCMethodDecl *clangAccessor,
                                  AbstractStorageDecl *storage,
@@ -7339,6 +7382,244 @@ SwiftDeclConverter::importAccessor(clang::ObjCMethodDecl *clangAccessor,
   return accessor;
 }
 
+static InOutExpr *
+createInOutSelfExpr(AccessorDecl *accessorDecl) {
+  ASTContext &ctx = accessorDecl->getASTContext();
+
+  auto inoutSelfDecl = accessorDecl->getImplicitSelfDecl();
+  auto inoutSelfRefExpr =
+      new (ctx) DeclRefExpr(inoutSelfDecl, DeclNameLoc(),
+                            /*implicit=*/ true);
+  inoutSelfRefExpr->setType(LValueType::get(inoutSelfDecl->getInterfaceType()));
+
+  auto inoutSelfExpr =
+      new (ctx) InOutExpr(SourceLoc(),
+                          inoutSelfRefExpr,
+                          accessorDecl->mapTypeIntoContext(
+                              inoutSelfDecl->getValueInterfaceType()),
+                          /*isImplicit=*/ true);
+  inoutSelfExpr->setType(InOutType::get(inoutSelfDecl->getInterfaceType()));
+  return inoutSelfExpr;
+}
+
+static DeclRefExpr *
+createParamRefExpr(AccessorDecl *accessorDecl, unsigned index) {
+  ASTContext &ctx = accessorDecl->getASTContext();
+
+  auto paramDecl = accessorDecl->getParameters()->get(index);
+  auto paramRefExpr = new (ctx) DeclRefExpr(paramDecl,
+                                            DeclNameLoc(),
+                                            /*Implicit=*/ true);
+  paramRefExpr->setType(paramDecl->getType());
+  return paramRefExpr;
+}
+
+static CallExpr *
+createAccessorImplCallExpr(FuncDecl *accessorImpl,
+                           InOutExpr *inoutSelfExpr,
+                           DeclRefExpr *keyRefExpr) {
+  ASTContext &ctx = accessorImpl->getASTContext();
+
+  auto accessorImplExpr =
+      new (ctx) DeclRefExpr(ConcreteDeclRef(accessorImpl),
+                            DeclNameLoc(),
+                            /*Implicit=*/ true);
+  accessorImplExpr->setType(accessorImpl->getInterfaceType());
+
+  auto accessorImplDotCallExpr =
+      new (ctx) DotSyntaxCallExpr(accessorImplExpr,
+                                  SourceLoc(),
+                                  inoutSelfExpr);
+  accessorImplDotCallExpr->setType(accessorImpl->getMethodInterfaceType());
+  accessorImplDotCallExpr->setThrows(false);
+
+  auto *accessorImplCallExpr =
+      CallExpr::createImplicit(ctx, accessorImplDotCallExpr,
+                               { keyRefExpr }, { Identifier() });
+  accessorImplCallExpr->setType(accessorImpl->getResultInterfaceType());
+  accessorImplCallExpr->setThrows(false);
+  return accessorImplCallExpr;
+}
+
+/// Synthesizer callback for a subscript getter.
+static std::pair<BraceStmt *, bool>
+synthesizeSubscriptGetterBody(AbstractFunctionDecl *afd, void *context) {
+  auto getterDecl = cast<AccessorDecl>(afd);
+  auto getterImpl = static_cast<FuncDecl *>(context);
+
+  ASTContext &ctx = getterDecl->getASTContext();
+
+  InOutExpr *inoutSelfExpr = createInOutSelfExpr(getterDecl);
+  DeclRefExpr *keyRefExpr = createParamRefExpr(getterDecl, 0);
+
+  Type elementTy = getterDecl->getResultInterfaceType();
+
+  auto *getterImplCallExpr = createAccessorImplCallExpr(getterImpl,
+                                                        inoutSelfExpr,
+                                                        keyRefExpr);
+
+  // `getterImpl` can return either UnsafePointer or UnsafeMutablePointer.
+  // Retrieve the corresponding `.pointee` declaration.
+  PointerTypeKind ptrKind;
+  getterImpl->getResultInterfaceType()->getAnyPointerElementType(ptrKind);
+  VarDecl *pointeePropertyDecl = ctx.getPointerPointeePropertyDecl(ptrKind);
+
+  SubstitutionMap subMap =
+      SubstitutionMap::get(ctx.getUnsafePointerDecl()->getGenericSignature(),
+                           { elementTy }, { });
+  auto pointeePropertyRefExpr =
+      new (ctx) MemberRefExpr(getterImplCallExpr,
+                              SourceLoc(),
+                              ConcreteDeclRef(pointeePropertyDecl, subMap),
+                              DeclNameLoc(),
+                              /*implicit=*/ true);
+  pointeePropertyRefExpr->setType(elementTy);
+
+  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(),
+                                         pointeePropertyRefExpr,
+                                         /*implicit=*/ true);
+
+  auto body = BraceStmt::create(ctx, SourceLoc(), { returnStmt }, SourceLoc(),
+                                /*implicit=*/ true);
+  return { body, /*isTypeChecked=*/true };
+}
+
+/// Synthesizer callback for a subscript setter.
+static std::pair<BraceStmt *, bool>
+synthesizeSubscriptSetterBody(AbstractFunctionDecl *afd, void *context) {
+  auto setterDecl = cast<AccessorDecl>(afd);
+  auto setterImpl = static_cast<FuncDecl *>(context);
+
+  ASTContext &ctx = setterDecl->getASTContext();
+
+  InOutExpr *inoutSelfExpr = createInOutSelfExpr(setterDecl);
+  DeclRefExpr *valueParamRefExpr = createParamRefExpr(setterDecl, 0);
+  DeclRefExpr *keyParamRefExpr = createParamRefExpr(setterDecl, 1);
+
+  Type elementTy = valueParamRefExpr->getDecl()->getInterfaceType();
+
+  auto *setterImplCallExpr = createAccessorImplCallExpr(setterImpl,
+                                                        inoutSelfExpr,
+                                                        keyParamRefExpr);
+
+  VarDecl *pointeePropertyDecl = ctx.getPointerPointeePropertyDecl(PTK_UnsafeMutablePointer);
+
+  SubstitutionMap subMap =
+      SubstitutionMap::get(ctx.getUnsafeMutablePointerDecl()->getGenericSignature(),
+                           { elementTy }, { });
+  auto pointeePropertyRefExpr =
+      new (ctx) MemberRefExpr(setterImplCallExpr,
+                              SourceLoc(),
+                              ConcreteDeclRef(pointeePropertyDecl, subMap),
+                              DeclNameLoc(),
+                              /*implicit=*/ true);
+  pointeePropertyRefExpr->setType(LValueType::get(elementTy));
+
+  auto assignExpr = new (ctx) AssignExpr(pointeePropertyRefExpr,
+                                         SourceLoc(),
+                                         valueParamRefExpr,
+                                         /*implicit*/ true);
+  assignExpr->setType(TupleType::getEmpty(ctx));
+
+  auto body = BraceStmt::create(ctx, SourceLoc(), { assignExpr, }, SourceLoc());
+  return { body, /*isTypeChecked=*/true };
+}
+
+SubscriptDecl *
+SwiftDeclConverter::makeSubscript(FuncDecl *getter, FuncDecl *setter) {
+  assert((getter || setter) && "getter or setter required to generate subscript");
+
+  // If only a setter (imported from non-const `operator[]`) is defined,
+  // generate both get & set accessors from it.
+  FuncDecl *getterImpl = getter ? getter : setter;
+  FuncDecl *setterImpl = setter;
+
+  // Get the return type wrapped in `Unsafe(Mutable)Pointer<T>`.
+  const auto rawElementTy = getterImpl->getResultInterfaceType();
+  // Unwrap `T`.
+  const auto elementTy = rawElementTy->getAnyPointerElementType();
+
+  auto &ctx = Impl.SwiftContext;
+  auto bodyParams = getterImpl->getParameters();
+  DeclName name(ctx, DeclBaseName::createSubscript(), bodyParams);
+  auto dc = getterImpl->getDeclContext();
+
+  SubscriptDecl *subscript = SubscriptDecl::createImported(ctx,
+                                                           name,
+                                                           getterImpl->getLoc(),
+                                                           bodyParams,
+                                                           getterImpl->getLoc(),
+                                                           elementTy,
+                                                           dc,
+                                                           getterImpl->getClangNode());
+  subscript->setAccess(AccessLevel::Public);
+
+  AccessorDecl *getterDecl = AccessorDecl::create(ctx,
+                                                  getterImpl->getLoc(),
+                                                  getterImpl->getLoc(),
+                                                  AccessorKind::Get,
+                                                  subscript,
+                                                  SourceLoc(),
+                                                  subscript->getStaticSpelling(),
+                                                  false,
+                                                  SourceLoc(),
+                                                  nullptr,
+                                                  bodyParams,
+                                                  elementTy,
+                                                  dc);
+  getterDecl->setAccess(AccessLevel::Public);
+  getterDecl->setImplicit();
+  getterDecl->setIsDynamic(false);
+  getterDecl->setIsTransparent(true);
+  getterDecl->setBodySynthesizer(synthesizeSubscriptGetterBody, getterImpl);
+
+  if (getterImpl->isMutating()) {
+    getterDecl->setSelfAccessKind(SelfAccessKind::Mutating);
+    subscript->setIsGetterMutating(true);
+  }
+
+  AccessorDecl *setterDecl = nullptr;
+  if (setterImpl) {
+    auto paramVarDecl =
+        new (ctx) ParamDecl(SourceLoc(), SourceLoc(),
+                            Identifier(), SourceLoc(),
+                            ctx.getIdentifier("newValue"), dc);
+    paramVarDecl->setSpecifier(ParamSpecifier::Default);
+    paramVarDecl->setInterfaceType(elementTy);
+
+    auto setterParamList =
+        ParameterList::create(ctx, { paramVarDecl, bodyParams->get(0) });
+
+    setterDecl = AccessorDecl::create(ctx,
+                                      setterImpl->getLoc(),
+                                      setterImpl->getLoc(),
+                                      AccessorKind::Set,
+                                      subscript,
+                                      SourceLoc(),
+                                      subscript->getStaticSpelling(),
+                                      false,
+                                      SourceLoc(),
+                                      nullptr,
+                                      setterParamList,
+                                      TupleType::getEmpty(ctx),
+                                      dc);
+    setterDecl->setAccess(AccessLevel::Public);
+    setterDecl->setImplicit();
+    setterDecl->setIsDynamic(false);
+    setterDecl->setIsTransparent(true);
+    setterDecl->setBodySynthesizer(synthesizeSubscriptSetterBody, setterImpl);
+
+    if (setterImpl->isMutating()) {
+      setterDecl->setSelfAccessKind(SelfAccessKind::Mutating);
+      subscript->setIsSetterMutating(true);
+    }
+  }
+
+  makeComputed(subscript, getterDecl, setterDecl);
+
+  return subscript;
+}
+
 void SwiftDeclConverter::addProtocols(
     ProtocolDecl *protocol, SmallVectorImpl<ProtocolDecl *> &protocols,
     llvm::SmallPtrSetImpl<ProtocolDecl *> &known) {

lib/ClangImporter/ImportName.cpp

@@ -1797,6 +1797,23 @@ ImportedName NameImporter::importNameImpl(const clang::NamedDecl *D,
       isFunction = true;
       addEmptyArgNamesForClangFunction(functionDecl, argumentNames);
       break;
+    case clang::OverloadedOperatorKind::OO_Subscript: {
+      auto returnType = functionDecl->getReturnType();
+      if (!returnType->isReferenceType()) {
+        // TODO: support non-reference return types (SR-14351)
+        return ImportedName();
+      }
+      if (returnType->getPointeeType().isConstQualified()) {
+        baseName = "__operatorSubscriptConst";
+        result.info.accessorKind = ImportedAccessorKind::SubscriptGetter;
+      } else {
+        baseName = "__operatorSubscript";
+        result.info.accessorKind = ImportedAccessorKind::SubscriptSetter;
+      }
+      isFunction = true;
+      addEmptyArgNamesForClangFunction(functionDecl, argumentNames);
+      break;
+    }
     default:
       // We don't import these yet.
       return ImportedName();

lib/ClangImporter/ImporterImpl.h

@@ -490,6 +490,15 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   /// pairs.
   llvm::DenseMap<std::pair<FuncDecl *, FuncDecl *>, SubscriptDecl *> Subscripts;
 
+  /// Keep track of getter/setter pairs for functions imported from C++
+  /// subscript operators based on the type in which they are declared and
+  /// the type of their parameter.
+  ///
+  /// `.first` corresponds to a getter
+  /// `.second` corresponds to a setter
+  llvm::MapVector<std::pair<NominalTypeDecl *, Type>,
+                  std::pair<FuncDecl *, FuncDecl *>> cxxSubscripts;
+
   /// Keeps track of the Clang functions that have been turned into
   /// properties.
   llvm::DenseMap<const clang::FunctionDecl *, VarDecl *> FunctionsAsProperties;

lib/Sema/TypeCheckStorage.cpp

@@ -2311,6 +2311,9 @@ IsAccessorTransparentRequest::evaluate(Evaluator &evaluator,
           break;
         }
       }
+      if (auto subscript = dyn_cast<SubscriptDecl>(storage)) {
+        break;
+      }
 
       // Anything else should not have a synthesized setter.
       LLVM_FALLTHROUGH;