Update StackCheck for memory64 (#4636)

3 files changed, 105 insertions, 8 deletions

diff --git a/src/passes/StackCheck.cpp b/src/passes/StackCheck.cpp
index 07cf46aee..4b1054597 100644
--- a/src/passes/StackCheck.cpp
+++ b/src/passes/StackCheck.cpp
@@ -21,6 +21,7 @@
 //
 
 #include "abi/js.h"
+#include "ir/abstract.h"
 #include "ir/import-utils.h"
 #include "ir/names.h"
 #include "pass.h"
@@ -87,16 +88,17 @@ struct EnforceStackLimits : public WalkerPass<PostWalker<EnforceStackLimits>> {
     } else {
       handlerExpr = builder.makeUnreachable();
     }
+
     // If it is >= the base or <= the limit, then error.
     auto check = builder.makeIf(
       builder.makeBinary(
         BinaryOp::OrInt32,
         builder.makeBinary(
-          BinaryOp::GtUInt32,
+          Abstract::getBinary(stackPointer->type, Abstract::GtU),
           builder.makeLocalTee(newSP, value, stackPointer->type),
           builder.makeGlobalGet(stackBase->name, stackBase->type)),
         builder.makeBinary(
-          BinaryOp::LtUInt32,
+          Abstract::getBinary(stackPointer->type, Abstract::LtU),
           builder.makeLocalGet(newSP, stackPointer->type),
           builder.makeGlobalGet(stackLimit->name, stackLimit->type))),
       handlerExpr);
@@ -147,12 +149,12 @@ struct StackCheck : public Pass {
     auto stackBase =
       module->addGlobal(builder.makeGlobal(stackBaseName,
                                            stackPointer->type,
-                                           builder.makeConst(int32_t(0)),
+                                           builder.makeConstPtr(0),
                                            Builder::Mutable));
     auto stackLimit =
       module->addGlobal(builder.makeGlobal(stackLimitName,
                                            stackPointer->type,
-                                           builder.makeConst(int32_t(0)),
+                                           builder.makeConstPtr(0),
                                            Builder::Mutable));
 
     // Instrument all the code.
@@ -162,10 +164,12 @@ struct StackCheck : public Pass {
 
     // Generate the exported function.
     auto limitsFunc = builder.makeFunction(
-      SET_STACK_LIMITS, Signature({Type::i32, Type::i32}, Type::none), {});
-    auto* getBase = builder.makeLocalGet(0, Type::i32);
+      SET_STACK_LIMITS,
+      Signature({stackPointer->type, stackPointer->type}, Type::none),
+      {});
+    auto* getBase = builder.makeLocalGet(0, stackPointer->type);
     auto* storeBase = builder.makeGlobalSet(stackBaseName, getBase);
-    auto* getLimit = builder.makeLocalGet(1, Type::i32);
+    auto* getLimit = builder.makeLocalGet(1, stackPointer->type);
     auto* storeLimit = builder.makeGlobalSet(stackLimitName, getLimit);
     limitsFunc->body = builder.makeBlock({storeBase, storeLimit});
     addExportedFunction(*module, std::move(limitsFunc));
diff --git a/test/lit/passes/stack-check-memory64.wast b/test/lit/passes/stack-check-memory64.wast
new file mode 100644
index 000000000..0eec0b036
--- /dev/null
+++ b/test/lit/passes/stack-check-memory64.wast
@@ -0,0 +1,93 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_test.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --stack-check --enable-memory64 -S -o - | filecheck %s
+
+(module
+ (memory i64 (data))
+ ;; CHECK:      (type $none_=>_i64 (func (result i64)))
+
+ ;; CHECK:      (type $i64_i64_=>_none (func (param i64 i64)))
+
+ ;; CHECK:      (global $sp (mut i64) (i64.const 0))
+ (global $sp (mut i64) (i64.const 0))
+ (func "use_stack" (result i64)
+  (global.set $sp (i64.const 42))
+  (global.get $sp)
+ )
+)
+;; CHECK:      (global $__stack_base (mut i64) (i64.const 0))
+
+;; CHECK:      (global $__stack_limit (mut i64) (i64.const 0))
+
+;; CHECK:      (memory $0 i64 0 65536)
+
+;; CHECK:      (data (i64.const 0) "")
+
+;; CHECK:      (export "use_stack" (func $0))
+
+;; CHECK:      (export "__set_stack_limits" (func $__set_stack_limits))
+
+;; CHECK:      (func $0 (result i64)
+;; CHECK-NEXT:  (local $0 i64)
+;; CHECK-NEXT:  (block
+;; CHECK-NEXT:   (if
+;; CHECK-NEXT:    (i32.or
+;; CHECK-NEXT:     (i64.gt_u
+;; CHECK-NEXT:      (local.tee $0
+;; CHECK-NEXT:       (i64.const 42)
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (global.get $__stack_base)
+;; CHECK-NEXT:     )
+;; CHECK-NEXT:     (i64.lt_u
+;; CHECK-NEXT:      (local.get $0)
+;; CHECK-NEXT:      (global.get $__stack_limit)
+;; CHECK-NEXT:     )
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (unreachable)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (global.set $sp
+;; CHECK-NEXT:    (local.get $0)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.get $sp)
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $__set_stack_limits (param $0 i64) (param $1 i64)
+;; CHECK-NEXT:  (global.set $__stack_base
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.set $__stack_limit
+;; CHECK-NEXT:   (local.get $1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+(module
+ ;; if the global names are taken we should not crash
+ (memory i64 (data))
+ ;; CHECK:      (type $i64_i64_=>_none (func (param i64 i64)))
+
+ ;; CHECK:      (global $sp (mut i64) (i64.const 0))
+ (global $sp (mut i64) (i64.const 0)))
+ ;; CHECK:      (global $__stack_base (mut i64) (i64.const 0))
+ (global $__stack_base (mut i64) (i64.const 0))
+ ;; CHECK:      (global $__stack_limit (mut i64) (i64.const 0))
+ (global $__stack_limit (mut i64) (i64.const 0))
+ (export "use_stack" (func $0))
+ (func $0 (result i64)
+  (unreachable)
+ )
+)
+;; CHECK:      (memory $0 i64 0 65536)
+
+;; CHECK:      (data (i64.const 0) "")
+
+;; CHECK:      (export "__set_stack_limits" (func $__set_stack_limits))
+
+;; CHECK:      (func $__set_stack_limits (param $0 i64) (param $1 i64)
+;; CHECK-NEXT:  (global.set $__stack_base
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.set $__stack_limit
+;; CHECK-NEXT:   (local.get $1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
diff --git a/test/lld/basic_safe_stack.s b/test/lld/basic_safe_stack.s
index 74614e798..6422bcda4 100644
--- a/test/lld/basic_safe_stack.s
+++ b/test/lld/basic_safe_stack.s
@@ -1,5 +1,5 @@
 # Test that wasm-binaryen-finalize --check-stack-overflow correctly
-# inserts stack chekc handlers.
+# inserts stack check handlers.
 
 .globl stackRestore
 .globl stackAlloc