Infer parameters type of the traced function from the original function

Author: loganek

src/passes/TraceCalls.cpp

@@ -37,8 +37,8 @@
 //     )
 //   )
 //
-// Instrumenting void* malloc(int32_t) with a user-defined tracer
-// (trace_allocation):
+// Instrumenting void* malloc(int32_t) with a user-defined name
+// of the tracer (trace_allocation):
 //  Before:
 //   (call $malloc
 //     (local.const 32))
@@ -58,7 +58,7 @@
 //   )
 //
 
-#include <set>
+#include <unordered_map>
 
 #include "asmjs/shared-constants.h"
 #include "ir/import-utils.h"
@@ -68,45 +68,18 @@
 
 namespace wasm {
 
-struct TracedFunction {
-  Name originName;
-  Name tracerName;
-  Type params;
-  Type results;
-};
-bool operator<(const TracedFunction& lhs, const TracedFunction& rhs) {
-  return lhs.originName < rhs.originName;
-}
-
-static Type strToType(const std::string& strType) {
-  if (strType == "i32") {
-    return Type::i32;
-  } else if (strType == "i64") {
-    return Type::i64;
-  } else if (strType == "f32") {
-    return Type::f32;
-  } else if (strType == "f64") {
-    return Type::f64;
-  } else if (strType == "v128") {
-    return Type::v128;
-  } else {
-    Fatal() << "Failed to parse type '" << strType << "'";
-  }
-}
+using TracedFunctions =
+  std::unordered_map<Name /* originName */, Name /* tracerName */>;
 
 struct AddTraceWrappers : public WalkerPass<PostWalker<AddTraceWrappers>> {
-  AddTraceWrappers(std::set<TracedFunction> tracedFunctions)
+  AddTraceWrappers(TracedFunctions tracedFunctions)
     : tracedFunctions(std::move(tracedFunctions)) {}
   void visitCall(Call* curr) {
     auto* target = getModule()->getFunction(curr->target);
 
-    auto iter = std::find_if(tracedFunctions.begin(),
-                             tracedFunctions.end(),
-                             [target](const TracedFunction& f) {
-                               return f.originName == target->name;
-                             });
+    auto iter = tracedFunctions.find(target->name);
     if (iter != tracedFunctions.end()) {
-      addInstrumentation(curr, target, iter->tracerName);
+      addInstrumentation(curr, target, iter->second);
     }
   }
 
@@ -143,7 +116,7 @@ struct AddTraceWrappers : public WalkerPass<PostWalker<AddTraceWrappers>> {
            wrapperName, trackerCallParams, Type::BasicType::none)}));
   }
 
-  std::set<TracedFunction> tracedFunctions;
+  TracedFunctions tracedFunctions;
 };
 
 struct TraceCalls : public Pass {
@@ -154,39 +127,55 @@ struct TraceCalls : public Pass {
     auto functionsDefinitions = getPassOptions().getArgument(
       "trace-calls",
       "TraceCalls usage: wasm-opt "
-      "--trace-calls=FUNCTION_TO_TRACE[:TRACER_NAME][,RESULT1_TYPE]"
-      "[,PARAM1_TYPE[,PARAM2_TYPE[,...]]][;...]");
+      "--trace-calls=FUNCTION_TO_TRACE[:TRACER_NAME][,...]");
 
     auto tracedFunctions = parseArgument(functionsDefinitions);
 
     for (const auto& tracedFunction : tracedFunctions) {
-      addImport(module, tracedFunction);
+      auto func = module->getFunctionOrNull(tracedFunction.first);
+      if (!func) {
+        std::cerr << "[TraceCalls] Function '" << tracedFunction.first
+                  << "' not found";
+      } else {
+        addImport(module, *func, tracedFunction.second);
+      }
     }
 
     AddTraceWrappers(std::move(tracedFunctions)).run(getPassRunner(), module);
   }
 
 private:
