implement serializable closures

Signed-off-by: Achille Roussel <[email protected]>

.gitignore

@@ -0,0 +1,19 @@
+# Binaries for programs and plugins
+*.exe
+*.exe~
+*.dll
+*.so
+*.dylib
+
+# Test binary, built with `go test -c`
+*.test
+*.wasm
+
+# Output of the go coverage tool, specifically when used with LiteIDE
+*.out
+
+# Dependency directories (remove the comment below to include it)
+# vendor/
+
+# Emacs
+*~

compiler/coroutine_test.go

@@ -119,6 +119,12 @@ func TestCoroutineYield(t *testing.T) {
 			yields: []int{0, 3, 6, 9},
 		},
 
+		{
+			name:   "range over closure",
+			coro:   func() { Range10Closure() },
+			yields: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
+		},
+
 		{
 			name:   "select",
 			coro:   func() { Select(8) },

compiler/function.go

@@ -55,17 +55,51 @@ func generateFunctypesInit(pkg *ssa.Package, fn *ssa.Function, init *ast.BlockSt
 		return // ignore non-instantiated generic functions
 	}
 
-	init.List = append(init.List, &ast.ExprStmt{
-		X: &ast.CallExpr{
-			Fun: &ast.IndexListExpr{
-				X: &ast.SelectorExpr{
-					X:   ast.NewIdent("_types"),
-					Sel: ast.NewIdent("RegisterFunc"),
-				},
-				Indices: []ast.Expr{
-					newFuncType(fn.Signature),
+	var register ast.Expr
+	if len(fn.FreeVars) == 0 {
+		register = &ast.IndexListExpr{
+			X: &ast.SelectorExpr{
+				X:   ast.NewIdent("_types"),
+				Sel: ast.NewIdent("RegisterFunc"),
+			},
+			Indices: []ast.Expr{
+				newFuncType(fn.Signature),
+			},
+		}
+	} else {
+		fields := make([]*ast.Field, 1+len(fn.FreeVars))
+		// first field is the function address (uintptr)
+		fields[0] = &ast.Field{
+			Names: []*ast.Ident{ast.NewIdent("_")},
+			Type:  ast.NewIdent("uintptr"),
+		}
+
+		for i, freeVar := range fn.FreeVars {
+			fields[i+1] = &ast.Field{
+				Names: []*ast.Ident{ast.NewIdent(freeVar.Name())},
+				Type:  typeExpr(freeVar.Type()),
+			}
+		}
+
+		register = &ast.IndexListExpr{
+			X: &ast.SelectorExpr{
+				X:   ast.NewIdent("_types"),
+				Sel: ast.NewIdent("RegisterClosure"),
+			},
+			Indices: []ast.Expr{
+				newFuncType(fn.Signature),
+				&ast.StructType{
+					Fields: &ast.FieldList{
+						List: fields,
+					},
 				},
 			},
+		}
+	}
+
+	init.List = append(init.List, &ast.ExprStmt{
+		X: &ast.CallExpr{
+			Fun: register,
 			Args: []ast.Expr{
 				&ast.BasicLit{
 					Kind:  token.STRING,

compiler/testdata/coroutine.go

@@ -279,6 +279,22 @@ func RangeTripleFuncValue(n int) {
 	Range(n, f)
 }
 
+func Range10Closure() {
+	i := 0
+	n := 10
+	f := func() bool {
+		if i < n {
+			coroutine.Yield[int, any](i)
+			i++
+			return true
+		}
+		return false
+	}
+
+	for f() {
+	}
+}
+
 func Select(n int) {
 	select {
 	default:

internal/serde/reflect.go

@@ -404,23 +404,39 @@ func deserializePointer(d *Deserializer, t reflect.Type, p unsafe.Pointer) {
 }
 
 func serializeStruct(s *Serializer, t reflect.Type, p unsafe.Pointer) {
-	n := t.NumField()
+	serializeStructFields(s, p, t.NumField(), t.Field)
+}
+
+func deserializeStruct(d *Deserializer, t reflect.Type, p unsafe.Pointer) {
+	deserializeStructFields(d, p, t.NumField(), t.Field)
+}
+
+func serializeStructFields(s *Serializer, p unsafe.Pointer, n int, field func(int) reflect.StructField) {
 	for i := 0; i < n; i++ {
-		ft := t.Field(i)
+		ft := field(i)
 		fp := unsafe.Add(p, ft.Offset)
 		SerializeAny(s, ft.Type, fp)
 	}
 }
 
-func deserializeStruct(d *Deserializer, t reflect.Type, p unsafe.Pointer) {
-	n := t.NumField()
+func deserializeStructFields(d *Deserializer, p unsafe.Pointer, n int, field func(int) reflect.StructField) {
 	for i := 0; i < n; i++ {
-		ft := t.Field(i)
+		ft := field(i)
 		fp := unsafe.Add(p, ft.Offset)
 		DeserializeAny(d, ft.Type, fp)
 	}
 }
 
+type closureType struct {
+	_ uintptr
+	a *int
+	b *int
+}
+
+func (c *closureType) String() string {
+	return fmt.Sprintf("{a:%v,b:%v}", *c.a, *c.b)
+}
+
 func serializeFunc(s *Serializer, t reflect.Type, p unsafe.Pointer) {
 	// p is a pointer to a function value, function values are pointers to a
 	// memory location starting with the address of the function, hence the
@@ -433,8 +449,11 @@ func serializeFunc(s *Serializer, t reflect.Type, p unsafe.Pointer) {
 	fn := types.FuncByAddr(*(*uintptr)(p))
 	SerializeString(s, &fn.Name)
 
-	if fn.Closure != nil { // TODO
-		panic("cannot serialize closures yet")
+	if fn.Closure != nil {
+		t := fn.Closure
+		serializeStructFields(s, p, t.NumField()-1, func(i int) reflect.StructField {
+			return t.Field(i + 1)
+		})
 	}
 }
 
@@ -454,7 +473,17 @@ func deserializeFunc(d *Deserializer, t reflect.Type, p unsafe.Pointer) {
 	}
 
 	if fn.Closure != nil {
-		panic(name + ": cannot deserialize closures yet")
+		t := fn.Closure
+		v := reflect.New(t)
+
+		closure := v.UnsafePointer()
+		*(*uintptr)(closure) = fn.Addr
+
+		deserializeStructFields(d, closure, t.NumField()-1, func(i int) reflect.StructField {
+			return t.Field(i + 1)
+		})
+
+		*(*unsafe.Pointer)(p) = closure
 	} else {
 		*(**types.Func)(p) = fn
 	}

types/func.go

@@ -59,6 +59,16 @@ func RegisterFunc[Type any](name string) {
 	}
 }
 
+// RegisterClosure is like RegisterFunc but the caller can specify the closure
+// type (see types.Func for details).
+func RegisterClosure[Type, Closure any](name string) {
+	if f := FuncByName(name); f != nil {
+		var signature Type
+		var closure Closure
+		f.Type, f.Closure = reflect.TypeOf(signature), reflect.TypeOf(closure)
+	}
+}
+
 // Go function values are pointers to an object starting with the function
 // address, whether they are referencing top-level functions or closures.
 //