From 573af9c19b5e8156ca27413d6f3ed6ef462a2031 Mon Sep 17 00:00:00 2001
From: Andy Wingo <wingo@igalia.com>
Date: Wed, 3 Apr 2019 12:09:38 +0200
Subject: [PATCH] Clean up stack after calls

---
 jit.h     |  2 +-
 jit/x86.c | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/jit.h b/jit.h
index ba8079d9f..67e107497 100644
--- a/jit.h
+++ b/jit.h
@@ -170,7 +170,7 @@ JIT_API void jit_patch_here(jit_state_t*, jit_reloc_t);
 JIT_API void jit_patch_there(jit_state_t*, jit_reloc_t, jit_pointer_t);
 
 /* Note that all functions that take jit_arg_t args[] use the args as scratch
-   space.  */
+   space while shuffling values into position.  */
 JIT_API void jit_calli(jit_state_t *, jit_pointer_t f,
                        size_t argc, const jit_arg_abi_t abi[],
                        jit_arg_t args[]);
diff --git a/jit/x86.c b/jit/x86.c
index d965afe9c..16bb5187a 100644
--- a/jit/x86.c
+++ b/jit/x86.c
@@ -735,6 +735,22 @@ prepare_args(jit_state_t *_jit, size_t argc, const jit_arg_abi_t abi[],
     }
 }
 
+static void
+cleanup_stack_after_call(jit_state_t *_jit, size_t argc,
+                         const jit_arg_abi_t abi[])
+{
+  jit_arg_t scratch;
+  struct abi_arg_iterator iter;
+
+  // Compute stack arg size.
+  reset_abi_arg_iterator(&iter, argc, abi);
+  for (size_t i = 0; i < argc; i++)
+    next_abi_arg(&iter, &scratch);
+
+  if (iter.stack_size)
+    jit_addi(_jit, JIT_SP, JIT_SP, iter.stack_size);
+}
+
 void
 jit_calli(jit_state_t *_jit, jit_pointer_t f,
           size_t argc, const jit_arg_abi_t abi[], jit_arg_t args[])
@@ -742,6 +758,8 @@ jit_calli(jit_state_t *_jit, jit_pointer_t f,
   prepare_args(_jit, argc, abi, args);
 
   calli(_jit, (jit_word_t)f);
+
+  cleanup_stack_after_call(_jit, argc, abi);
 }
 
 void
@@ -751,6 +769,8 @@ jit_callr(jit_state_t *_jit, jit_gpr_t f,
   prepare_args(_jit, argc, abi, args);
 
   callr(_jit, rn(f));
+
+  cleanup_stack_after_call(_jit, argc, abi);
 }
 
 void