compiler: add //go:noescape pragma

This only works on declarations, not definitions. This is intentional:
it follows the upstream Go implemetation.

However, we might want to loosen this requirement at some point: TinyGo
sometimes stores pointers in memory mapped I/O knowing they won't
actually escape, but the compiler doesn't know about this.

4 files changed, 44 insertions, 8 deletions

diff --git a/compiler/calls.go b/compiler/calls.go
index f4b76a513..6400e634b 100644
--- a/compiler/calls.go
+++ b/compiler/calls.go
@@ -19,8 +19,9 @@ const maxFieldsPerParam = 3
 // useful while declaring or defining a function.
 type paramInfo struct {
 	llvmType llvm.Type
-	name     string // name, possibly with suffixes for e.g. struct fields
-	elemSize uint64 // size of pointer element type, or 0 if this isn't a pointer
+	name     string     // name, possibly with suffixes for e.g. struct fields
+	elemSize uint64     // size of pointer element type, or 0 if this isn't a pointer
+	flags    paramFlags // extra flags for this parameter
 }
 
 // paramFlags identifies parameter attributes for flags. Most importantly, it
@@ -28,9 +29,9 @@ type paramInfo struct {
 type paramFlags uint8
 
 const (
-	// Parameter may have the deferenceable_or_null attribute. This attribute
-	// cannot be applied to unsafe.Pointer and to the data pointer of slices.
-	paramIsDeferenceableOrNull = 1 << iota
+	// Whether this is a full or partial Go parameter (int, slice, etc).
+	// The extra context parameter is not a Go parameter.
+	paramIsGoParam = 1 << iota
 )
 
 // createRuntimeCallCommon creates a runtime call. Use createRuntimeCall or
@@ -195,6 +196,7 @@ func (c *compilerContext) getParamInfo(t llvm.Type, name string, goType types.Ty
 	info := paramInfo{
 		llvmType: t,
 		name:     name,
+		flags:    paramIsGoParam,
 	}
 	if goType != nil {
 		switch underlying := goType.Underlying().(type) {
diff --git a/compiler/symbol.go b/compiler/symbol.go
index e53df30f1..93c27803e 100644
--- a/compiler/symbol.go
+++ b/compiler/symbol.go
@@ -33,6 +33,7 @@ type functionInfo struct {
 	exported      bool       // go:export, CGo
 	interrupt     bool       // go:interrupt
 	nobounds      bool       // go:nobounds
+	noescape      bool       // go:noescape
 	variadic      bool       // go:variadic (CGo only)
 	inline        inlineType // go:inline
 }
@@ -127,11 +128,20 @@ func (c *compilerContext) getFunction(fn *ssa.Function) (llvm.Type, llvm.Value)
 	c.addStandardDeclaredAttributes(llvmFn)
 
 	dereferenceableOrNullKind := llvm.AttributeKindID("dereferenceable_or_null")
-	for i, info := range paramInfos {
-		if info.elemSize != 0 {
-			dereferenceableOrNull := c.ctx.CreateEnumAttribute(dereferenceableOrNullKind, info.elemSize)
+	for i, paramInfo := range paramInfos {
+		if paramInfo.elemSize != 0 {
+			dereferenceableOrNull := c.ctx.CreateEnumAttribute(dereferenceableOrNullKind, paramInfo.elemSize)
 			llvmFn.AddAttributeAtIndex(i+1, dereferenceableOrNull)
 		}
+		if info.noescape && paramInfo.flags&paramIsGoParam != 0 && paramInfo.llvmType.TypeKind() == llvm.PointerTypeKind {
+			// Parameters to functions with a //go:noescape parameter should get
+			// the nocapture attribute. However, the context parameter should
+			// not.
+			// (It may be safe to add the nocapture parameter to the context
+			// parameter, but I'd like to stay on the safe side here).
+			nocapture := c.ctx.CreateEnumAttribute(llvm.AttributeKindID("nocapture"), 0)
+			llvmFn.AddAttributeAtIndex(i+1, nocapture)
+		}
 	}
 
 	// Set a number of function or parameter attributes, depending on the
@@ -394,6 +404,13 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 			if hasUnsafeImport(f.Pkg.Pkg) {
 				info.nobounds = true
 			}
+		case "//go:noescape":
+			// Don't let pointer parameters escape.
+			// Following the upstream Go implementation, we only do this for
+			// declarations, not definitions.
+			if len(f.Blocks) == 0 {
+				info.noescape = true
+			}
 		case "//go:variadic":
 			// The //go:variadic pragma is emitted by the CGo preprocessing
 			// pass for C variadic functions. This includes both explicit
diff --git a/compiler/testdata/pragma.go b/compiler/testdata/pragma.go
index 1f6badf7f..1e6e967f5 100644
--- a/compiler/testdata/pragma.go
+++ b/compiler/testdata/pragma.go
@@ -106,3 +106,12 @@ var undefinedGlobalNotInSection uint32
 //go:align 1024
 //go:section .global_section
 var multipleGlobalPragmas uint32
+
+//go:noescape
+func doesNotEscapeParam(a *int, b []int, c chan int, d *[0]byte)
+
+// The //go:noescape pragma only works on declarations, not definitions.
+//
+//go:noescape
+func stillEscapes(a *int, b []int, c chan int, d *[0]byte) {
+}
diff --git a/compiler/testdata/pragma.ll b/compiler/testdata/pragma.ll
index 28e678359..fa8015a90 100644
--- a/compiler/testdata/pragma.ll
+++ b/compiler/testdata/pragma.ll
@@ -85,6 +85,14 @@ entry:
 
 declare void @main.undefinedFunctionNotInSection(ptr) #1
 
+declare void @main.doesNotEscapeParam(ptr nocapture dereferenceable_or_null(4), ptr nocapture, i32, i32, ptr nocapture dereferenceable_or_null(32), ptr nocapture, ptr) #1
+
+; Function Attrs: nounwind
+define hidden void @main.stillEscapes(ptr dereferenceable_or_null(4) %a, ptr %b.data, i32 %b.len, i32 %b.cap, ptr dereferenceable_or_null(32) %c, ptr %d, ptr %context) unnamed_addr #2 {
+entry:
+  ret void
+}
+
 attributes #0 = { allockind("alloc,zeroed") allocsize(0) "alloc-family"="runtime.alloc" "target-features"="+bulk-memory,+mutable-globals,+nontrapping-fptoint,+sign-ext" }
 attributes #1 = { "target-features"="+bulk-memory,+mutable-globals,+nontrapping-fptoint,+sign-ext" }
 attributes #2 = { nounwind "target-features"="+bulk-memory,+mutable-globals,+nontrapping-fptoint,+sign-ext" }