Merge pull request #68524 from kateinoigakukun/katei/expose-wasm-sym

kateinoigakukun · web-flow · commit 8cfcc245b532 · 2023-09-28T09:17:48.000-07:00
[wasm] Add `@_expose(wasm)` attribute for top-level functions
diff --git a/docs/ReferenceGuides/UnderscoredAttributes.md b/docs/ReferenceGuides/UnderscoredAttributes.md
@@ -427,6 +427,16 @@ in the generated C++ binding header.
 The optional "cxxName" string will be used as the name of
 the generated C++ declaration.
 
+### `_expose(wasm[, <"wasmExportName">])`
+
+Indicates that a particular function declaration should be
+exported from the linked WebAssembly.
+
+The optional "wasmExportName" string will be used as the
+the export name.
+
+It's the equivalent of clang's `__attribute__((export_name))`.
+
 ## `@_fixed_layout`
 
 Same as `@frozen` but also works for classes.
diff --git a/include/swift/AST/Attr.def b/include/swift/AST/Attr.def
@@ -398,7 +398,7 @@ SIMPLE_DECL_ATTR(_alwaysEmitConformanceMetadata, AlwaysEmitConformanceMetadata,
   OnProtocol | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   132)
 DECL_ATTR(_expose, Expose,
-  OnFunc | OnNominalType | OnVar | OnConstructor | LongAttribute | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
+  OnFunc | OnNominalType | OnVar | AllowMultipleAttributes | OnConstructor | LongAttribute | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   133)
 SIMPLE_DECL_ATTR(_spiOnly, SPIOnly,
   OnImport | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -172,6 +172,10 @@ class DeclAttribute : public AttributeBase {
       kind : 1
     );
 
+    SWIFT_INLINE_BITFIELD(ExposeAttr, DeclAttribute, NumExposureKindBits,
+      kind : NumExposureKindBits
+    );
+
     SWIFT_INLINE_BITFIELD(SynthesizedProtocolAttr, DeclAttribute, 1,
       isUnchecked : 1
     );
@@ -2309,15 +2313,23 @@ class BackDeployedAttr : public DeclAttribute {
 /// header used by C/C++ to interoperate with Swift.
 class ExposeAttr : public DeclAttribute {
 public:
-  ExposeAttr(StringRef Name, SourceLoc AtLoc, SourceRange Range, bool Implicit)
-      : DeclAttribute(DAK_Expose, AtLoc, Range, Implicit), Name(Name) {}
+  ExposeAttr(StringRef Name, SourceLoc AtLoc, SourceRange Range,
+             ExposureKind Kind, bool Implicit)
+      : DeclAttribute(DAK_Expose, AtLoc, Range, Implicit), Name(Name) {
+    Bits.ExposeAttr.kind = static_cast<unsigned>(Kind);
+  }
 
-  ExposeAttr(StringRef Name, bool Implicit)
-      : ExposeAttr(Name, SourceLoc(), SourceRange(), Implicit) {}
+  ExposeAttr(StringRef Name, ExposureKind Kind, bool Implicit)
+      : ExposeAttr(Name, SourceLoc(), SourceRange(), Kind, Implicit) {}
 
   /// The exposed declaration name.
   const StringRef Name;
 
+  /// Returns the kind of exposure.
+  ExposureKind getExposureKind() const {
+    return static_cast<ExposureKind>(Bits.ExposeAttr.kind);
+  }
+
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DAK_Expose;
   }
diff --git a/include/swift/AST/AttrKind.h b/include/swift/AST/AttrKind.h
@@ -102,6 +102,15 @@ enum class EffectsKind : uint8_t {
 enum : unsigned { NumEffectsKindBits =
   countBitsUsed(static_cast<unsigned>(EffectsKind::Last_EffectsKind)) };
 
+/// This enum represents the possible values of the @_expose attribute.
+enum class ExposureKind: uint8_t {
+  Cxx,
+  Wasm,
+  Last_ExposureKind = Wasm
+};
+
+enum : unsigned { NumExposureKindBits =
+  countBitsUsed(static_cast<unsigned>(ExposureKind::Last_ExposureKind)) };
   
 enum DeclAttrKind : unsigned {
 #define DECL_ATTR(_, NAME, ...) DAK_##NAME,
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -1838,6 +1838,9 @@ ERROR(section_not_at_top_level,none,
 ERROR(section_empty_name,none,
       "@_section section name cannot be empty", ())
 
+ERROR(expose_wasm_not_at_top_level_func,none,
+      "@_expose attribute with 'wasm' can only be applied to global functions", ())
+
 ERROR(expose_only_non_other_attr,none,
       "@_expose attribute cannot be applied to an '%0' declaration", (StringRef))
 ERROR(expose_inside_unexposed_decl,none,
diff --git a/include/swift/SIL/SILFunction.h b/include/swift/SIL/SILFunction.h
@@ -299,6 +299,10 @@ class SILFunction
   /// Name of a section if @_section attribute was used, otherwise empty.
   StringRef Section;
 
+  /// Name of a Wasm export if @_expose(wasm) attribute was used, otherwise
+  /// empty.
+  StringRef WasmExportName;
+
   /// Has value if there's a profile for this function
   /// Contains Function Entry Count
   ProfileCounter EntryCount;
@@ -1270,6 +1274,10 @@ class SILFunction
   StringRef section() const { return Section; }
   void setSection(StringRef value) { Section = value; }
 
+  /// Return Wasm export name if @_expose(wasm) was used, otherwise empty
+  StringRef wasmExportName() const { return WasmExportName; }
+  void setWasmExportName(StringRef value) { WasmExportName = value; }
+
   /// Returns true if this function belongs to a declaration that returns
   /// an opaque result type with one or more availability conditions that are
   /// allowed to produce a different underlying type at runtime.
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -1122,14 +1122,22 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer << "@_cdecl(\"" << cast<CDeclAttr>(this)->Name << "\")";
     break;
 
-  case DAK_Expose:
+  case DAK_Expose: {
     Printer.printAttrName("@_expose");
-    Printer << "(Cxx";
+    auto Attr = cast<ExposeAttr>(this);
+    switch (Attr->getExposureKind()) {
+    case ExposureKind::Wasm:
+      Printer << "(wasm";
+      break;
+    case ExposureKind::Cxx:
+      Printer << "(Cxx";
+      break;
+    }
     if (!cast<ExposeAttr>(this)->Name.empty())
       Printer << ", \"" << cast<ExposeAttr>(this)->Name << "\"";
     Printer << ")";
     break;
-
+  }
   case DAK_Section:
     Printer.printAttrName("@_section");
     Printer << "(\"" << cast<SectionAttr>(this)->Name << "\")";
diff --git a/lib/AST/SwiftNameTranslation.cpp b/lib/AST/SwiftNameTranslation.cpp
@@ -158,9 +158,9 @@ swift::cxx_translation::getNameForCxx(const ValueDecl *VD,
                                       CustomNamesOnly_t customNamesOnly) {
   ASTContext& ctx = VD->getASTContext();
 
-  if (const auto *Expose = VD->getAttrs().getAttribute<ExposeAttr>()) {
-    if (!Expose->Name.empty())
-      return Expose->Name;
+  for (auto *EA : VD->getAttrs().getAttributes<ExposeAttr>()) {
+    if (EA->getExposureKind() == ExposureKind::Cxx && !EA->Name.empty())
+      return EA->Name;
   }
 
   if (customNamesOnly)
diff --git a/lib/IRGen/GenDecl.cpp b/lib/IRGen/GenDecl.cpp
@@ -3561,6 +3561,11 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
     addUsedGlobal(fn);
   if (!f->section().empty())
     fn->setSection(f->section());
+  if (!f->wasmExportName().empty()) {
+    llvm::AttrBuilder attrBuilder(getLLVMContext());
+    attrBuilder.addAttribute("wasm-export-name", f->wasmExportName());
+    fn->addFnAttrs(attrBuilder);
+  }
 
   // If `hasCReferences` is true, then the function is either marked with
   // @_silgen_name OR @_cdecl.  If it is the latter, it must have a definition
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -3015,12 +3015,24 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
     }
 
     bool ParseSymbolName = true;
+
+    ExposureKind ExpKind;
     if (DK == DAK_Expose) {
-      if (Tok.isNot(tok::identifier) || Tok.getText() != "Cxx") {
+      auto diagnoseExpectOption = [&]() {
         diagnose(Tok.getLoc(), diag::attr_expected_option_such_as, AttrName,
                  "Cxx");
-        if (Tok.isNot(tok::identifier))
-          return makeParserSuccess();
+        ParseSymbolName = false;
+      };
+      if (Tok.isNot(tok::identifier)) {
+        diagnoseExpectOption();
+        return makeParserSuccess();
+      }
+      if (Tok.getText() == "Cxx") {
+        ExpKind = ExposureKind::Cxx;
+      } else if (Tok.getText() == "wasm") {
+        ExpKind = ExposureKind::Wasm;
+      } else {
+        diagnoseExpectOption();
         DiscardAttribute = true;
       }
       consumeToken(tok::identifier);
@@ -3076,10 +3088,19 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
       else if (DK == DAK_CDecl)
         Attributes.add(new (Context) CDeclAttr(AsmName.value(), AtLoc,
                                                AttrRange, /*Implicit=*/false));
-      else if (DK == DAK_Expose)
+      else if (DK == DAK_Expose) {
+        for (auto *EA : Attributes.getAttributes<ExposeAttr>()) {
+          // A single declaration cannot have two @_exported attributes with
+          // the same exposure kind.
+          if (EA->getExposureKind() == ExpKind) {
+            diagnose(Loc, diag::duplicate_attribute, false);
+            break;
+          }
+        }
         Attributes.add(new (Context) ExposeAttr(
             AsmName ? AsmName.value() : StringRef(""), AtLoc, AttrRange,
-            /*Implicit=*/false));
+            ExpKind, /*Implicit=*/false));
+      }
       else
         llvm_unreachable("out of sync with switch");
     }
diff --git a/lib/PrintAsClang/DeclAndTypePrinter.cpp b/lib/PrintAsClang/DeclAndTypePrinter.cpp
@@ -2825,8 +2825,10 @@ static bool hasExposeAttr(const ValueDecl *VD, bool isExtension = false) {
   // Clang decls don't need to be explicitly exposed.
   if (VD->hasClangNode())
     return true;
-  if (VD->getAttrs().hasAttribute<ExposeAttr>())
-    return true;
+  for (auto *EA : VD->getAttrs().getAttributes<ExposeAttr>()) {
+    if (EA->getExposureKind() == ExposureKind::Cxx)
+      return true;
+  }
   if (const auto *NMT = dyn_cast<NominalTypeDecl>(VD->getDeclContext()))
     return hasExposeAttr(NMT);
   if (const auto *ED = dyn_cast<ExtensionDecl>(VD->getDeclContext())) {
diff --git a/lib/SIL/IR/SILFunctionBuilder.cpp b/lib/SIL/IR/SILFunctionBuilder.cpp
@@ -218,6 +218,24 @@ void SILFunctionBuilder::addFunctionAttributes(
   if (Attrs.hasAttribute<SILGenNameAttr>() || Attrs.hasAttribute<CDeclAttr>())
     F->setHasCReferences(true);
 
+  for (auto *EA : Attrs.getAttributes<ExposeAttr>()) {
+    bool shouldExportDecl = true;
+    if (Attrs.hasAttribute<CDeclAttr>()) {
+      // If the function is marked with @cdecl, expose only C compatible
+      // thunk function.
+      shouldExportDecl = constant.isNativeToForeignThunk();
+    }
+    if (EA->getExposureKind() == ExposureKind::Wasm && shouldExportDecl) {
+      // A wasm-level exported function must be retained if it appears in a
+      // compilation unit.
+      F->setMarkedAsUsed(true);
+      if (EA->Name.empty())
+        F->setWasmExportName(F->getName());
+      else
+        F->setWasmExportName(EA->Name);
+    }
+  }
+
   if (Attrs.hasAttribute<UsedAttr>())
     F->setMarkedAsUsed(true);
 
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -2026,36 +2026,47 @@ void AttributeChecker::visitCDeclAttr(CDeclAttr *attr) {
 }
 
 void AttributeChecker::visitExposeAttr(ExposeAttr *attr) {
-  auto *VD = cast<ValueDecl>(D);
-  // Expose cannot be mixed with '@_cdecl' declarations.
-  if (VD->getAttrs().hasAttribute<CDeclAttr>())
-    diagnose(attr->getLocation(), diag::expose_only_non_other_attr, "@_cdecl");
-
-  // Nested exposed declarations are expected to be inside
-  // of other exposed declarations.
-  bool hasExpose = true;
-  const ValueDecl *decl = VD;
-  while (const NominalTypeDecl *NMT =
-             dyn_cast<NominalTypeDecl>(decl->getDeclContext())) {
-    decl = NMT;
-    hasExpose = NMT->getAttrs().hasAttribute<ExposeAttr>();
-  }
-  if (!hasExpose) {
-    diagnose(attr->getLocation(), diag::expose_inside_unexposed_decl, decl);
-  }
-
-  // Verify that the declaration is exposable.
-  auto repr = cxx_translation::getDeclRepresentation(VD);
-  if (repr.isUnsupported())
-    diagnose(attr->getLocation(),
-             cxx_translation::diagnoseRepresenationError(*repr.error, VD));
-
-  // Verify that the name mentioned in the expose
-  // attribute matches the supported name pattern.
-  if (!attr->Name.empty()) {
-    if (isa<ConstructorDecl>(VD) && !attr->Name.startswith("init"))
-      diagnose(attr->getLocation(), diag::expose_invalid_name_pattern_init,
-               attr->Name);
+  switch (attr->getExposureKind()) {
+  case ExposureKind::Wasm: {
+    // Only top-level func decls are currently supported.
+    if (!isa<FuncDecl>(D) || D->getDeclContext()->isTypeContext())
+      diagnose(attr->getLocation(), diag::expose_wasm_not_at_top_level_func);
+    break;
+  }
+  case ExposureKind::Cxx: {
+    auto *VD = cast<ValueDecl>(D);
+    // Expose cannot be mixed with '@_cdecl' declarations.
+    if (VD->getAttrs().hasAttribute<CDeclAttr>())
+      diagnose(attr->getLocation(), diag::expose_only_non_other_attr, "@_cdecl");
+
+    // Nested exposed declarations are expected to be inside
+    // of other exposed declarations.
+    bool hasExpose = true;
+    const ValueDecl *decl = VD;
+    while (const NominalTypeDecl *NMT =
+               dyn_cast<NominalTypeDecl>(decl->getDeclContext())) {
+      decl = NMT;
+      hasExpose = NMT->getAttrs().hasAttribute<ExposeAttr>();
+    }
+    if (!hasExpose) {
+      diagnose(attr->getLocation(), diag::expose_inside_unexposed_decl, decl);
+    }
+
+    // Verify that the declaration is exposable.
+    auto repr = cxx_translation::getDeclRepresentation(VD);
+    if (repr.isUnsupported())
+      diagnose(attr->getLocation(),
+               cxx_translation::diagnoseRepresenationError(*repr.error, VD));
+
+    // Verify that the name mentioned in the expose
+    // attribute matches the supported name pattern.
+    if (!attr->Name.empty()) {
+      if (isa<ConstructorDecl>(VD) && !attr->Name.startswith("init"))
+        diagnose(attr->getLocation(), diag::expose_invalid_name_pattern_init,
+                 attr->Name);
+    }
+    break;
+  }
   }
 }
 
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -5795,10 +5795,11 @@ llvm::Error DeclDeserializer::deserializeDeclCommon() {
       }
 
       case decls_block::Expose_DECL_ATTR: {
+        unsigned kind;
         bool isImplicit;
         serialization::decls_block::ExposeDeclAttrLayout::readRecord(
-            scratch, isImplicit);
-        Attr = new (ctx) ExposeAttr(blobData, isImplicit);
+            scratch, kind, isImplicit);
+        Attr = new (ctx) ExposeAttr(blobData, (ExposureKind)kind, isImplicit);
         break;
       }
 
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 807; // keypath accessor
+const uint16_t SWIFTMODULE_VERSION_MINOR = 808; // @_expose(wasm)
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2317,6 +2317,7 @@ namespace decls_block {
   >;
 
   using ExposeDeclAttrLayout = BCRecordLayout<Expose_DECL_ATTR,
+                                              BCFixed<1>, // exposure kind
                                               BCFixed<1>, // implicit flag
                                               BCBlob      // declaration name
                                               >;
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -3092,7 +3092,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       auto *theAttr = cast<ExposeAttr>(DA);
       auto abbrCode = S.DeclTypeAbbrCodes[ExposeDeclAttrLayout::Code];
       ExposeDeclAttrLayout::emitRecord(S.Out, S.ScratchRecord, abbrCode,
-                                       theAttr->isImplicit(), theAttr->Name);
+                                       (unsigned)theAttr->getExposureKind(), theAttr->isImplicit(), theAttr->Name);
       return;
     }
 
diff --git a/test/Interop/WebAssembly/expose-swift-decls-to-wasm.swift b/test/Interop/WebAssembly/expose-swift-decls-to-wasm.swift
diff --git a/test/attr/attr_expose.swift b/test/attr/attr_expose.swift
