Support range over int

compiler/coroutine_test.go

@@ -249,6 +249,12 @@ func TestCoroutineYield(t *testing.T) {
 			coro:   func() { IndirectClosure(1) },
 			yields: []int{-1, 1, 2, 3},
 		},
+
+		{
+			name:   "range over int",
+			coro:   func() { RangeOverInt(3) },
+			yields: []int{0, 1, 2},
+		},
 	}
 
 	// This emulates the installation of function type information by the

compiler/desugar.go

@@ -246,7 +246,39 @@ func (d *desugarer) desugar(stmt ast.Stmt, breakTo, continueTo, userLabel *ast.I
 		}
 		prologue := d.desugarList([]ast.Stmt{init}, nil, nil)
 
+		intType := types.Typ[types.Int]
+
 		switch rangeElemType := d.info.TypeOf(s.X).(type) {
+		case *types.Basic:
+			switch rangeElemType.Kind() {
+			case types.Int:
+				// Rewrite for range loops over int:
+				// - `for range x {}` => `{ _x := x; for _i := 0; _i < _x; _i++ {} }`
+				// - `for _ := range x {}` => `{ _x := x; for _i := 0; _i < _x; _i++ {} }`
+				// - `for i := range x {}` => `{ _x := x; for i := 0; i < _x; i++ {} }`
+				var i *ast.Ident
+				if s.Key == nil || isUnderscore(s.Key) {
+					i = d.newVar(intType)
+				} else {
+					i = s.Key.(*ast.Ident)
+				}
+				forStmt := &ast.ForStmt{
+					Init: &ast.AssignStmt{Lhs: []ast.Expr{i}, Tok: token.DEFINE, Rhs: []ast.Expr{&ast.BasicLit{Kind: token.INT, Value: "0"}}},
+					Post: &ast.IncDecStmt{X: i, Tok: token.INC},
+					Cond: &ast.BinaryExpr{X: i, Op: token.LSS, Y: x},
+					Body: s.Body,
+				}
+				if d.mayYield(s.Body) {
+					d.nodesThatMayYield[forStmt] = struct{}{}
+				}
+				stmt = &ast.BlockStmt{
+					List: append(prologue, d.desugar(forStmt, breakTo, continueTo, userLabel)),
+				}
+
+			default:
+				panic(fmt.Sprintf("not implemented: for range over %T", rangeElemType))
+			}
+
 		case *types.Array, *types.Slice:
 			// Rewrite for range loops over arrays/slices:
 			// - `for range x {}` => `{ _x := x; for _i := 0; _i < len(_x); _i++ {} }`
@@ -259,7 +291,7 @@ func (d *desugarer) desugar(stmt ast.Stmt, breakTo, continueTo, userLabel *ast.I
 			// Then, desugar loops further (see ast.ForStmt case above).
 			var i *ast.Ident
 			if s.Key == nil || isUnderscore(s.Key) {
-				i = d.newVar(types.Typ[types.Int])
+				i = d.newVar(intType)
 			} else {
 				i = s.Key.(*ast.Ident)
 			}

compiler/desugar_test.go

@@ -1144,6 +1144,63 @@ defer func() {
 	_v0 := time.Now()
 	_v0.UTC()
 }
+`,
+		},
+		{
+			name: "for range over int",
+			body: "for range n { foo }",
+			info: func(stmts []ast.Stmt, info *types.Info) {
+				n := stmts[0].(*ast.RangeStmt).X
+				info.Types[n] = types.TypeAndValue{Type: intType}
+			},
+			expect: `
+{
+	_v0 := n
+	{
+		_v1 := 0
+		for ; _v1 < _v0; _v1++ {
+			foo
+		}
+	}
+}
+`,
+		},
+		{
+			name: "for range over int, with underscore index",
+			body: "for _ := range n { foo }",
+			info: func(stmts []ast.Stmt, info *types.Info) {
+				n := stmts[0].(*ast.RangeStmt).X
+				info.Types[n] = types.TypeAndValue{Type: intType}
+			},
+			expect: `
+{
+	_v0 := n
+	{
+		_v1 := 0
+		for ; _v1 < _v0; _v1++ {
+			foo
+		}
+	}
+}
+`,
+		},
+		{
+			name: "for range over int, with index",
+			body: "for i := range n { foo }",
+			info: func(stmts []ast.Stmt, info *types.Info) {
+				n := stmts[0].(*ast.RangeStmt).X
+				info.Types[n] = types.TypeAndValue{Type: intType}
+			},
+			expect: `
+{
+	_v0 := n
+	{
+		i := 0
+		for ; i < _v0; i++ {
+			foo
+		}
+	}
+}
 `,
 		},
 	} {

compiler/testdata/coroutine.go

@@ -656,3 +656,9 @@ func indirectClosure(m interface{ YieldAndInc() }) func() {
 		m.YieldAndInc()
 	}
 }
+
+func RangeOverInt(n int) {
+	for i := range n {
+		coroutine.Yield[int, any](i)
+	}
+}

compiler/testdata/coroutine_durable.go

@@ -3634,6 +3634,53 @@ func indirectClosure(_fn0 interface{ YieldAndInc() }) (_ func()) {
 	}
 	panic("unreachable")
 }
+
+//go:noinline
+func RangeOverInt(_fn0 int) {
+	_c := coroutine.LoadContext[int, any]()
+	var _f0 *struct {
+		IP int
+		X0 int
+		X1 int
+		X2 int
+	} = coroutine.Push[struct {
+		IP int
+		X0 int
+		X1 int
+		X2 int
+	}](&_c.Stack)
+	if _f0.IP == 0 {
+		*_f0 = struct {
+			IP int
+			X0 int
+			X1 int
+			X2 int
+		}{X0: _fn0}
+	}
+	defer func() {
+		if !_c.Unwinding() {
+			coroutine.Pop(&_c.Stack)
+		}
+	}()
+	switch {
+	case _f0.IP < 2:
+		_f0.X1 = _f0.X0
+		_f0.IP = 2
+		fallthrough
+	case _f0.IP < 4:
+		switch {
+		case _f0.IP < 3:
+			_f0.X2 = 0
+			_f0.IP = 3
+			fallthrough
+		case _f0.IP < 4:
+			for ; _f0.X2 < _f0.X1; _f0.X2, _f0.IP = _f0.X2+1, 3 {
+
+				coroutine.Yield[int, any](_f0.X2)
+			}
+		}
+	}
+}
 func init() {
 	_types.RegisterFunc[func(_fn1 int) (_ func(int))]("github.com/stealthrocket/coroutine/compiler/testdata.(*Box).Closure")
 	_types.RegisterClosure[func(_fn0 int), struct {
@@ -3741,6 +3788,7 @@ func init() {
 	}]("github.com/stealthrocket/coroutine/compiler/testdata.Range10ClosureHeterogenousCapture.func3")
 	_types.RegisterFunc[func()]("github.com/stealthrocket/coroutine/compiler/testdata.Range10Heterogenous")
 	_types.RegisterFunc[func(_ int)]("github.com/stealthrocket/coroutine/compiler/testdata.RangeArrayIndexValueGenerator")
+	_types.RegisterFunc[func(_fn0 int)]("github.com/stealthrocket/coroutine/compiler/testdata.RangeOverInt")
 	_types.RegisterFunc[func(_fn0 int)]("github.com/stealthrocket/coroutine/compiler/testdata.RangeOverMaps")
 	_types.RegisterFunc[func(_fn0 int)]("github.com/stealthrocket/coroutine/compiler/testdata.RangeReverseClosureCaptureByValue")
 	_types.RegisterClosure[func(), struct {

go.mod

@@ -1,6 +1,6 @@
 module github.com/stealthrocket/coroutine
 
-go 1.21.0
+go 1.22
 
 require (
 	golang.org/x/sync v0.5.0