From 93fb47cccccda3efb654d73acbc7384260f70506 Mon Sep 17 00:00:00 2001
From: Steven Hartland <stevenmhartland@gmail.com>
Date: Tue, 29 Nov 2022 11:09:56 +0000
Subject: [PATCH] fix(ast): walk missing nodes (#471)

Fix ast.Walk missing out:
* DotExpressions.Identifier
* LabelledStatement.Label

Fixes #299
---
 ast/walk.go      |  2 ++
 ast/walk_test.go | 57 +++++++++++++++++++++++++++++-------------------
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/ast/walk.go b/ast/walk.go
index 7580a82b..521445b8 100644
--- a/ast/walk.go
+++ b/ast/walk.go
@@ -93,6 +93,7 @@ func Walk(v Visitor, n Node) {
 	case *DotExpression:
 		if n != nil {
 			Walk(v, n.Left)
+			Walk(v, n.Identifier)
 		}
 	case *EmptyExpression:
 	case *EmptyStatement:
@@ -134,6 +135,7 @@ func Walk(v Visitor, n Node) {
 		}
 	case *LabelledStatement:
 		if n != nil {
+			Walk(v, n.Label)
 			Walk(v, n.Statement)
 		}
 	case *NewExpression:
diff --git a/ast/walk_test.go b/ast/walk_test.go
index 14440998..ad0e2f21 100644
--- a/ast/walk_test.go
+++ b/ast/walk_test.go
@@ -1,18 +1,20 @@
 package ast_test
 
 import (
-	"log"
 	"testing"
 
 	"github.com/robertkrimen/otto/ast"
 	"github.com/robertkrimen/otto/file"
 	"github.com/robertkrimen/otto/parser"
+	"github.com/stretchr/testify/require"
 )
 
 type walker struct {
-	stack  []ast.Node
-	source string
-	shift  file.Idx
+	stack     []ast.Node
+	source    string
+	shift     file.Idx
+	seen      map[ast.Node]struct{}
+	duplicate int
 }
 
 // push and pop below are to prove the symmetry of Enter/Exit calls
@@ -38,18 +40,25 @@ func (w *walker) pop(n ast.Node) {
 
 func (w *walker) Enter(n ast.Node) ast.Visitor {
 	w.push(n)
+	if _, ok := w.seen[n]; ok {
+		// Skip items we've already seen which occurs due to declarations.
+		w.duplicate++
+		return w
+	}
+
+	w.seen[n] = struct{}{}
 
 	if id, ok := n.(*ast.Identifier); ok && id != nil {
 		idx := n.Idx0() + w.shift - 1
-		s := w.source[:idx] + "new_" + w.source[idx:]
+		s := w.source[:idx] + "IDENT_" + w.source[idx:]
 		w.source = s
-		w.shift += 4
+		w.shift += 6
 	}
 	if v, ok := n.(*ast.VariableExpression); ok && v != nil {
 		idx := n.Idx0() + w.shift - 1
-		s := w.source[:idx] + "varnew_" + w.source[idx:]
+		s := w.source[:idx] + "VAR_" + w.source[idx:]
 		w.source = s
-		w.shift += 7
+		w.shift += 4
 	}
 
 	return w
@@ -60,23 +69,27 @@ func (w *walker) Exit(n ast.Node) {
 }
 
 func TestVisitorRewrite(t *testing.T) {
-	source := `var b = function() {test(); try {} catch(e) {} var test = "test(); var test = 1"} // test`
+	source := `var b = function() {
+		test();
+		try {} catch(e) {}
+		var test = "test(); var test = 1"
+	} // test`
 	program, err := parser.ParseFile(nil, "", source, 0)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	w := &walker{source: source}
+	require.NoError(t, err)
 
+	w := &walker{
+		source: source,
+		seen:   make(map[ast.Node]struct{}),
+	}
 	ast.Walk(w, program)
 
-	xformed := `var varnew_b = function() {new_test(); try {} catch(new_e) {} var varnew_test = "test(); var test = 1"} // test`
+	xformed := `var VAR_b = function() {
+		IDENT_test();
+		try {} catch(IDENT_e) {}
+		var VAR_test = "test(); var test = 1"
+	} // test`
 
-	if w.source != xformed {
-		t.Errorf("source is `%s` not `%s`", w.source, xformed)
-	}
-
-	if len(w.stack) != 0 {
-		t.Errorf("stack should be empty, but is length: %d", len(w.stack))
-	}
+	require.Equal(t, xformed, w.source)
+	require.Len(t, w.stack, 0)
+	require.Equal(t, w.duplicate, 0)
 }