From 36e3059be7201e143f38a18c43fa65d9baf6b578 Mon Sep 17 00:00:00 2001
From: Graham Kelly <gkgoat6700@gmail.com>
Date: Sun, 2 Jun 2024 19:33:57 -0400
Subject: [PATCH] return call ref

---
 src/backend/mod.rs      |  6 ++++++
 src/backend/stackify.rs | 15 +++++++++++++++
 src/frontend.rs         | 17 ++++++++++++++++-
 src/interp.rs           |  1 +
 src/ir/func.rs          | 26 ++++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/backend/mod.rs b/src/backend/mod.rs
index 9b6bcaa..e9a9160 100644
--- a/src/backend/mod.rs
+++ b/src/backend/mod.rs
@@ -214,6 +214,12 @@ impl<'a> WasmFuncBackend<'a> {
                         table: table.index() as u32,
                     });
                 }
+                WasmBlock::ReturnCallRef { sig, values } => {
+                    for &value in &values[..] {
+                        self.lower_value(value, func);
+                    }
+                    func.instruction(&wasm_encoder::Instruction::ReturnCallRef(sig.index() as u32));
+                }
                 WasmBlock::Unreachable => {
                     func.instruction(&wasm_encoder::Instruction::Unreachable);
                 }
diff --git a/src/backend/stackify.rs b/src/backend/stackify.rs
index c61e985..0a37b62 100644
--- a/src/backend/stackify.rs
+++ b/src/backend/stackify.rs
@@ -61,6 +61,12 @@ pub enum WasmBlock<'a> {
         table: Table,
         values: &'a [Value],
     },
+    /// A function refernce tail call instruction
+    ReturnCallRef {
+        sig: Signature,
+        // table: Table,
+        values: &'a [Value],
+    },
     /// An unreachable instruction.
     Unreachable,
 }
@@ -465,6 +471,15 @@ impl<'a, 'b> Context<'a, 'b> {
                     table,
                     values: args,
                 }),
+                &Terminator::ReturnCallRef {
+                    sig,
+                    // table,
+                    ref args,
+                } => into.push(WasmBlock::ReturnCallRef {
+                    sig,
+                    // table,
+                    values: args,
+                }),
             }
         }
     }
diff --git a/src/frontend.rs b/src/frontend.rs
index 2bde9d7..40c82c5 100644
--- a/src/frontend.rs
+++ b/src/frontend.rs
@@ -1531,7 +1531,22 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
                 });
                 self.reachable = false;
             }
-
+            wasmparser::Operator::ReturnCallRef {
+                type_index,
+                // table_index,
+            } => {
+                let retvals = self.pop_n(
+                    self.module.signatures[Signature::new(*type_index as usize)]
+                        .params
+                        .len(),
+                );
+                self.emit_term(Terminator::ReturnCallRef {
+                    sig: Signature::new(*type_index as usize),
+                    // table: Table::new(*table_index as usize),
+                    args: retvals,
+                });
+                self.reachable = false;
+            }
             _ => bail!(FrontendError::UnsupportedFeature(format!(
                 "Unsupported operator: {:?}",
                 op
diff --git a/src/interp.rs b/src/interp.rs
index b4b0f49..a4b84f0 100644
--- a/src/interp.rs
+++ b/src/interp.rs
@@ -325,6 +325,7 @@ impl InterpContext {
                         log::trace!("returning from {}: {:?}", func, values);
                         return InterpResult::Ok(values);
                     }
+                    Terminator::ReturnCallRef { sig, args } => todo!(),
                 }
             }
         }
diff --git a/src/ir/func.rs b/src/ir/func.rs
index 7d63fcf..f160ece 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -545,6 +545,10 @@ pub enum Terminator {
         table: Table,
         args: Vec<Value>,
     },
+    ReturnCallRef {
+        sig: Signature,
+        args: Vec<Value>,
+    },
     Unreachable,
     None,
 }
@@ -609,6 +613,16 @@ impl std::fmt::Display for Terminator {
                     .collect::<Vec<_>>()
                     .join(", ")
             )?,
+            Terminator::ReturnCallRef { sig, args } =>write!(
+                f,
+                "return_call_ref ({})({})",
+                sig,
+                // table,
+                args.iter()
+                    .map(|val| format!("{}", val))
+                    .collect::<Vec<_>>()
+                    .join(", ")
+            )?,
         }
         Ok(())
     }
@@ -641,6 +655,7 @@ impl Terminator {
             Terminator::Unreachable => {}
             Terminator::ReturnCall { func, args } => {}
             Terminator::ReturnCallIndirect { sig, table, args } => {}
+            Terminator::ReturnCallRef { sig, args } => {},
         }
     }
 
@@ -670,6 +685,9 @@ impl Terminator {
             Terminator::Unreachable => {}
             Terminator::ReturnCall { func, args } => {}
             Terminator::ReturnCallIndirect { sig, table, args } => {}
+            Terminator::ReturnCallRef { sig, args } => {
+                
+            },
         }
     }
 
@@ -787,6 +805,14 @@ impl Terminator {
                     f(value);
                 }
             }
+            &mut Terminator::ReturnCallRef {
+                sig,
+                ref mut args,
+            } => {
+                for value in args {
+                    f(value);
+                }
+            }
             _ => {}
         }
     }