Support cross-module calls in the interpreter (#7787)

Previously a runtime instance of a function was just the function name.
That is fine for optimizations, and even for executing a single module,
but not cross-module calls.

To support that, add a FuncData (parallel to ExnData, ContData, etc.). The
constructor now requires that (but keep makeFunc as a simple helper
for cases when we just want the name).

A FuncData contains info about the module instance the function is in,
and a std::function to actually call it. We create a proper thunk when we
create a function literal, and in call_ref, we use it to do the actual call.

Author: kripken

src/ir/possible-contents.cpp

@@ -641,9 +641,9 @@ struct InfoCollector
     addRoot(curr);
   }
   void visitRefFunc(RefFunc* curr) {
-    addRoot(
-      curr,
-      PossibleContents::literal(Literal(curr->func, curr->type.getHeapType())));
+    addRoot(curr,
+            PossibleContents::literal(
+              Literal::makeFunc(curr->func, curr->type.getHeapType())));
 
     // The presence of a RefFunc indicates the function may be called
     // indirectly, so add the relevant connections for this particular function.
@@ -1800,7 +1800,7 @@ void TNHOracle::infer() {
         //       lot of other optimizations become possible anyhow.
         auto target = possibleTargets[0]->name;
         info.inferences[call->target] = PossibleContents::literal(
-          Literal(target, wasm.getFunction(target)->type));
+          Literal::makeFunc(target, wasm.getFunction(target)->type));
         continue;
       }
 

src/ir/properties.h

