[Custom Descriptors] Exact function imports

src/parser/context-decls.cpp

@@ -67,6 +67,7 @@ Result<> ParseDeclsCtx::addFunc(Name name,
                                 const std::vector<Name>& exports,
                                 ImportNames* import,
                                 TypeUseT type,
+                                Exactness exact,
                                 std::optional<LocalsT>,
                                 std::vector<Annotation>&& annotations,
                                 Index pos) {

src/parser/contexts.h

@@ -1078,6 +1078,7 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
                    const std::vector<Name>& exports,
                    ImportNames* import,
                    TypeUseT type,
+                   Exactness exact,
                    std::optional<LocalsT>,
                    std::vector<Annotation>&&,
                    Index pos);
@@ -1411,17 +1412,15 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
                    const std::vector<Name>&,
                    ImportNames*,
                    TypeUse type,
+                   Exactness exact,
                    std::optional<LocalsT> locals,
                    std::vector<Annotation>&&,
                    Index pos) {
     auto& f = wasm.functions[index];
     if (!type.type.isSignature()) {
       return in.err(pos, "expected signature type");
     }
-    f->type = f->type.with(type.type);
-    if (f->imported()) {
-      f->type = f->type.with(Inexact);
-    }
+    f->type = Type(type.type, NonNullable, exact);
     // If we are provided with too many names (more than the function has), we
     // will error on that later when we check the signature matches the type.
     // For now, avoid asserting in setLocalName.
@@ -1803,6 +1802,7 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                    const std::vector<Name>&,
                    ImportNames*,
                    TypeUseT,
+                   Exactness,
                    std::optional<LocalsT>,
                    std::vector<Annotation>&&,
                    Index) {

src/parser/parsers.h

@@ -370,6 +370,8 @@ Result<typename Ctx::LabelIdxT> labelidx(Ctx&, bool inDelegate = false);
 template<typename Ctx> Result<typename Ctx::TagIdxT> tagidx(Ctx&);
 template<typename Ctx>
 Result<typename Ctx::TypeUseT> typeuse(Ctx&, bool allowNames = true);
+template<typename Ctx>
+Result<std::pair<typename Ctx::TypeUse, Exactness>> exacttypeuse(Ctx&);
 MaybeResult<ImportNames> inlineImport(Lexer&);
 Result<std::vector<Name>> inlineExports(Lexer&);
 template<typename Ctx> Result<> comptype(Ctx&);
@@ -3007,6 +3009,24 @@ Result<typename Ctx::TypeUseT> typeuse(Ctx& ctx, bool allowNames) {
   return ctx.makeTypeUse(pos, type, namedParams.getPtr(), resultTypes.getPtr());
 }
 
+// exacttypeuse ::= typeuse |
+//                  '(' 'exact' typeuse ')'
+template<typename Ctx>
+Result<std::pair<typename Ctx::TypeUseT, Exactness>> exacttypeuse(Ctx& ctx) {
+  auto exact = Inexact;
+  if (ctx.in.takeSExprStart("exact"sv)) {
+    exact = Exact;
+  }
+  auto type = typeuse(ctx, true);
+  CHECK_ERR(type);
+  if (exact == Exact) {
+    if (!ctx.in.takeRParen()) {
+      return ctx.in.err("expected end of exact type use");
+    }
+  }
+  return std::make_pair(*type, exact);
+}
+
 // ('(' 'import' mod:name nm:name ')')?
 inline MaybeResult<ImportNames> inlineImport(Lexer& in) {
   if (!in.takeSExprStart("import"sv)) {
@@ -3221,7 +3241,7 @@ template<typename Ctx> MaybeResult<typename Ctx::LocalsT> locals(Ctx& ctx) {
 }
 
 // import ::= '(' 'import' mod:name nm:name importdesc ')'
-// importdesc ::= '(' 'func' id? typeuse ')'
+// importdesc ::= '(' 'func' id? exacttypeuse ')'
 //              | '(' 'table' id? tabletype ')'
 //              | '(' 'memory' id? memtype ')'
 //              | '(' 'global' id? globaltype ')'
@@ -3246,11 +3266,12 @@ template<typename Ctx> MaybeResult<> import_(Ctx& ctx) {
 
   if (ctx.in.takeSExprStart("func"sv)) {
     auto name = ctx.in.takeID();
-    auto type = typeuse(ctx);
-    CHECK_ERR(type);
+    auto use = exacttypeuse(ctx);
+    CHECK_ERR(use);
+    auto [type, exact] = *use;
     // TODO: function import annotations
     CHECK_ERR(ctx.addFunc(
-      name ? *name : Name{}, {}, &names, *type, std::nullopt, {}, pos));
+      name ? *name : Name{}, {}, &names, type, exact, std::nullopt, {}, pos));
   } else if (ctx.in.takeSExprStart("table"sv)) {
     auto name = ctx.in.takeID();
     auto type = tabletype(ctx);
@@ -3289,7 +3310,7 @@ template<typename Ctx> MaybeResult<> import_(Ctx& ctx) {
 // func ::= '(' 'func' id? ('(' 'export' name ')')*
 //              x,I:typeuse t*:vec(local) (in:instr)* ')'
 //        | '(' 'func' id? ('(' 'export' name ')')*
-//              '(' 'import' mod:name nm:name ')' typeuse ')'
+//              '(' 'import' mod:name nm:name ')' exacttypeuse ')'
 template<typename Ctx> MaybeResult<> func(Ctx& ctx) {
   auto pos = ctx.in.getPos();
   auto annotations = ctx.in.getAnnotations();
@@ -3309,11 +3330,19 @@ template<typename Ctx> MaybeResult<> func(Ctx& ctx) {
   auto import = inlineImport(ctx.in);
   CHECK_ERR(import);
 
-  auto type = typeuse(ctx);
-  CHECK_ERR(type);
-
+  typename Ctx::TypeUseT type;
+  Exactness exact = Exact;
   std::optional<typename Ctx::LocalsT> localVars;
-  if (!import) {
+
+  if (import) {
+    auto use = exacttypeuse(ctx);
+    CHECK_ERR(use);
+    type = use->first;
+    exact = use->second;
+  } else {
+    auto use = typeuse(ctx);
+    CHECK_ERR(use);
+    type = *use;
     if (auto l = locals(ctx)) {
       CHECK_ERR(l);
       localVars = *l;
@@ -3331,7 +3360,8 @@ template<typename Ctx> MaybeResult<> func(Ctx& ctx) {
   CHECK_ERR(ctx.addFunc(name,
                         *exports,
                         import.getPtr(),
-                        *type,
+                        type,
+                        exact,
                         localVars,
                         std::move(annotations),
                         pos));

src/passes/Print.cpp

@@ -3078,6 +3078,9 @@ void PrintSExpression::handleSignature(Function* curr,
   o << '(';
   printMajor(o, "func ");
   curr->name.print(o);
+  if (curr->imported() && curr->type.isExact()) {
+    o << " (exact";
+  }
   if ((currModule && currModule->features.hasGC()) ||
       requiresExplicitFuncType(curr->type.getHeapType())) {
     o << " (type ";
@@ -3124,6 +3127,9 @@ void PrintSExpression::handleSignature(Function* curr,
     o << maybeSpace;
     printResultType(curr->getResults());
   }
+  if (curr->imported() && curr->type.isExact()) {
+    o << ')';
+  }
 }
 
 void PrintSExpression::visitExport(Export* curr) {

src/passes/StackCheck.cpp

@@ -41,7 +41,8 @@ importStackOverflowHandler(Module& module, Name name, Signature sig) {
   ImportInfo info(module);
 
   if (!info.getImportedFunction(ENV, name)) {
-    auto import = Builder::makeFunction(name, sig, {});
+    auto import =
+      Builder::makeFunction(name, Type(sig, NonNullable, Inexact), {});
     import->module = ENV;
     import->base = name;
     module.addFunction(std::move(import));

src/wasm-binary.h

@@ -356,6 +356,8 @@ enum BrOnCastFlag {
   OutputNullable = 1 << 1,
 };
 
+constexpr uint32_t ExactImport = 1 << 5;
+
 enum EncodedType {
   // value types
   i32 = -0x1,  // 0x7f

src/wasm-builder.h

@@ -47,7 +47,7 @@ class Builder {
                                                 Type type,
                                                 std::vector<Type>&& vars,
                                                 Expression* body = nullptr) {
-    assert(type.isSignature());
+    assert(type.isSignature() && type.isNonNullable());
     auto func = std::make_unique<Function>();
     func->name = name;
     func->type = type;

src/wasm.h

@@ -2342,15 +2342,20 @@ class Function : public Importable {
   void clearDebugInfo();
 };
 
-// The kind of an import or export.
-enum class ExternalKind {
+// The kind of an import or export. Use a namespace to avoid polluting the wasm
+// namespace while maintaining implicit conversion to int, which an enum class
+// would not have.
+namespace ExternalKindImpl {
+enum Kind : uint32_t {
   Function = 0,
   Table = 1,
   Memory = 2,
   Global = 3,
   Tag = 4,
-  Invalid = -1
+  Invalid = uint32_t(-1)
 };
+} // namespace ExternalKindImpl
+using ExternalKind = ExternalKindImpl::Kind;
 
 // The kind of a top-level module item. (This overlaps with ExternalKind, but
 // C++ has no good way to extend an enum.) All such items are referred to by

src/wasm/wasm-binary.cpp

@@ -332,8 +332,11 @@ void WasmBinaryWriter::writeImports() {
   };
   ModuleUtils::iterImportedFunctions(*wasm, [&](Function* func) {
     writeImportHeader(func);
-    o << U32LEB(int32_t(ExternalKind::Function));
-    o << U32LEB(getTypeIndex(func->type.getHeapType()));
+    uint32_t kind = ExternalKind::Function;
+    if (func->type.isExact()) {
+      kind |= BinaryConsts::ExactImport;
+    }
+    o << U32LEB(kind) << U32LEB(getTypeIndex(func->type.getHeapType()));
   });
   ModuleUtils::iterImportedGlobals(*wasm, [&](Global* global) {
     writeImportHeader(global);
@@ -2875,12 +2878,13 @@ void WasmBinaryReader::readImports() {
   for (size_t i = 0; i < num; i++) {
     auto module = getInlineString();
     auto base = getInlineString();
-    auto kind = (ExternalKind)getU32LEB();
+    auto kind = getU32LEB();
     // We set a unique prefix for the name based on the kind. This ensures no
     // collisions between them, which can't occur here (due to the index i) but
     // could occur later due to the names section.
     switch (kind) {
-      case ExternalKind::Function: {
+      case ExternalKind::Function:
+      case ExternalKind::Function | BinaryConsts::ExactImport: {
         auto [name, isExplicit] =
           getOrMakeName(functionNames,
                         wasm.functions.size(),
@@ -2894,8 +2898,9 @@ void WasmBinaryReader::readImports() {
                      '.' + base.toString() +
                      "'s type must be a signature. Given: " + type.toString());
         }
+        auto exact = (kind & BinaryConsts::ExactImport) ? Exact : Inexact;
         auto curr =
-          builder.makeFunction(name, Type(type, NonNullable, Inexact), {});
+          builder.makeFunction(name, Type(type, NonNullable, exact), {});
         curr->hasExplicitName = isExplicit;
         curr->module = module;
         curr->base = base;
@@ -4691,8 +4696,8 @@ void WasmBinaryReader::readExports() {
     if (!names.emplace(name).second) {
       throwError("duplicate export name");
     }
-    ExternalKind kind = (ExternalKind)getU32LEB();
-    std::variant<Name, HeapType> value;
+    auto kind = getU32LEB();
+    std::optional<std::variant<Name, HeapType>> value;
     auto index = getU32LEB();
     switch (kind) {
       case ExternalKind::Function:
@@ -4711,9 +4716,12 @@ void WasmBinaryReader::readExports() {
         value = getTagName(index);
         break;
       case ExternalKind::Invalid:
-        throwError("invalid export kind");
+        break;
+    }
+    if (!value) {
+      throwError("invalid export kind");
     }
-    wasm.addExport(new Export(name, kind, value));
+    wasm.addExport(new Export(name, ExternalKind(kind), *value));
   }
 }
 

src/wasm/wasm-validator.cpp

@@ -4159,10 +4159,7 @@ void FunctionValidator::visitFunction(Function* curr) {
                curr->name,
                "all used types should be allowed");
 
-  if (curr->imported()) {
-    shouldBeTrue(
-      !curr->type.isExact(), curr->name, "imported function should be inexact");
-  } else {
+  if (!curr->imported()) {
     shouldBeTrue(
       curr->type.isExact(), curr->name, "defined function should be exact");
   }

test/lit/basic/exact-imports.wast

@@ -0,0 +1,97 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt %s -all -o %t.text.wast -g -S
+;; RUN: wasm-as %s -all -g -o %t.wasm
+;; RUN: wasm-dis %t.wasm -all -o %t.bin.wast
+;; RUN: wasm-as %s -all -o %t.nodebug.wasm
+;; RUN: wasm-dis %t.nodebug.wasm -all -o %t.bin.nodebug.wast
+;; RUN: cat %t.text.wast | filecheck %s --check-prefix=CHECK-TEXT
+;; RUN: cat %t.bin.wast | filecheck %s --check-prefix=CHECK-BIN
+;; RUN: cat %t.bin.nodebug.wast | filecheck %s --check-prefix=CHECK-BIN-NODEBUG
+
+(module
+  ;; CHECK-TEXT:      (type $f (func (param i32) (result i64)))
+  ;; CHECK-BIN:      (type $f (func (param i32) (result i64)))
+  (type $f (func (param i32) (result i64)))
+  ;; CHECK-TEXT:      (import "" "" (func $1 (exact (type $f) (param i32) (result i64))))
+  ;; CHECK-BIN:      (import "" "" (func $1 (exact (type $f) (param i32) (result i64))))
+  (import "" "" (func $1 (exact (type $f))))
+  ;; CHECK-TEXT:      (import "" "" (func $2 (exact (type $f) (param i32) (result i64))))
+  ;; CHECK-BIN:      (import "" "" (func $2 (exact (type $f) (param i32) (result i64))))
+  (import "" "" (func $2 (exact (type $f) (param i32) (result i64))))
+  ;; CHECK-TEXT:      (import "" "" (func $3 (exact (type $f) (param i32) (result i64))))
+  ;; CHECK-BIN:      (import "" "" (func $3 (exact (type $f) (param i32) (result i64))))
+  (import "" "" (func $3 (exact (param i32) (result i64))))
+
+  (func $4 (import "" "") (exact (type $f)))
+  (func $5 (import "" "") (exact (type $f) (param i32) (result i64)))
+  (func $6 (import "" "") (exact (param i32) (result i64)))
+
+  (global (ref (exact $f)) (ref.func $1))
+  (global (ref (exact $f)) (ref.func $2))
+  (global (ref (exact $f)) (ref.func $3))
+  (global (ref (exact $f)) (ref.func $4))
+  (global (ref (exact $f)) (ref.func $5))
+  (global (ref (exact $f)) (ref.func $6))
+)
+;; CHECK-TEXT:      (import "" "" (func $4 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-TEXT:      (import "" "" (func $5 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-TEXT:      (import "" "" (func $6 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-TEXT:      (global $global$0 (ref (exact $f)) (ref.func $1))
+
+;; CHECK-TEXT:      (global $global$1 (ref (exact $f)) (ref.func $2))
+
+;; CHECK-TEXT:      (global $global$2 (ref (exact $f)) (ref.func $3))
+
+;; CHECK-TEXT:      (global $global$3 (ref (exact $f)) (ref.func $4))
+
+;; CHECK-TEXT:      (global $global$4 (ref (exact $f)) (ref.func $5))
+
+;; CHECK-TEXT:      (global $global$5 (ref (exact $f)) (ref.func $6))
+
+;; CHECK-BIN:      (import "" "" (func $4 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-BIN:      (import "" "" (func $5 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-BIN:      (import "" "" (func $6 (exact (type $f) (param i32) (result i64))))
+
+;; CHECK-BIN:      (global $global$0 (ref (exact $f)) (ref.func $1))
+
+;; CHECK-BIN:      (global $global$1 (ref (exact $f)) (ref.func $2))
+
+;; CHECK-BIN:      (global $global$2 (ref (exact $f)) (ref.func $3))
+
+;; CHECK-BIN:      (global $global$3 (ref (exact $f)) (ref.func $4))
+
+;; CHECK-BIN:      (global $global$4 (ref (exact $f)) (ref.func $5))
+
+;; CHECK-BIN:      (global $global$5 (ref (exact $f)) (ref.func $6))
+
+;; CHECK-BIN-NODEBUG:      (type $0 (func (param i32) (result i64)))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$0 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$1 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$2 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$3 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$4 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (import "" "" (func $fimport$5 (exact (type $0) (param i32) (result i64))))
+
+;; CHECK-BIN-NODEBUG:      (global $global$0 (ref (exact $0)) (ref.func $fimport$0))
+
+;; CHECK-BIN-NODEBUG:      (global $global$1 (ref (exact $0)) (ref.func $fimport$1))
+
+;; CHECK-BIN-NODEBUG:      (global $global$2 (ref (exact $0)) (ref.func $fimport$2))
+
+;; CHECK-BIN-NODEBUG:      (global $global$3 (ref (exact $0)) (ref.func $fimport$3))
+
+;; CHECK-BIN-NODEBUG:      (global $global$4 (ref (exact $0)) (ref.func $fimport$4))
+
+;; CHECK-BIN-NODEBUG:      (global $global$5 (ref (exact $0)) (ref.func $fimport$5))