Improve cursor states, add test

cursor.go

@@ -9,14 +9,15 @@ var ErrCursourOutObFounds = errors.New("out of bounds")
 
 type Cursor struct {
 	currentIdx int
-	savedIdx   int
 	statements []ast.Stmt
+	saves      []int
 }
 
 func NewCursor(i int, statements []ast.Stmt) *Cursor {
 	return &Cursor{
 		currentIdx: i,
 		statements: statements,
+		saves:      []int{},
 	}
 }
 
@@ -45,9 +46,15 @@ func (c *Cursor) Stmt() ast.Stmt {
 }
 
 func (c *Cursor) Save() {
-	c.savedIdx = c.currentIdx
+	c.saves = append(c.saves, c.currentIdx)
 }
 
 func (c *Cursor) Reset() {
-	c.currentIdx = c.savedIdx
+	if len(c.saves) == 0 {
+		return
+	}
+
+	idx := c.saves[len(c.saves)-1]
+	c.saves = c.saves[:len(c.saves)-1]
+	c.currentIdx = idx
 }

cursor_test.go

@@ -0,0 +1,49 @@
+package wsl
+
+import (
+	"go/ast"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSavestate(t *testing.T) {
+	cursor := NewCursor(0, []ast.Stmt{
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+		&ast.AssignStmt{},
+	})
+
+	cursor.Save()
+	cursor.Next()
+	cursor.Next()
+
+	func() {
+		cursor.Save()
+		defer cursor.Reset()
+
+		cursor.Next()
+		cursor.Next()
+
+		func() {
+			cursor.Save()
+			defer cursor.Reset()
+
+			cursor.Next()
+			cursor.Next()
+
+			assert.Equal(t, 6, cursor.currentIdx)
+		}()
+
+		assert.Equal(t, 4, cursor.currentIdx)
+	}()
+
+	assert.Equal(t, 2, cursor.currentIdx)
+
+	cursor.Reset()
+	assert.Equal(t, 0, cursor.currentIdx)
+}

go.mod

@@ -6,6 +6,9 @@ require golang.org/x/tools v0.19.0
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/testify v1.9.0 // indirect
 	golang.org/x/mod v0.16.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,5 +1,9 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/mod v0.8.0 h1:LUYupSeNrTNCGzR/hVBk2NHZO4hXcVaW1k4Qx7rjPx8=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/mod v0.16.0 h1:QX4fJ0Rr5cPQCF7O9lh9Se4pmwfwskqZfq5moyldzic=
@@ -13,3 +17,6 @@ golang.org/x/tools v0.6.0 h1:BOw41kyTf3PuCW1pVQf8+Cyg8pMlkYB1oo9iJ6D/lKM=
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/tools v0.19.0 h1:tfGCXNR1OsFG+sVdLAitlpjAvD/I6dHDKnYrpEZUHkw=
 golang.org/x/tools v0.19.0/go.mod h1:qoJWxmGSIBmAeriMx19ogtrEPrGtDbPK634QFIcLAhc=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

wsl.go

@@ -68,11 +68,8 @@ func (w *WSL) CheckFunc(funcDecl *ast.FuncDecl) {
 }
 
 func (w *WSL) CheckIf(stmt *ast.IfStmt, cursor *Cursor) {
-	// TODO: This seems finicky, we don't want to do this.
-	current := cursor.currentIdx
+	cursor.Save()
 
-	// TODO: If we're _not_ cuddled, check if the previous variable implements
-	// an error and if that's used in the if statement.
 	currentIdents := allIdents(cursor.Stmt())
 
 	previousIdents := []*ast.Ident{}
@@ -146,7 +143,7 @@ func (w *WSL) CheckIf(stmt *ast.IfStmt, cursor *Cursor) {
 		}
 	}
 
-	cursor.currentIdx = current
+	cursor.Reset()
 
 	// if
 	w.CheckBlock(stmt.Body)