@@ -116,7 +116,7 @@ inline Literal getLiteral(const Expression* curr) {
   } else if (auto* n = curr->dynCast<RefNull>()) {
     return Literal(n->type);
   } else if (auto* r = curr->dynCast<RefFunc>()) {
-    return Literal(r->func, r->type.getHeapType());
+    return Literal::makeFunc(r->func, r->type.getHeapType());
   } else if (auto* i = curr->dynCast<RefI31>()) {
     if (auto* c = i->value->dynCast<Const>()) {
       return Literal::makeI31(c->value.geti32(),

src/literal.h

@@ -31,6 +31,7 @@
 namespace wasm {
 
 class Literals;
+struct FuncData;
 struct GCData;
 struct ExnData;
 struct ContData;
@@ -45,9 +46,8 @@ class Literal {
     int32_t i32;
     int64_t i64;
     uint8_t v128[16];
-    // funcref function name. `isNull()` indicates a `null` value.
-    // TODO: handle cross-module calls using something other than a Name here.
-    Name func;
+    // A reference to Function data.
+    std::shared_ptr<FuncData> funcData;
     // A reference to GC data, either a Struct or an Array. For both of those we
     // store the referred data as a Literals object (which is natural for an
     // Array, and for a Struct, is just the fields in order). The type is used
@@ -90,10 +90,7 @@ class Literal {
   explicit Literal(const std::array<Literal, 8>&);
   explicit Literal(const std::array<Literal, 4>&);
   explicit Literal(const std::array<Literal, 2>&);
-  explicit Literal(Name func, HeapType type)
-    : func(func), type(type, NonNullable, Exact) {
-    assert(type.isSignature());
-  }
+  explicit Literal(std::shared_ptr<FuncData> funcData, HeapType type);
   explicit Literal(std::shared_ptr<GCData> gcData, HeapType type);
   explicit Literal(std::shared_ptr<ExnData> exnData);
   explicit Literal(std::shared_ptr<ContData> contData);
@@ -253,9 +250,9 @@ class Literal {
   static Literal makeNull(HeapType type) {
     return Literal(Type(type.getBottom(), Nullable));
   }
-  static Literal makeFunc(Name func, HeapType type) {
-    return Literal(func, type);
-  }
+  // Simple way to create a function from the name and type, without a full
+  // FuncData.
+  static Literal makeFunc(Name func, HeapType type);
   static Literal makeI31(int32_t value, Shareability share) {
     auto lit = Literal(Type(HeapTypes::i31.getBasic(share), NonNullable));
     lit.i32 = value | 0x80000000;
@@ -311,10 +308,8 @@ class Literal {
     return bit_cast<double>(i64);
   }
   std::array<uint8_t, 16> getv128() const;
-  Name getFunc() const {
-    assert(type.isFunction() && !func.isNull());
-    return func;
-  }
+  Name getFunc() const;
+  std::shared_ptr<FuncData> getFuncData() const;
   std::shared_ptr<GCData> getGCData() const;
   std::shared_ptr<ExnData> getExnData() const;
   std::shared_ptr<ContData> getContData() const;

src/tools/wasm-ctor-eval.cpp

@@ -92,6 +92,18 @@ class EvallingModuleRunner : public ModuleRunnerBase<EvallingModuleRunner> {
     // internally), but we want to disable table.get for now.
     throw FailToEvalException("TODO: table.get");
   }
+
+  // This needs to be duplicated from ModuleRunner, unfortunately.
+  Literal makeFuncData(Name name, HeapType type) {
+    auto allocation =
+      std::make_shared<FuncData>(name, this, [this, name](Literals arguments) {
+        return callFunction(name, arguments);
+      });
+#if __has_feature(leak_sanitizer) || __has_feature(address_sanitizer)
+    __lsan_ignore_object(allocation.get());
+#endif
+    return Literal(allocation, type);
+  }
 };
 
 // Build an artificial `env` module based on a module's imports, so that the

src/wasm-interpreter.h

@@ -130,6 +130,35 @@ class Flow {
   }
 };
 
+struct FuncData {
+  // Name of the function in the module.
+  Name name;
+
+  // The interpreter instance we are in. This is only used for equality
+  // comparisons, as two functions are equal iff they have the same name and are
+  // in the same instance (in particular, we do *not* compare the |call| field
+  // below, which is an execution detail).
+  void* self;
+
+  // A way to execute this function. We use this when it is called.
+  using Call = std::function<Flow(Literals)>;
+  std::optional<Call> call;
+
+  FuncData(Name name,
+           void* self = nullptr,
+           std::optional<Call> call = std::nullopt)
+    : name(name), self(self), call(call) {}
+
+  bool operator==(const FuncData& other) const {
+    return name == other.name && self == other.self;
+  }
+
+  Flow doCall(Literals arguments) {
+    assert(call);
+    return (*call)(arguments);
+  }
+};
+
 // Suspend/resume support.
 //
 // As we operate directly on our structured IR, we do not have a program counter
@@ -277,6 +306,16 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
     Execute,
   };
 
+  Literal makeFuncData(Name name, HeapType type) {
+    // Identify the interpreter, but do not provide a way to actually call the
+    // function.
+    auto allocation = std::make_shared<FuncData>(name, this);
+#if __has_feature(leak_sanitizer) || __has_feature(address_sanitizer)
+    __lsan_ignore_object(allocation.get());
+#endif
+    return Literal(allocation, type);
+  }
+
 protected:
   RelaxedBehavior relaxedBehavior = RelaxedBehavior::NonConstant;
 
@@ -1801,7 +1840,7 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
     return Literal(int32_t(value.isNull()));
   }
   Flow visitRefFunc(RefFunc* curr) {
-    return Literal::makeFunc(curr->func, curr->type.getHeapType());
+    return self()->makeFuncData(curr->func, curr->type.getHeapType());
   }
   Flow visitRefEq(RefEq* curr) {
     Flow flow = visit(curr->left);
@@ -3545,7 +3584,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     if (curr->isReturn) {
       // Return calls are represented by their arguments followed by a reference
       // to the function to be called.
-      arguments.push_back(Literal::makeFunc(target, funcType));
+      arguments.push_back(self()->makeFuncData(target, funcType));
       return Flow(RETURN_CALL_FLOW, std::move(arguments));
     }
 
@@ -3651,7 +3690,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     std::cout << self()->indent() << "(calling ref " << targetRef.getFunc()
               << ")\n";
 #endif
-    Flow ret = callFunction(targetRef.getFunc(), arguments);
+    Flow ret = targetRef.getFuncData()->doCall(arguments);
 #if WASM_INTERPRETER_DEBUG
     std::cout << self()->indent() << "(returned to " << scope->function->name
               << ")\n";
@@ -5057,6 +5096,19 @@ class ModuleRunner : public ModuleRunnerBase<ModuleRunner> {
     ExternalInterface* externalInterface,
     std::map<Name, std::shared_ptr<ModuleRunner>> linkedInstances = {})
     : ModuleRunnerBase(wasm, externalInterface, linkedInstances) {}
+
+  Literal makeFuncData(Name name, HeapType type) {
+    // As the super's |makeFuncData|, but here we also provide a way to
+    // actually call the function.
+    auto allocation =
+      std::make_shared<FuncData>(name, this, [this, name](Literals arguments) {
+        return callFunction(name, arguments);
+      });
+#if __has_feature(leak_sanitizer) || __has_feature(address_sanitizer)
+    __lsan_ignore_object(allocation.get());
+#endif
+    return Literal(allocation, type);
+  }
 };
 
 } // namespace wasm

src/wasm/literal.cpp

@@ -71,6 +71,17 @@ Literal::Literal(const uint8_t init[16]) : type(Type::v128) {
   memcpy(&v128, init, 16);
 }
 
+Literal::Literal(std::shared_ptr<FuncData> funcData, HeapType type)
+  : funcData(funcData), type(type, NonNullable, Exact) {
+  assert(funcData);
+  assert(type.isSignature());
+}
+
+Literal Literal::makeFunc(Name func, HeapType type) {
+  // Provide only the name of the function, without execution info.
+  return Literal(std::make_shared<FuncData>(func), type);
+}
+
 Literal::Literal(std::shared_ptr<GCData> gcData, HeapType type)
   : gcData(gcData), type(type,
                          gcData ? NonNullable : Nullable,
@@ -140,7 +151,7 @@ Literal::Literal(const Literal& other) : type(other.type) {
     return;
   }
   if (type.isFunction()) {
-    func = other.func;
+    new (&funcData) std::shared_ptr<FuncData>(other.funcData);
     return;
   }
   if (type.isContinuation()) {
@@ -187,6 +198,8 @@ Literal::~Literal() {
   if (isNull() || isData() || type.getHeapType().isMaybeShared(HeapType::ext) ||
       type.getHeapType().isMaybeShared(HeapType::any)) {
     gcData.~shared_ptr();
+  } else if (isFunction()) {
+    funcData.~shared_ptr();
   } else if (isExn()) {
     exnData.~shared_ptr();
   } else if (isContinuation()) {
@@ -335,6 +348,16 @@ std::array<uint8_t, 16> Literal::getv128() const {
   return ret;
 }
 
+Name Literal::getFunc() const {
+  assert(isFunction());
+  return getFuncData()->name;
+}
+
+std::shared_ptr<FuncData> Literal::getFuncData() const {
+  assert(isFunction());
+  return funcData;
+}
+
 std::shared_ptr<GCData> Literal::getGCData() const {
   assert(isNull() || isData());
   return gcData;
@@ -458,8 +481,7 @@ bool Literal::operator==(const Literal& other) const {
       return true;
     }
     if (type.isFunction()) {
-      assert(func.is() && other.func.is());
-      return func == other.func;
+      return *funcData == *other.funcData;
     }
     if (type.isString()) {
       return gcData->values == other.gcData->values;

test/gtest/possible-contents.cpp

@@ -90,7 +90,7 @@ class PossibleContentsTest : public testing::Test {
     PossibleContents::global("funcGlobal", Type(HeapType::func, NonNullable));
 
   PossibleContents nonNullFunc = PossibleContents::literal(
-    Literal("func", Signature(Type::none, Type::none)));
+    Literal::makeFunc("func", Signature(Type::none, Type::none)));
 
   PossibleContents exactI32 = PossibleContents::exactType(Type::i32);
   PossibleContents exactAnyref = PossibleContents::exactType(anyref);

test/lit/ctor-eval/call_ref.wast

@@ -0,0 +1,34 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-ctor-eval %s --ctors=test --kept-exports=test --ignore-external-input --quiet -all -g -S -o - | filecheck %s
+
+;; Test we execute call_ref properly. After the call, the global will be set to
+;; 1.
+(module
+  ;; CHECK:      (type $f (func))
+  (type $f (func))
+
+  ;; CHECK:      (global $g (mut i32) (i32.const 1))
+  (global $g (export "gg") (mut i32) (i32.const 0))
+
+  (elem declare $test2)
+
+  (func $test (export "test")
+    (call_ref $f
+      (ref.func $called)
+    )
+  )
+
+  (func $called (type $f)
+    (global.set $g
+      (i32.const 1)
+    )
+  )
+)
+
+;; CHECK:      (export "gg" (global $g))
+
+;; CHECK:      (export "test" (func $test_2))
+
+;; CHECK:      (func $test_2 (type $f)
+;; CHECK-NEXT:  (nop)
+;; CHECK-NEXT: )

test/spec/call_ref_linked.wast

@@ -0,0 +1,39 @@
+;; Test that we can create a function in one module and call it using call_ref
+;; from another.
+
+(module
+  (type $func-i32 (func (result i32)))
+
+  (func $inner (result i32)
+    (i32.const 42)
+  )
+  (func $func-i32 (result i32)
+    ;; This call must execute in this module, i.e., call the right $inner.
+    (call $inner)
+  )
+  (func (export "get-func-i32") (result (ref $func-i32))
+    (return (ref.func $func-i32))
+  )
+)
+
+(register "first")
+
+(module
+  (type $func-i32 (func (result i32)))
+
+  (import "first" "get-func-i32" (func $imported (result (ref $func-i32))))
+
+  (func (export "run") (result i32)
+    (call_ref $func-i32
+      (call $imported)
+    )
+  )
+
+  (func $inner (result i32)
+    ;; We should not get here.
+    (unreachable)
+  )
+)
+
+(assert_return (invoke "run") (i32.const 42))
+