rewriter: distinguish address-taken call gates

This commit distinguishes between call gates for static address-taken,
non-static address-taken in the same and different DSOs.

tools/rewriter/SourceRewriter.cpp

@@ -102,7 +102,7 @@ static std::string RootDirectory;
 static std::string OutputDirectory;
 static std::string OutputPrefix;
 
-// Map each translation unit's filename to its pkey.
+// Map each translation unit's original filename to its pkey.
 static std::map<Filename, Pkey> file_pkeys;
 
 static std::map<Filename, Filename> rel_path_to_full;
@@ -662,7 +662,7 @@ class FnPtrExpr : public RefactoringCallback {
 
     auto linkage = fn_decl->getFormalLinkage();
     auto [it, new_fn] = addr_taken_fns[filename].insert(
-        std::make_pair(fn_name, new_type));
+        std::make_tuple(fn_name, new_type, linkage));
 
     // TODO: Note that this only checks if a function is added to the
     // addr_taken_fns map. To make the rewriter idempotent we should
@@ -710,7 +710,7 @@ class FnPtrExpr : public RefactoringCallback {
     return;
   }
 
-  std::map<Filename, std::set<std::pair<Function, OpaqueStruct>>>
+  std::map<Filename, std::set<std::tuple<Function, OpaqueStruct, clang::Linkage>>>
       addr_taken_fns;
 
 private:
@@ -1422,8 +1422,20 @@ int main(int argc, const char **argv) {
 
     // For each static function, define a macro to define the wrapper in the
     // output header and invoke it in the source file
-    for (const auto [fn_name, opaque] : addr_taken_fns_in_file) {
+    for (const auto [fn_name, opaque, linkage] : addr_taken_fns_in_file) {
       std::string wrapper_name = "__ia2_"s + fn_name;
+      std::string target_fn = fn_name;
+
+      if (clang::isExternallyVisible(linkage)) {// && filename pkey does not fn_name pkey) {
+        llvm::SmallString<256> input_file(filename);
+        llvm::sys::path::replace_path_prefix(input_file, OutputDirectory, RootDirectory);
+        Pkey file_pkey = file_pkeys.at(input_file.str().str());
+        Pkey fn_pkey = fn_decl_pass.fn_pkeys.at(fn_name);
+        if (file_pkey != fn_pkey) {
+          target_fn = "__real_"s + fn_name;
+        }
+      }
+
       // TODO: These wrapper go from pkey 0 to the target pkey so if the target
       // also has pkey 0 then it just needs call the original function
       Pkey target_pkey = fn_decl_pass.fn_pkeys[fn_name];
@@ -1432,14 +1444,14 @@ int main(int argc, const char **argv) {
         AbiSignature c_abi_sig = fn_decl_pass.abi_signatures[fn_name];
 
         std::string asm_wrapper = emit_asm_wrapper(
-            c_abi_sig, wrapper_name, fn_name, WrapperKind::Pointer, 0,
+            c_abi_sig, wrapper_name, target_fn, WrapperKind::Pointer, 0,
             target_pkey, Target, true /* as_macro */);
-        static_wrappers += "#define IA2_DEFINE_WRAPPER_"s + fn_name + " \\\n";
+        static_wrappers += "#define IA2_DEFINE_WRAPPER_"s + target_fn + " \\\n";
         static_wrappers += "asm(\\\n";
         static_wrappers += asm_wrapper;
         static_wrappers += ");\n";
 
-        source_file << "IA2_DEFINE_WRAPPER(" << fn_name << ")\n";
+        source_file << "IA2_DEFINE_WRAPPER(" << target_fn << ")\n";
       } else {
         header_out << "asm(\n";
         header_out << "  \".set " << wrapper_name << ", __real_" << fn_name << "\\n\"\n";