Thunks: Carry annotations in callback wrappers of host functions

Previously, two functions with the same signature would always be wrapped
in the same logic. This change allows customizing one function with
annotations while leaving the other one unchanged.

ThunkLibs/Generator/analysis.cpp

@@ -170,6 +170,8 @@ static ParameterAnnotations GetParameterAnnotations(clang::ASTContext& context,
 void AnalysisAction::ParseInterface(clang::ASTContext& context) {
     ErrorReporter report_error { context };
 
+    const std::unordered_map<unsigned, ParameterAnnotations> no_param_annotations {};
+
     // TODO: Assert fex_gen_type is not declared at non-global namespaces
     if (auto template_decl = FindClassTemplateDeclByName(*context.getTranslationUnitDecl(), "fex_gen_type")) {
         for (auto* decl : template_decl->specializations()) {
@@ -184,7 +186,7 @@ void AnalysisAction::ParseInterface(clang::ASTContext& context) {
             type = type->getLocallyUnqualifiedSingleStepDesugaredType();
 
             if (type->isFunctionPointerType() || type->isFunctionType()) {
-                thunked_funcptrs.insert(type.getTypePtr());
+                thunked_funcptrs[type.getAsString()] = std::pair { type.getTypePtr(), no_param_annotations };
             } else {
                 [[maybe_unused]] auto [it, inserted] = types.emplace(context.getCanonicalType(type.getTypePtr()), RepackedType { });
                 assert(inserted);
@@ -349,7 +351,7 @@ void AnalysisAction::ParseInterface(clang::ASTContext& context) {
 
                             data.callbacks.emplace(param_idx, callback);
                             if (!callback.is_stub && !callback.is_guest && !data.custom_host_impl) {
-                                thunked_funcptrs.insert(context.getCanonicalType(funcptr));
+                                thunked_funcptrs[emitted_function->getNameAsString() + "_cb" + std::to_string(param_idx)] = std::pair { context.getCanonicalType(funcptr), no_param_annotations };
                             }
 
                             if (data.callbacks.size() != 1) {
@@ -417,7 +419,7 @@ void AnalysisAction::ParseInterface(clang::ASTContext& context) {
 
                     // For indirect calls, register the function signature as a function pointer type
                     if (namespace_info.indirect_guest_calls) {
-                        thunked_funcptrs.insert(context.getCanonicalType(emitted_function->getFunctionType()));
+                        thunked_funcptrs[emitted_function->getNameAsString()] = std::pair { context.getCanonicalType(emitted_function->getFunctionType()), data.param_annotations };
                     }
 
                     thunks.push_back(std::move(data));

ThunkLibs/Generator/analysis.h

@@ -133,7 +133,9 @@ protected:
     std::vector<ThunkedAPIFunction> thunked_api;
 
     // Set of function types for which to generate Guest->Host thunking trampolines.
-    std::unordered_set<const clang::Type*> thunked_funcptrs;
+    // The map key is a unique identifier that must be consistent between guest/host processing passes.
+    // The map value is a pair of the function pointer's clang::Type and the mapping of parameter annotations
+    std::unordered_map<std::string, std::pair<const clang::Type*, std::unordered_map<unsigned, ParameterAnnotations>>> thunked_funcptrs;
 
     std::unordered_map<const clang::Type*, RepackedType> types;
     std::optional<unsigned> lib_version;

ThunkLibs/Generator/data_layout.cpp

@@ -151,10 +151,34 @@ ABI GetStableLayout(const clang::ASTContext& context, const std::unordered_map<c
     return stable_layout;
 }
 
+static std::array<uint8_t, 32> GetSha256(const std::string& function_name) {
+    std::array<uint8_t, 32> sha256;
+    SHA256(reinterpret_cast<const unsigned char*>(function_name.data()),
+           function_name.size(),
+           sha256.data());
+    return sha256;
+};
+
 void AnalyzeDataLayoutAction::OnAnalysisComplete(clang::ASTContext& context) {
     if (StrictModeEnabled(context)) {
     type_abi = GetStableLayout(context, ComputeDataLayout(context, types));
     }
+
+    // Register functions that must be guest-callable through host function pointers
+    for (auto funcptr_type_it = thunked_funcptrs.begin(); funcptr_type_it != thunked_funcptrs.end(); ++funcptr_type_it) {
+        auto& funcptr_id = funcptr_type_it->first;
+        auto& [type, param_annotations] = funcptr_type_it->second;
+        auto func_type = type->getAs<clang::FunctionProtoType>();
+        std::string mangled_name = clang::QualType { type, 0 }.getAsString();
+        auto cb_sha256 = GetSha256("fexcallback_" + mangled_name);
+        FuncPtrInfo info = { cb_sha256 };
+
+        info.result = func_type->getReturnType().getAsString();
+        for (auto arg : func_type->getParamTypes()) {
+            info.args.push_back(arg.getAsString());
+        }
+        type_abi.thunked_funcptrs[funcptr_id] = std::move(info);
+    }
 }
 
 TypeCompatibility DataLayoutCompareAction::GetTypeCompatibility(
@@ -297,6 +321,10 @@ TypeCompatibility DataLayoutCompareAction::GetTypeCompatibility(
     return compat;
 }
 
+FuncPtrInfo DataLayoutCompareAction::LookupGuestFuncPtrInfo(const char* funcptr_id) {
+    return guest_abi.thunked_funcptrs.at(funcptr_id);
+}
+
 DataLayoutCompareActionFactory::DataLayoutCompareActionFactory(const ABI& abi) : abi(abi) {
 
 }

ThunkLibs/Generator/data_layout.h

@@ -84,6 +84,7 @@ struct FuncPtrInfo {
 };
 
 struct ABI : std::unordered_map<std::string, TypeInfo> {
+    std::unordered_map<std::string, FuncPtrInfo> thunked_funcptrs;
     int pointer_size; // in bytes
 };
 
@@ -111,6 +112,8 @@ public:
             const std::unordered_map<const clang::Type*, TypeInfo> host_abi,
             std::unordered_map<const clang::Type*, TypeCompatibility>& type_compat);
 
+    FuncPtrInfo LookupGuestFuncPtrInfo(const char* funcptr_id);
+
 private:
     const ABI& guest_abi;
 };

ThunkLibs/Generator/gen.cpp

@@ -120,8 +120,8 @@ void GenerateThunkLibsAction::OnAnalysisComplete(clang::ASTContext& context) {
         return ret;
     };
 
-    auto get_sha256 = [this](const std::string& function_name) {
-        std::string sha256_message = libname + ":" + function_name;
+    auto get_sha256 = [this](const std::string& function_name, bool include_libname) {
+        std::string sha256_message = (include_libname ? libname + ":" : "") + function_name;
         std::vector<unsigned char> sha256(SHA256_DIGEST_LENGTH);
         SHA256(reinterpret_cast<const unsigned char*>(sha256_message.data()),
                sha256_message.size(),
@@ -141,18 +141,26 @@ void GenerateThunkLibsAction::OnAnalysisComplete(clang::ASTContext& context) {
         file << "extern \"C\" {\n";
         for (auto& thunk : thunks) {
             const auto& function_name = thunk.function_name;
-            auto sha256 = get_sha256(function_name);
+            auto sha256 = get_sha256(function_name, true);
             fmt::print( file, "MAKE_THUNK({}, {}, \"{:#02x}\")\n",
                         libname, function_name, fmt::join(sha256, ", "));
         }
         file << "}\n";
 
         // Guest->Host transition points for invoking runtime host-function pointers based on their signature
+        std::vector<std::vector<unsigned char>> sha256s;
         for (auto type_it = thunked_funcptrs.begin(); type_it != thunked_funcptrs.end(); ++type_it) {
-            auto* type = *type_it;
+            auto* type = type_it->second.first;
             std::string funcptr_signature = clang::QualType { type, 0 }.getAsString();
 
-            auto cb_sha256 = get_sha256("fexcallback_" + funcptr_signature);
+            auto cb_sha256 = get_sha256("fexcallback_" + funcptr_signature, false);
+            auto it = std::find(sha256s.begin(), sha256s.end(), cb_sha256);
+            if (it != sha256s.end()) {
+                // TODO: Avoid this ugly way of avoiding duplicates
+                continue;
+            } else {
+                sha256s.push_back(cb_sha256);
+            }
 
             // Thunk used for guest-side calls to host function pointers
             file << "  // " << funcptr_signature << "\n";
@@ -382,18 +390,30 @@ void GenerateThunkLibsAction::OnAnalysisComplete(clang::ASTContext& context) {
         file << "static ExportEntry exports[] = {\n";
         for (auto& thunk : thunks) {
             const auto& function_name = thunk.function_name;
-            auto sha256 = get_sha256(function_name);
+            auto sha256 = get_sha256(function_name, true);
             fmt::print( file, "  {{(uint8_t*)\"\\x{:02x}\", (void(*)(void *))&fexfn_unpack_{}_{}}}, // {}:{}\n",
                         fmt::join(sha256, "\\x"), libname, function_name, libname, function_name);
         }
 
         // Endpoints for Guest->Host invocation of runtime host-function pointers
-        for (auto& type : thunked_funcptrs) {
+        for (auto& host_funcptr_entry : thunked_funcptrs) {
+            auto& [type, param_annotations] = host_funcptr_entry.second;
             std::string mangled_name = clang::QualType { type, 0 }.getAsString();
-            auto cb_sha256 = get_sha256("fexcallback_" + mangled_name);
-            fmt::print( file, "  {{(uint8_t*)\"\\x{:02x}\", (void(*)(void *))&CallbackUnpack<{}>::ForIndirectCall}},\n",
-                        fmt::join(cb_sha256, "\\x"), mangled_name);
+            auto info = LookupGuestFuncPtrInfo(host_funcptr_entry.first.c_str());
+
+            std::string annotations;
+            for (int param_idx = 0; param_idx < info.args.size(); ++param_idx) {
+                if (param_idx != 0) {
+                    annotations += ", ";
+                }
+
+                // TODO: Add annotations as needed
+                annotations += "ParameterAnnotations {}";
+            }
+            fmt::print( file, "  {{(uint8_t*)\"\\x{:02x}\", (void(*)(void *))&GuestWrapperForHostFunction<{}({})>::Call<{}>}}, // {}\n",
+                        fmt::join(info.sha256, "\\x"), info.result, fmt::join(info.args, ", "), annotations, host_funcptr_entry.first);
         }
+
         file << "  { nullptr, nullptr }\n";
         file << "};\n";
 

ThunkLibs/include/common/Host.h

@@ -103,6 +103,9 @@ struct GuestcallInfo {
   asm volatile("mov %0, x11" : "=r" (target_variable))
 #endif
 
+struct ParameterAnnotations {
+};
+
 template<typename>
 struct CallbackUnpack;
 
@@ -120,8 +123,23 @@ struct CallbackUnpack<Result(Args...)> {
       return packed_args.rv;
     }
   }
+};
+
+template<ParameterAnnotations Annotation, typename T>
+auto Projection(T& data) {
+  return data;
+}
+
+template<typename>
+struct GuestWrapperForHostFunction;
+
+template<typename Result, typename... Args>
+struct GuestWrapperForHostFunction<Result(Args...)> {
+  // Host functions called from Guest
+  template<ParameterAnnotations... Annotations>
+  static void Call(void* argsv) {
+    static_assert(sizeof...(Annotations) == sizeof...(Args));
 
-  static void ForIndirectCall(void* argsv) {
     auto args = reinterpret_cast<PackedArguments<Result, Args..., uintptr_t>*>(argsv);
     constexpr auto CBIndex = sizeof...(Args);
     uintptr_t cb;
@@ -168,8 +186,14 @@ struct CallbackUnpack<Result(Args...)> {
       cb = args->a23;
     }
 
+    // This is almost the same type as "Result func(Args..., uintptr_t)", but
+    // individual parameters annotated as passthrough are replaced by guest_layout<GuestArgs>
     auto callback = reinterpret_cast<Result(*)(Args..., uintptr_t)>(cb);
-    Invoke(callback, *args);
+
+    auto f = [&callback](Args... args, uintptr_t target) -> Result {
+      return callback(Projection<Annotations, Args>(args)..., target);
+    };
+    Invoke(f, *args);
   }
 };
 

ThunkLibs/include/common/PackedArguments.h

@@ -113,8 +113,8 @@ T&& operator,(T&& t, Regularize) {
     return std::forward<T>(t);
 }
 
-template<typename Result, typename... Args>
-void Invoke(Result(*func)(Args...), PackedArguments<Result, Args...>& args) {
+template<typename Result, typename... Args, typename Func>
+void Invoke(Func&& func, PackedArguments<Result, Args...>& args) requires(std::is_invocable_r_v<Result, Func, Args...>) {
     constexpr auto NumArgs = sizeof...(Args);
     static_assert(NumArgs <= 19 || NumArgs == 24);
 

unittests/ThunkLibs/generator.cpp

@@ -270,9 +270,10 @@ SourceWithAST Fixture::run_thunkgen_host(std::string_view prelude, std::string_v
         "  uintptr_t GuestUnpacker;\n"
         "  uintptr_t GuestTarget;\n"
         "};\n"
+        "struct ParameterAnnotations {};\n"
         "template<typename>\n"
-        "struct CallbackUnpack {\n"
-        "  static void ForIndirectCall(void* argsv);\n"
+        "struct GuestWrapperForHostFunction {\n"
+        "  template<ParameterAnnotations...> static void Call(void*);\n"
         "};\n"
         "template<typename F>\n"
         "void FinalizeHostTrampolineForGuestFunction(F*);\n"
@@ -377,7 +378,7 @@ TEST_CASE_METHOD(Fixture, "FunctionPointerViaType") {
         matches(varDecl(
             hasName("exports"),
             hasType(constantArrayType(hasElementType(asString("struct ExportEntry")), hasSize(2))),
-            hasInitializer(hasDescendant(declRefExpr(to(cxxMethodDecl(hasName("ForIndirectCall"), ofClass(hasName("CallbackUnpack"))).bind("funcptr")))))
+            hasInitializer(hasDescendant(declRefExpr(to(cxxMethodDecl(hasName("Call"), ofClass(hasName("GuestWrapperForHostFunction"))).bind("funcptr")))))
             )).check_binding("funcptr", +[](const clang::CXXMethodDecl* decl) {
                 auto parent = llvm::cast<clang::ClassTemplateSpecializationDecl>(decl->getParent());
                 return parent->getTemplateArgs().get(0).getAsType().getAsString() == "int (char, char)";
@@ -415,7 +416,7 @@ TEST_CASE_METHOD(Fixture, "FunctionPointerParameter") {
         matches(varDecl(
             hasName("exports"),
             hasType(constantArrayType(hasElementType(asString("struct ExportEntry")), hasSize(3))),
-            hasInitializer(hasDescendant(declRefExpr(to(cxxMethodDecl(hasName("ForIndirectCall"), ofClass(hasName("CallbackUnpack")))))))
+            hasInitializer(hasDescendant(declRefExpr(to(cxxMethodDecl(hasName("Call"), ofClass(hasName("GuestWrapperForHostFunction")))))))
             )));
 }