Prevent infinite loop within serialization methods (#287)

Author: atelpis

testutils/go/predefined.go

@@ -2,10 +2,13 @@ package goutils
 
 import (
 	"encoding/json"
+	"errors"
 	"strconv"
 	"strings"
 )
 
+var ErrInfiniteLoop = errors.New("infinite loop detected")
+
 /*
 Much appreciated to EndlessCheng
 Adapted from https://github.com/EndlessCheng/codeforces-go/blob/ae5b312f3f/leetcode/testutil/leetcode.go
@@ -36,14 +39,23 @@ func DeserializeListNode(s string) (*ListNode, error) {
 	return root, nil
 }
 
+// ToString returns a string representation of the linked list.
+// It panics with ErrInfiniteLoop if a cycle is detected.
 func (l *ListNode) ToString() string {
+	seen := make(map[*ListNode]bool, 10)
+
 	sb := &strings.Builder{}
 	sb.WriteByte('[')
 	for ; l != nil; l = l.Next {
 		if sb.Len() > 1 {
 			sb.WriteByte(',')
 		}
 		sb.WriteString(strconv.Itoa(l.Val))
+
+		if seen[l] {
+			panic(ErrInfiniteLoop)
+		}
+		seen[l] = true
 	}
 	sb.WriteByte(']')
 	return sb.String()
@@ -104,13 +116,21 @@ func DeserializeTreeNode(s string) (*TreeNode, error) {
 	return root, nil
 }
 
+// ToString returns a string representation of the binary tree.
+// It panics with ErrInfiniteLoop if a cycle is detected.
 func (t *TreeNode) ToString() string {
 	nodes := []*TreeNode{}
 	queue := []*TreeNode{t}
+	seen := make(map[*TreeNode]bool, 10)
 	for len(queue) > 0 {
 		t, queue = queue[0], queue[1:]
 		nodes = append(nodes, t)
 		if t != nil {
+			if seen[t] {
+				panic(ErrInfiniteLoop)
+			}
+			seen[t] = true
+
 			queue = append(queue, t.Left, t.Right)
 		}
 	}
@@ -164,16 +184,24 @@ func DeserializeNaryTreeNode(s string) (*NaryTreeNode, error) {
 	return root.Children[0], nil
 }
 
+// ToString returns a string representation of the nary tree.
+// It panics with ErrInfiniteLoop if a cycle is detected.
 func (t *NaryTreeNode) ToString() string {
 	nodes := []*NaryTreeNode{}
 	q := []*NaryTreeNode{{Children: []*NaryTreeNode{t}}}
+	seen := make(map[*NaryTreeNode]bool, 10)
 
 	for len(q) > 0 {
 		node := q[0]
 		q = q[1:]
 		nodes = append(nodes, node)
 
 		if node != nil {
+			if seen[node] {
+				panic(ErrInfiniteLoop)
+			}
+			seen[node] = true
+
 			if len(node.Children) > 0 {
 				q = append(q, node.Children...)
 			}

testutils/go/serialize_test.go

@@ -0,0 +1,33 @@
+package goutils
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInfiniteLoopDetect(t *testing.T) {
+
+	linkedList := &ListNode{Val: 1}
+	linkedList.Next = &ListNode{Val: 2, Next: linkedList}
+
+	tree := &TreeNode{Val: 1}
+	tree.Left = &TreeNode{Val: 2, Right: tree}
+
+	naryTree := &NaryTreeNode{Val: 1}
+	naryTree.Children = []*NaryTreeNode{{Val: 2, Children: []*NaryTreeNode{naryTree}}}
+
+	type toStringer interface {
+		ToString() string
+	}
+
+	tests := []toStringer{
+		linkedList,
+		tree,
+		naryTree,
+	}
+
+	for _, tc := range tests {
+		assert.PanicsWithValue(t, ErrInfiniteLoop, func() { tc.ToString() })
+	}
+}