Intrinsic: js.called (#8324)

When a function is annotated with this:

```wat
(@binaryen.js.called)
(func $foo ..)
```

then it is assumed to be called from JS (if a reference is taken).
In closed world, we assume such references are not actually
called from the outside, and this annotation can be used to mark
the exceptions that are.

Such functions are not normally exported, but "js called" in
that JS may call them, but not other wasm code (so rec
group type identity does not matter).

Author: kripken

src/ir/intrinsics.cpp

@@ -101,29 +101,29 @@ std::vector<Name> Intrinsics::getConfigureAllFunctions(Call* call) {
   return ret;
 }
 
-std::vector<Name> Intrinsics::getConfigureAllFunctions() {
-  // ConfigureAll in a start function makes its functions callable.
+std::vector<Name> Intrinsics::getJSCalledFunctions() {
+  std::vector<Name> ret;
+  for (auto& func : module.functions) {
+    if (getAnnotations(func.get()).jsCalled) {
+      ret.push_back(func->name);
+    }
+  }
+
+  // ConfigureAlls in a start function make their functions callable.
   if (module.start) {
     auto* start = module.getFunction(module.start);
     if (!start->imported()) {
       FindAll<Call> calls(start->body);
-      // Look for the (single) configureAll.
-      Call* configureAll = nullptr;
       for (auto* call : calls.list) {
         if (isConfigureAll(call)) {
-          if (configureAll) {
-            Fatal() << "Multiple configureAlls";
-          } else {
-            configureAll = call;
+          for (auto name : getConfigureAllFunctions(call)) {
+            ret.push_back(name);
           }
         }
       }
-      if (configureAll) {
-        return getConfigureAllFunctions(configureAll);
-      }
     }
   }
-  return {};
+  return ret;
 }
 
 } // namespace wasm

src/ir/intrinsics.h

@@ -109,8 +109,11 @@ class Intrinsics {
   //
   // where the segment $seg is of size N.
   std::vector<Name> getConfigureAllFunctions(Call* call);
-  // As above, but looks through the module to find the configureAll.
-  std::vector<Name> getConfigureAllFunctions();
+
+  // Returns the names of all functions that are JS-called. That includes ones
+  // in configureAll (which we look through the module for), and also those
+  // annotated with @binaryen.js.called.
+  std::vector<Name> getJSCalledFunctions();
 
   // Get the code annotations for an expression in a function.
   CodeAnnotation getAnnotations(Expression* curr, Function* func) {
@@ -149,6 +152,9 @@ class Intrinsics {
       if (!ret.removableIfUnused) {
         ret.removableIfUnused = funcAnnotations.removableIfUnused;
       }
+      if (!ret.jsCalled) {
+        ret.jsCalled = funcAnnotations.jsCalled;
+      }
     }
 
     return ret;

src/ir/possible-contents.cpp

@@ -2454,8 +2454,8 @@ Flower::Flower(Module& wasm, const PassOptions& options)
     }
   }
 
-  // configureAll functions are called from outside the module, as if exported.
-  for (auto func : Intrinsics(wasm).getConfigureAllFunctions()) {
+  // JS-called functions are called from outside the module, as if exported.
+  for (auto func : Intrinsics(wasm).getJSCalledFunctions()) {
     calledFromOutside.insert(func);
   }
 

src/parser/contexts.h

@@ -1317,6 +1317,8 @@ struct AnnotationParserCtx {
         inlineHint = &a;
       } else if (a.kind == Annotations::RemovableIfUnusedHint) {
         ret.removableIfUnused = true;
+      } else if (a.kind == Annotations::JSCalledHint) {
+        ret.jsCalled = true;
       }
     }
 

src/passes/GlobalTypeOptimization.cpp

@@ -420,7 +420,7 @@ struct GlobalTypeOptimization : public Pass {
     if (!wasm.features.hasCustomDescriptors()) {
       return;
     }
-    for (auto func : Intrinsics(wasm).getConfigureAllFunctions()) {
+    for (auto func : Intrinsics(wasm).getJSCalledFunctions()) {
       // Look at the result types being returned to JS and make sure we preserve
       // any configured prototypes they might expose.
       for (auto type : wasm.getFunction(func)->getResults()) {

src/passes/Print.cpp

@@ -2797,6 +2797,12 @@ void PrintSExpression::printCodeAnnotations(Expression* curr) {
       restoreNormalColor(o);
       doIndent(o, indent);
     }
+    if (annotation.jsCalled) {
+      Colors::grey(o);
+      o << "(@" << Annotations::JSCalledHint << ")\n";
+      restoreNormalColor(o);
+      doIndent(o, indent);
+    }
   }
 }
 

src/passes/RemoveUnusedModuleElements.cpp

@@ -184,7 +184,16 @@ struct Noter : public PostWalker<Noter, UnifiedExpressionVisitor<Noter>> {
     noteCallRef(curr->target->type.getHeapType());
   }
 
-  void visitRefFunc(RefFunc* curr) { noteRefFunc(curr->func); }
+  void visitRefFunc(RefFunc* curr) {
+    // If the target is js-called then a reference is as strong as a use.
+    auto target = curr->func;
+    Intrinsics intrinsics(*getModule());
+    if (intrinsics.getAnnotations(getModule()->getFunction(target)).jsCalled) {
+      use({ModuleElementKind::Function, target});
+    } else {
+      noteRefFunc(target);
+    }
+  }
 
   void visitStructGet(StructGet* curr) {
     if (curr->ref->type == Type::unreachable || curr->ref->type.isNull()) {

src/passes/SignaturePruning.cpp

@@ -176,9 +176,10 @@ struct SignaturePruning : public Pass {
       allInfo[tag->type].optimizable = false;
     }
 
-    // configureAll functions are signature-called, and must also not be
-    // modified.
-    for (auto func : Intrinsics(*module).getConfigureAllFunctions()) {
+    // Signature-called functions must also not be modified.
+    // TODO: Explore whether removing parameters from the end could be
+    //       beneficial (check if it does not regress call performance with JS).
+    for (auto func : Intrinsics(*module).getJSCalledFunctions()) {
       allInfo[module->getFunction(func)->type.getHeapType()].optimizable =
         false;
     }

src/passes/SignatureRefining.cpp

@@ -162,9 +162,8 @@ struct SignatureRefining : public Pass {
       }
     }
 
-    // configureAll functions are signature-called, which means their params
-    // must not be refined.
-    for (auto func : Intrinsics(*module).getConfigureAllFunctions()) {
+    // Signature-called functions must not have params refined.
+    for (auto func : Intrinsics(*module).getJSCalledFunctions()) {
       allInfo[module->getFunction(func)->type.getHeapType()].canModifyParams =
         false;
     }

src/passes/StripToolchainAnnotations.cpp

@@ -43,6 +43,7 @@ struct StripToolchainAnnotations
       // Remove the toolchain-specific annotations.
       auto& annotation = iter->second;
       annotation.removableIfUnused = false;
+      annotation.jsCalled = false;
 
       // If nothing remains, remove the entire annotation.
       if (annotation == CodeAnnotation()) {