From e845217706749e61fa6fa6fceff5fdccfaa73c0f Mon Sep 17 00:00:00 2001
From: Congcong Cai <congcong.cai@bmw.com>
Date: Tue, 14 Nov 2023 12:33:37 +0800
Subject: [PATCH] fix: handle nest tostack in shadowstack pass

---
 src/passes/shadowstack.ts                | 15 +++++++++------
 tests/compiler/resolve-unary.debug.wat   |  8 ++------
 tests/compiler/resolve-unary.release.wat |  8 --------
 3 files changed, 11 insertions(+), 20 deletions(-)

diff --git a/src/passes/shadowstack.ts b/src/passes/shadowstack.ts
index 4d08e39971..d34997ed08 100644
--- a/src/passes/shadowstack.ts
+++ b/src/passes/shadowstack.ts
@@ -159,18 +159,21 @@ type TempMap = Map<TypeRef,LocalIndex>;
 
 /** Attempts to match the `__tostack(value)` pattern. Returns `value` if a match, otherwise `0`.  */
 function matchPattern(module: Module, expr: ExpressionRef): ExpressionRef {
-  if (
+  let isFound = false;
+  while (
     _BinaryenExpressionGetId(expr) == ExpressionId.Call &&
     module.readStringCached(_BinaryenCallGetTarget(expr)) == BuiltinNames.tostack
   ) {
     assert(_BinaryenCallGetNumOperands(expr) == 1);
-    return _BinaryenCallGetOperandAt(expr, 0);
+    expr = _BinaryenCallGetOperandAt(expr, 0);
+    isFound = true;
   }
-  return 0;
+  if (!isFound) return 0;
+  return expr;
 }
 
 /** Tests whether a `value` matched by `matchTostack` needs a slot. */
-function needsSlot(module: Module, value: ExpressionRef): bool {
+function needsSlot(value: ExpressionRef): bool {
   switch (_BinaryenExpressionGetId(value)) {
     // no need to stack null pointers
     case ExpressionId.Const: return !isConstZero(value);
@@ -344,7 +347,7 @@ export class ShadowStackPass extends Pass {
       let operand = operands[i];
       let match = matchPattern(module, operand);
       if (!match) continue;
-      if (!needsSlot(module, match)) {
+      if (!needsSlot(match)) {
         operands[i] = match;
         continue;
       }
@@ -434,7 +437,7 @@ export class ShadowStackPass extends Pass {
     let value = _BinaryenLocalSetGetValue(localSet);
     let match = matchPattern(module, value);
     if (!match) return;
-    if (!needsSlot(module, match)) {
+    if (!needsSlot(match)) {
       _BinaryenLocalSetSetValue(localSet, match);
       return;
     }
diff --git a/tests/compiler/resolve-unary.debug.wat b/tests/compiler/resolve-unary.debug.wat
index 1eff1a8a06..f22d05a02b 100644
--- a/tests/compiler/resolve-unary.debug.wat
+++ b/tests/compiler/resolve-unary.debug.wat
@@ -3753,10 +3753,8 @@
    call $~lib/builtins/abort
    unreachable
   end
-  global.get $~lib/memory/__stack_pointer
   global.get $resolve-unary/bar
-  local.tee $6
-  i32.store
+  local.set $6
   global.get $~lib/memory/__stack_pointer
   local.get $6
   i32.store offset=12
@@ -3789,10 +3787,8 @@
    call $~lib/builtins/abort
    unreachable
   end
-  global.get $~lib/memory/__stack_pointer
   global.get $resolve-unary/bar
-  local.tee $6
-  i32.store
+  local.set $6
   global.get $~lib/memory/__stack_pointer
   local.get $6
   i32.store offset=12
diff --git a/tests/compiler/resolve-unary.release.wat b/tests/compiler/resolve-unary.release.wat
index 97dd0c9459..1998ea2c48 100644
--- a/tests/compiler/resolve-unary.release.wat
+++ b/tests/compiler/resolve-unary.release.wat
@@ -2626,10 +2626,6 @@
    global.get $~lib/memory/__stack_pointer
    local.tee $0
    global.get $resolve-unary/bar
-   local.tee $1
-   i32.store
-   local.get $0
-   local.get $1
    i32.store offset=12
    local.get $0
    i32.const 3680
@@ -2655,10 +2651,6 @@
    global.get $~lib/memory/__stack_pointer
    local.tee $0
    global.get $resolve-unary/bar
-   local.tee $1
-   i32.store
-   local.get $0
-   local.get $1
    i32.store offset=12
    local.get $0
    i32.const 3712