From 6a68130d5b0296379fae0b8de5fbb8a1499b67a5 Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <public@benj.me>
Date: Wed, 1 Apr 2020 16:05:34 +0200
Subject: [PATCH] cranelift codegen: add a supplementary method add_call_site
 to CodeSink;

This allows keeping track of indirect call sites, for instance.
---
 cranelift/codegen/meta/src/isa/x86/recipes.rs |  3 +++
 cranelift/codegen/src/binemit/memorysink.rs   | 15 ++++++++++++++-
 cranelift/codegen/src/binemit/mod.rs          |  9 ++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/cranelift/codegen/meta/src/isa/x86/recipes.rs b/cranelift/codegen/meta/src/isa/x86/recipes.rs
index d4fea0a9b1..c42979f5ac 100644
--- a/cranelift/codegen/meta/src/isa/x86/recipes.rs
+++ b/cranelift/codegen/meta/src/isa/x86/recipes.rs
@@ -2417,6 +2417,7 @@ pub(crate) fn define<'shared>(
                                 &func.dfg.ext_funcs[func_ref].name,
                                 -4);
             sink.put4(0);
+            sink.add_call_site(opcode, func.srclocs[inst]);
         "#,
         ),
     );
@@ -2431,6 +2432,7 @@ pub(crate) fn define<'shared>(
                                 &func.dfg.ext_funcs[func_ref].name,
                                 -4);
             sink.put4(0);
+            sink.add_call_site(opcode, func.srclocs[inst]);
         "#,
         ),
     );
@@ -2443,6 +2445,7 @@ pub(crate) fn define<'shared>(
                     sink.trap(TrapCode::StackOverflow, func.srclocs[inst]);
                     {{PUT_OP}}(bits, rex1(in_reg0), sink);
                     modrm_r_bits(in_reg0, bits, sink);
+                    sink.add_call_site(opcode, func.srclocs[inst]);
                 "#,
             ),
     );
diff --git a/cranelift/codegen/src/binemit/memorysink.rs b/cranelift/codegen/src/binemit/memorysink.rs
index 45b73131f8..110017bf55 100644
--- a/cranelift/codegen/src/binemit/memorysink.rs
+++ b/cranelift/codegen/src/binemit/memorysink.rs
@@ -16,7 +16,7 @@
 use super::{Addend, CodeInfo, CodeOffset, CodeSink, Reloc};
 use crate::binemit::stackmap::Stackmap;
 use crate::ir::entities::Value;
-use crate::ir::{ConstantOffset, ExternalName, Function, JumpTable, SourceLoc, TrapCode};
+use crate::ir::{ConstantOffset, ExternalName, Function, JumpTable, Opcode, SourceLoc, TrapCode};
 use crate::isa::TargetIsa;
 use core::ptr::write_unaligned;
 
@@ -92,6 +92,10 @@ pub trait RelocSink {
 
     /// Add a relocation referencing a jump table.
     fn reloc_jt(&mut self, _: CodeOffset, _: Reloc, _: JumpTable);
+
+    /// Track a call site whose return address is the given CodeOffset, for the given opcode. Does
+    /// nothing in general, only useful for certain embedders (SpiderMonkey).
+    fn add_call_site(&mut self, _: Opcode, _: CodeOffset, _: SourceLoc) {}
 }
 
 /// A trait for receiving trap codes and offsets.
@@ -183,6 +187,15 @@ impl<'a> CodeSink for MemoryCodeSink<'a> {
         let stackmap = Stackmap::from_values(&val_list, func, isa);
         self.stackmaps.add_stackmap(ofs, stackmap);
     }
+
+    fn add_call_site(&mut self, opcode: Opcode, loc: SourceLoc) {
+        debug_assert!(
+            opcode.is_call(),
+            "adding call site info for a non-call instruction."
+        );
+        let ret_addr = self.offset();
+        self.relocs.add_call_site(opcode, ret_addr, loc);
+    }
 }
 
 /// A `RelocSink` implementation that does nothing, which is convenient when
diff --git a/cranelift/codegen/src/binemit/mod.rs b/cranelift/codegen/src/binemit/mod.rs
index 872f33ab28..3a33649d4d 100644
--- a/cranelift/codegen/src/binemit/mod.rs
+++ b/cranelift/codegen/src/binemit/mod.rs
@@ -16,7 +16,9 @@ pub use self::relaxation::relax_branches;
 pub use self::shrink::shrink_instructions;
 pub use self::stackmap::Stackmap;
 use crate::ir::entities::Value;
-use crate::ir::{ConstantOffset, ExternalName, Function, Inst, JumpTable, SourceLoc, TrapCode};
+use crate::ir::{
+    ConstantOffset, ExternalName, Function, Inst, JumpTable, Opcode, SourceLoc, TrapCode,
+};
 use crate::isa::TargetIsa;
 pub use crate::regalloc::RegDiversions;
 use core::fmt;
@@ -162,6 +164,11 @@ pub trait CodeSink {
 
     /// Add a stackmap at the current code offset.
     fn add_stackmap(&mut self, _: &[Value], _: &Function, _: &dyn TargetIsa);
+
+    /// Add a call site for a call with the given opcode, returning at the current offset.
+    fn add_call_site(&mut self, _: Opcode, _: SourceLoc) {
+        // Default implementation doesn't need to do anything.
+    }
 }
 
 /// Type of the frame unwind information.