3 files changed, 139 insertions, 19 deletions
diff --git a/test/lit/wasm-split/call_exports.mjs b/test/lit/wasm-split/call_exports.mjs
index 546564d2e..592fbad48 100644
--- a/test/lit/wasm-split/call_exports.mjs
+++ b/test/lit/wasm-split/call_exports.mjs
@@ -20,6 +20,6 @@ for (let i = 4; i < process.argv.length; i++) {
 }
 
 // Create and read the profile
-let profileSize = instance.exports['__write_profile'](0, 2**32 - 1);
-let profileData = Buffer.from(instance.exports.memory.buffer, 0, profileSize);
+let profileSize = instance.exports['__write_profile'](1024, 2**32 - 1024);
+let profileData = Buffer.from(instance.exports.memory.buffer, 1024, profileSize);
 fs.writeFileSync(outFile, profileData);
diff --git a/test/lit/wasm-split/instrument-in-memory.wast b/test/lit/wasm-split/instrument-in-memory.wast
new file mode 100644
index 000000000..568ce2207
--- /dev/null
+++ b/test/lit/wasm-split/instrument-in-memory.wast
@@ -0,0 +1,88 @@
+;; RUN: wasm-split %s --instrument --in-memory -all -S -o - | filecheck %s
+
+;; Check that the output round trips and validates as well
+;; RUN: wasm-split %s --instrument --in-memory -all -g -o %t.wasm
+;; RUN: wasm-opt -all %t.wasm -S -o -
+
+(module
+  (import "env" "foo" (func $foo))
+  (export "bar" (func $bar))
+  (func $bar
+    (call $foo)
+  )
+  (func $baz (param i32) (result i32)
+    (local.get 0)
+  )
+)
+
+;; Check that a memory has been added
+;; CHECK: (memory $0 1 1)
+
+;; And the profiling function exported
+;; CHECK: (export "__write_profile" (func $__write_profile))
+
+;; Check that the function instrumentation is correct
+
+;; CHECK:      (func $bar
+;; CHECK-NEXT:  (i32.atomic.store8
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $foo)
+;; CHECK-NEXT: )
+
+;; CHECK-NEXT: (func $baz (param $0 i32) (result i32)
+;; CHECK-NEXT:  (i32.atomic.store8 offset=1
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (local.get $0)
+;; CHECK-NEXT: )
+
+;; Check that the profiling function is correct.
+
+;; CHECK:      (func $__write_profile (param $addr i32) (param $size i32) (result i32)
+;; CHECK-NEXT:  (local $funcIdx i32)
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.ge_u
+;; CHECK-NEXT:    (local.get $size)
+;; CHECK-NEXT:    (i32.const 16)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (block
+;; CHECK-NEXT:    (i64.store align=1
+;; CHECK-NEXT:     (local.get $addr)
+;; CHECK-NEXT:     (i64.const {{.*}})
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (block $outer
+;; CHECK-NEXT:     (loop $l
+;; CHECK-NEXT:      (br_if $outer
+;; CHECK-NEXT:       (i32.eq
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:        (i32.const 2)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (i32.store offset=8
+;; CHECK-NEXT:       (i32.add
+;; CHECK-NEXT:        (local.get $addr)
+;; CHECK-NEXT:        (i32.mul
+;; CHECK-NEXT:         (local.get $funcIdx)
+;; CHECK-NEXT:         (i32.const 4)
+;; CHECK-NEXT:        )
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:       (i32.atomic.load8_u
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (local.set $funcIdx
+;; CHECK-NEXT:       (i32.add
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:        (i32.const 1)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (br $l)
+;; CHECK-NEXT:     )
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (i32.const 16)
+;; CHECK-NEXT: )
diff --git a/test/lit/wasm-split/profile-guided.wast b/test/lit/wasm-split/profile-guided.wast
index f2d79f99c..4a96ef55e 100644
--- a/test/lit/wasm-split/profile-guided.wast
+++ b/test/lit/wasm-split/profile-guided.wast
@@ -1,43 +1,75 @@
+;; =====================
 ;; Instrument the binary
+;; =====================
 
-;; RUN: wasm-split --instrument %s -o %t.instrumented.wasm
+;; RUN: wasm-split -all --instrument %s -o %t.instrumented.wasm
 
 ;; Create profiles
 
-;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
-;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
-;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
-;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.none.prof
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.none.prof
 
 ;; Create profile-guided splits
 
-;; RUN: wasm-split %s --profile=%t.foo.prof -v -o1 %t.foo.1.wasm -o2 %t.foo.2.wasm \
+;; RUN: wasm-split -all %s --profile=%t.foo.prof -v -o1 %t.foo.1.wasm -o2 %t.foo.2.wasm \
 ;; RUN:   | filecheck %s --check-prefix FOO
 
-;; FOO: Keeping functions: deep_foo_callee, foo, foo_callee, shared_callee
-;; FOO: Splitting out functions: bar, bar_callee, uncalled
+;; RUN: wasm-split -all %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BAR
+
+;; RUN: wasm-split -all %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BOTH
+
+;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
+;; RUN:   | filecheck %s --check-prefix NONE
+
+;; =================================
+;; Do it all again using --in-memory
+;; =================================
 
-;; RUN: wasm-split %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN: wasm-split -all --instrument --in-memory %s -o %t.instrumented.wasm
+
+;; Create profiles
+
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.none.prof
+
+;; Create profile-guided splits
+
+;; RUN: wasm-split -all %s --profile=%t.foo.prof -v -o1 %t.foo.1.wasm -o2 %t.foo.2.wasm \
+;; RUN:   | filecheck %s --check-prefix FOO
+
+;; RUN: wasm-split -all %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
 ;; RUN:   | filecheck %s --check-prefix BAR
 
+;; RUN: wasm-split -all %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BOTH
+
+;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
+;; RUN:   | filecheck %s --check-prefix NONE
+
+;; =======
+;; Results
+;; =======
+
+;; FOO: Keeping functions: deep_foo_callee, foo, foo_callee, shared_callee
+;; FOO: Splitting out functions: bar, bar_callee, uncalled
+
 ;; BAR: Keeping functions: bar, bar_callee, shared_callee
 ;; BAR: Splitting out functions: deep_foo_callee, foo, foo_callee, uncalled
 
-;; RUN: wasm-split %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
-;; RUN:   | filecheck %s --check-prefix BOTH
-
 ;; BOTH: Keeping functions: bar, bar_callee, deep_foo_callee, foo, foo_callee, shared_callee
 ;; BOTH: Splitting out functions: uncalled
 
-;; RUN: wasm-split %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
-;; RUN:   | filecheck %s --check-prefix NONE
-
 ;; NONE: Keeping functions:
 ;; NONE: Splitting out functions: bar, bar_callee, deep_foo_callee, foo, foo_callee, shared_callee, uncalled
 
-
 (module
-  (memory $mem 1 1)
+  (memory $mem (shared 1 1))
   (export "memory" (memory $mem))
   (export "foo" (func $foo))
   (export "bar" (func $bar))