From af67c21522824cdeb31ad527cd9a136b1551c396 Mon Sep 17 00:00:00 2001
From: Chris O'Hara <cohara87@gmail.com>
Date: Tue, 25 Jun 2024 08:17:57 +1000
Subject: [PATCH] Allow go statements in certain circumstances

---
 compiler/compile.go     |  4 ++--
 compiler/unsupported.go | 22 ++++++++++++++++++++--
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/compiler/compile.go b/compiler/compile.go
index 2c7c756..b50a819 100644
--- a/compiler/compile.go
+++ b/compiler/compile.go
@@ -379,7 +379,7 @@ func (c *compiler) compilePackage(p *packages.Package, colors functionColors) er
 				compiled := false
 				if color != nil || containsColoredFuncLit(decl, colorsByFunc) {
 					// Reject certain language features for now.
-					if err := unsupported(decl, p.TypesInfo); err != nil {
+					if err := c.unsupported(decl, p.TypesInfo, colorsByFunc); err != nil {
 						return err
 					}
 					scope := &scope{compiler: c, colors: colorsByFunc}
@@ -433,7 +433,7 @@ func (c *compiler) compilePackage(p *packages.Package, colors functionColors) er
 	return nil
 }
 
-func containsColoredFuncLit(decl *ast.FuncDecl, colorsByFunc map[ast.Node]*types.Signature) (yes bool) {
+func containsColoredFuncLit(decl ast.Node, colorsByFunc map[ast.Node]*types.Signature) (yes bool) {
 	ast.Inspect(decl, func(n ast.Node) bool {
 		if lit, ok := n.(*ast.FuncLit); ok {
 			if _, ok := colorsByFunc[lit]; ok {
diff --git a/compiler/unsupported.go b/compiler/unsupported.go
index 16644e9..48956b2 100644
--- a/compiler/unsupported.go
+++ b/compiler/unsupported.go
@@ -5,17 +5,35 @@ import (
 	"go/ast"
 	"go/token"
 	"go/types"
+	"log"
 )
 
 // unsupported checks a function for unsupported language features.
-func unsupported(decl ast.Node, info *types.Info) (err error) {
+func (c *compiler) unsupported(decl ast.Node, info *types.Info, colorsByFunc map[ast.Node]*types.Signature) (err error) {
 	ast.Inspect(decl, func(node ast.Node) bool {
 		switch nn := node.(type) {
 		case ast.Stmt:
 			switch n := nn.(type) {
 			// Not yet supported:
 			case *ast.GoStmt:
-				err = fmt.Errorf("not implemented: go")
+				if _, inColoredFunc := colorsByFunc[decl]; inColoredFunc {
+					err = fmt.Errorf("not implemented: go")
+				} else {
+					// Allow go statement in certain limited circumstances.
+					_, goColoredFunc := colorsByFunc[n.Call.Fun]
+					if !goColoredFunc {
+						switch fn := n.Call.Fun.(type) {
+						case *ast.FuncLit:
+							goColoredFunc = containsColoredFuncLit(fn, colorsByFunc)
+						}
+					}
+					if goColoredFunc {
+						err = fmt.Errorf("not implemented: go")
+					} else {
+						pos := c.fset.Position(n.Pos())
+						log.Printf("warning: goroutine mutations at %s may not be durable", pos)
+					}
+				}
 
 			// Partially supported:
 			case *ast.BranchStmt: