Refactored the Encode function

Author: udan-jayanith

serializer.go

@@ -9,21 +9,13 @@ import (
 	"golang.org/x/net/html"
 )
 
-func wrapAttributeValue(value string) string {
-	if isDigit(value) {
-		return value
-	}
-
-	return `"` + strings.ReplaceAll(value, `"`, """) + `"`
-}
-
 func encodeListAttributes(node *Node) string {
 	w := strings.Builder{}
 	node.IterateAttributes(func(attribute, value string) {
-		if strings.TrimSpace(attribute) == "" {
+		if strings.TrimSpace(value) == "" {
 			w.Write(fmt.Appendf(nil, " %s", attribute))
 		} else {
-			w.Write(fmt.Appendf(nil, " %s=%s", attribute, wrapAttributeValue(value)))
+			w.Write(fmt.Appendf(nil, " %s=%s", attribute, `"`+html.EscapeString(value)+`"`))
 		}
 
 	})
@@ -32,43 +24,55 @@ func encodeListAttributes(node *Node) string {
 
 // Encode writes to w encoding of the node tree from rootNode.
 func Encode(w io.Writer, rootNode *Node) {
+	if rootNode == nil {
+		return
+	}
+
 	type stackFrame struct {
-		node      *Node
-		openedTag bool
+		node         *Node
+		isClosingTag bool
 	}
 
 	stack := linkedliststack.New()
-	stack.Push(stackFrame{node: rootNode, openedTag: false})
+	stack.Push(stackFrame{
+		node: rootNode,
+	})
 
 	for stack.Size() > 0 {
-		t, _ := stack.Pop()
-		top := t.(stackFrame)
-		current := top.node
+		v, _ := stack.Pop()
+		currentStackFrame := v.(stackFrame)
 
-		if current == nil {
+		if currentStackFrame.isClosingTag {
+			fmt.Fprintf(w, "</%s>", currentStackFrame.node.GetTagName())
 			continue
+		} else if currentStackFrame.node.IsTextNode() {
+			fmt.Fprint(w, html.EscapeString(currentStackFrame.node.GetText()))
+		} else {
+			fmt.Fprintf(w, "<%s%s>", func() string {
+				tagName := currentStackFrame.node.GetTagName()
+				tagNameUpperCased := strings.ToUpper(tagName)
+				if tagNameUpperCased == DOCTYPEDTD {
+					tagName = tagNameUpperCased
+				}
+				return tagName
+			}(), encodeListAttributes(currentStackFrame.node))
 		}
 
-		tagName := current.GetTagName()
-		if tagName == "" {
-			w.Write([]byte(html.EscapeString(current.GetText())))
-		} else if IsVoidTag(tagName) {
-			fmt.Fprintf(w, "<%s%s>", tagName, encodeListAttributes(current))
-			if current.GetNextNode() != nil {
-				stack.Push(stackFrame{node: current.GetNextNode(), openedTag: false})
-			}
-		} else if !top.openedTag {
-			fmt.Fprintf(w, "<%s%s>", tagName, encodeListAttributes(current))
-			stack.Push(stackFrame{node: current, openedTag: true})
-
-			if current.GetChildNode() != nil {
-				stack.Push(stackFrame{node: current.GetChildNode(), openedTag: false})
-			}
-		} else {
-			fmt.Fprintf(w, "</%s>", tagName)
-			if current.GetNextNode() != nil {
-				stack.Push(stackFrame{node: current.GetNextNode(), openedTag: false})
-			}
+		if currentStackFrame.node.GetNextNode() != nil {
+			stack.Push(stackFrame{
+				node: currentStackFrame.node.GetNextNode(),
+			})
+		}
+		if !IsVoidTag(currentStackFrame.node.GetTagName()) && !currentStackFrame.node.IsTextNode(){
+			stack.Push(stackFrame{
+				node:         currentStackFrame.node,
+				isClosingTag: true,
+			})
+		}
+		if currentStackFrame.node.GetChildNode() != nil {
+			stack.Push(stackFrame{
+				node: currentStackFrame.node.GetChildNode(),
+			})
 		}
 	}
 }