-  std::set<TracedFunction> parseArgument(const std::string& arg) {
-    std::set<TracedFunction> tracedFunctions;
+  Type getTracerParamsType(ImportInfo& info, const Function& func) {
+    auto resultsType = func.type.getSignature().results;
+    if (resultsType.isTuple()) {
+      Fatal() << "Failed to instrument function '" << func.name
+              << "': Multi-value result type is not supported";
+    }
+
+    std::vector<Type> tracerParamTypes;
+    if (resultsType.isConcrete()) {
+      tracerParamTypes.push_back(resultsType);
+    }
+    for (auto& op : func.type.getSignature().params) {
+      tracerParamTypes.push_back(op);
+    }
+
+    return Type(tracerParamTypes);
+  }
 
-    for (const auto& definition : String::Split(arg, ";")) {
-      auto parts = String::Split(definition, ",");
-      if (parts.size() == 0) {
+  TracedFunctions parseArgument(const std::string& arg) {
+    TracedFunctions tracedFunctions;
+
+    for (const auto& definition : String::Split(arg, ",")) {
+      if (definition.empty()) {
         // Empty definition, ignore.
         continue;
       }
 
       std::string originName, traceName;
-      parseFunctionName(parts[0], originName, traceName);
-
-      std::vector<Type> paramsAndResults;
-      for (size_t i = 1; i < parts.size(); i++) {
-        paramsAndResults.push_back(strToType(parts[i]));
-      }
+      parseFunctionName(definition, originName, traceName);
 
-      tracedFunctions.emplace(TracedFunction{
-        Name(originName), Name(traceName), paramsAndResults, Type::none});
+      tracedFunctions[Name(originName)] = Name(traceName);
     }
 
     return tracedFunctions;
@@ -211,14 +200,14 @@ struct TraceCalls : public Pass {
     }
   }
 
-  void addImport(Module* wasm, const TracedFunction& f) {
+  void addImport(Module* wasm, const Function& f, const Name& tracerName) {
     ImportInfo info(*wasm);
 
-    if (!info.getImportedFunction(ENV, f.tracerName)) {
-      auto import =
-        Builder::makeFunction(f.tracerName, Signature(f.params, f.results), {});
+    if (!info.getImportedFunction(ENV, tracerName)) {
+      auto import = Builder::makeFunction(
+        tracerName, Signature(getTracerParamsType(info, f), Type::none), {});
       import->module = ENV;
-      import->base = f.tracerName;
+      import->base = tracerName;
       wasm->addFunction(std::move(import));
     }
   }

test/lit/passes/trace-calls.wast

@@ -1,6 +1,6 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 
-;; RUN: foreach %s %t wasm-opt --enable-simd --trace-calls="noparamsnoresults;singleparamnoresults,f64;multiparamsnoresults:tracempnr,i32,i64,f32;noparamssingleresult,v128;multiparamssingleresult,v128,i32,v128" -S -o - | filecheck %s
+;; RUN: wasm-opt --enable-simd --trace-calls="noparamsnoresults,singleparamnoresults,multiparamsnoresults:tracempnr,noparamssingleresult,multiparamssingleresult" %s -S -o - | filecheck %s
 
 (module
   ;; CHECK:      (type $0 (func))
@@ -39,16 +39,16 @@
   (import "env" "dont_trace_me" (func $donttraceme))
 
 
-  ;; CHECK:      (import "env" "tracempnr" (func $tracempnr (param i32 i64 f32)))
-
   ;; CHECK:      (import "env" "trace_multiparamssingleresult" (func $trace_multiparamssingleresult (param v128 i32 v128)))
 
-  ;; CHECK:      (import "env" "trace_noparamsnoresults" (func $trace_noparamsnoresults))
-
   ;; CHECK:      (import "env" "trace_noparamssingleresult" (func $trace_noparamssingleresult (param v128)))
 
+  ;; CHECK:      (import "env" "tracempnr" (func $tracempnr (param i32 i64 f32)))
+
   ;; CHECK:      (import "env" "trace_singleparamnoresults" (func $trace_singleparamnoresults (param f64)))
 
+  ;; CHECK:      (import "env" "trace_noparamsnoresults" (func $trace_noparamsnoresults))
+
   ;; CHECK:      (func $test_no_params_no_results
   ;; CHECK-NEXT:  (call $noparamsnoresults)
   ;; CHECK-NEXT:  (call $trace_noparamsnoresults)

test/lit/passes/trace-calls_multi-value-result.wast

@@ -0,0 +1,9 @@
+;; Test that a traced function with a multi-value result type
+;; results in a useful error message
+
+;; RUN: not wasm-opt --enable-simd --enable-multivalue --trace-calls=multi_param_result %s 2>&1 | filecheck %s
+
+;; CHECK: Fatal: Failed to instrument function 'multi_param_result': Multi-value result type is not supported
+(module
+  (import "env" "multi_param_result" (func $multi_param_result (result i32 i32)))
+)