4 files changed, 111 insertions, 127 deletions
diff --git a/tools/go_marshal/test/BUILD b/tools/go_marshal/test/BUILD
index 3b839799d..2fbcc8a03 100644
--- a/tools/go_marshal/test/BUILD
+++ b/tools/go_marshal/test/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_binary", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library", "go_test")
 
 licenses(["notice"])
 
@@ -25,21 +25,10 @@ go_library(
     testonly = 1,
     srcs = ["test.go"],
     marshal = True,
+    visibility = ["//tools/go_marshal/test:__subpackages__"],
     deps = ["//tools/go_marshal/test/external"],
 )
 
-go_binary(
-    name = "escape",
-    testonly = 1,
-    srcs = ["escape.go"],
-    gc_goopts = ["-m"],
-    deps = [
-        ":test",
-        "//pkg/usermem",
-        "//tools/go_marshal/marshal",
-    ],
-)
-
 go_test(
     name = "marshal_test",
     size = "small",
diff --git a/tools/go_marshal/test/escape.go b/tools/go_marshal/test/escape.go
deleted file mode 100644
index 184f05ea3..000000000
--- a/tools/go_marshal/test/escape.go
+++ /dev/null
@@ -1,114 +0,0 @@
-// Copyright 2020 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-// This binary provides a convienient target for analyzing how the go-marshal
-// API causes its various arguments to escape to the heap. To use, build and
-// observe the output from the go compiler's escape analysis:
-//
-// $ bazel build :escape
-// ...
-// escape.go:67:2: moved to heap: task
-// escape.go:77:31: make([]byte, size) escapes to heap
-// escape.go:87:31: make([]byte, size) escapes to heap
-// escape.go:96:6: moved to heap: stat
-// ...
-//
-// This is not an automated test, but simply a minimal binary for easy analysis.
-package main
-
-import (
-	"gvisor.dev/gvisor/pkg/usermem"
-	"gvisor.dev/gvisor/tools/go_marshal/marshal"
-	"gvisor.dev/gvisor/tools/go_marshal/test"
-)
-
-// dummyTask implements marshal.Task.
-type dummyTask struct {
-}
-
-func (*dummyTask) CopyScratchBuffer(size int) []byte {
-	return make([]byte, size)
-}
-
-func (*dummyTask) CopyOutBytes(addr usermem.Addr, b []byte) (int, error) {
-	return len(b), nil
-}
-
-func (*dummyTask) CopyInBytes(addr usermem.Addr, b []byte) (int, error) {
-	return len(b), nil
-}
-
-func (task *dummyTask) MarshalBytes(addr usermem.Addr, marshallable marshal.Marshallable) {
-	buf := task.CopyScratchBuffer(marshallable.SizeBytes())
-	marshallable.MarshalBytes(buf)
-	task.CopyOutBytes(addr, buf)
-}
-
-func (task *dummyTask) MarshalUnsafe(addr usermem.Addr, marshallable marshal.Marshallable) {
-	buf := task.CopyScratchBuffer(marshallable.SizeBytes())
-	marshallable.MarshalUnsafe(buf)
-	task.CopyOutBytes(addr, buf)
-}
-
-// Expected escapes:
-// - task: passed to marshal.Marshallable.CopyOut as the marshal.Task interface.
-func doCopyOut() {
-	task := dummyTask{}
-	var stat test.Stat
-	stat.CopyOut(&task, usermem.Addr(0xf000ba12))
-}
-
-// Expected escapes:
-// - buf: make allocates on the heap.
-func doMarshalBytesDirect() {
-	task := dummyTask{}
-	var stat test.Stat
-	buf := task.CopyScratchBuffer(stat.SizeBytes())
-	stat.MarshalBytes(buf)
-	task.CopyOutBytes(usermem.Addr(0xf000ba12), buf)
-}
-
-// Expected escapes:
-// - buf: make allocates on the heap.
-func doMarshalUnsafeDirect() {
-	task := dummyTask{}
-	var stat test.Stat
-	buf := task.CopyScratchBuffer(stat.SizeBytes())
-	stat.MarshalUnsafe(buf)
-	task.CopyOutBytes(usermem.Addr(0xf000ba12), buf)
-}
-
-// Expected escapes:
-// - stat: passed to dummyTask.MarshalBytes as the marshal.Marshallable interface.
-func doMarshalBytesViaMarshallable() {
-	task := dummyTask{}
-	var stat test.Stat
-	task.MarshalBytes(usermem.Addr(0xf000ba12), &stat)
-}
-
-// Expected escapes:
-// - stat: passed to dummyTask.MarshalUnsafe as the marshal.Marshallable interface.
-func doMarshalUnsafeViaMarshallable() {
-	task := dummyTask{}
-	var stat test.Stat
-	task.MarshalUnsafe(usermem.Addr(0xf000ba12), &stat)
-}
-
-func main() {
-	doCopyOut()
-	doMarshalBytesDirect()
-	doMarshalUnsafeDirect()
-	doMarshalBytesViaMarshallable()
-	doMarshalUnsafeViaMarshallable()
-}
diff --git a/tools/go_marshal/test/escape/BUILD b/tools/go_marshal/test/escape/BUILD
new file mode 100644
index 000000000..f74e6ffae
--- /dev/null
+++ b/tools/go_marshal/test/escape/BUILD
@@ -0,0 +1,14 @@
+load("//tools:defs.bzl", "go_library")
+
+licenses(["notice"])
+
+go_library(
+    name = "escape",
+    testonly = 1,
+    srcs = ["escape.go"],
+    deps = [
+        "//pkg/usermem",
+        "//tools/go_marshal/marshal",
+        "//tools/go_marshal/test",
+    ],
+)
diff --git a/tools/go_marshal/test/escape/escape.go b/tools/go_marshal/test/escape/escape.go
new file mode 100644
index 000000000..6a46ddbf8
--- /dev/null
+++ b/tools/go_marshal/test/escape/escape.go
@@ -0,0 +1,95 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package escape
+
+import (
+	"gvisor.dev/gvisor/pkg/usermem"
+	"gvisor.dev/gvisor/tools/go_marshal/marshal"
+	"gvisor.dev/gvisor/tools/go_marshal/test"
+)
+
+// dummyTask implements marshal.Task.
+type dummyTask struct {
+}
+
+func (*dummyTask) CopyScratchBuffer(size int) []byte {
+	return make([]byte, size)
+}
+
+func (*dummyTask) CopyOutBytes(addr usermem.Addr, b []byte) (int, error) {
+	return len(b), nil
+}
+
+func (*dummyTask) CopyInBytes(addr usermem.Addr, b []byte) (int, error) {
+	return len(b), nil
+}
+
+func (t *dummyTask) MarshalBytes(addr usermem.Addr, marshallable marshal.Marshallable) {
+	buf := t.CopyScratchBuffer(marshallable.SizeBytes())
+	marshallable.MarshalBytes(buf)
+	t.CopyOutBytes(addr, buf)
+}
+
+func (t *dummyTask) MarshalUnsafe(addr usermem.Addr, marshallable marshal.Marshallable) {
+	buf := t.CopyScratchBuffer(marshallable.SizeBytes())
+	marshallable.MarshalUnsafe(buf)
+	t.CopyOutBytes(addr, buf)
+}
+
+// +checkescape:all
+//go:nosplit
+func doCopyIn(t *dummyTask) {
+	var stat test.Stat
+	stat.CopyIn(t, usermem.Addr(0xf000ba12))
+}
+
+// +checkescape:all
+//go:nosplit
+func doCopyOut(t *dummyTask) {
+	var stat test.Stat
+	stat.CopyOut(t, usermem.Addr(0xf000ba12))
+}
+
+// +mustescape:builtin
+// +mustescape:stack
+func doMarshalBytesDirect(t *dummyTask) {
+	var stat test.Stat
+	buf := t.CopyScratchBuffer(stat.SizeBytes())
+	stat.MarshalBytes(buf)
+	t.CopyOutBytes(usermem.Addr(0xf000ba12), buf)
+}
+
+// +mustescape:builtin
+// +mustescape:stack
+func doMarshalUnsafeDirect(t *dummyTask) {
+	var stat test.Stat
+	buf := t.CopyScratchBuffer(stat.SizeBytes())
+	stat.MarshalUnsafe(buf)
+	t.CopyOutBytes(usermem.Addr(0xf000ba12), buf)
+}
+
+// +mustescape:local,heap
+// +mustescape:stack
+func doMarshalBytesViaMarshallable(t *dummyTask) {
+	var stat test.Stat
+	t.MarshalBytes(usermem.Addr(0xf000ba12), &stat)
+}
+
+// +mustescape:local,heap
+// +mustescape:stack
+func doMarshalUnsafeViaMarshallable(t *dummyTask) {
+	var stat test.Stat
+	t.MarshalUnsafe(usermem.Addr(0xf000ba12), &stat)
+}