Cache reader usage improvement

Author: ssratkevich

cached_reader.go

@@ -5,45 +5,48 @@ import (
 )
 
 type cachedReader struct {
-	buffer *bufio.Reader
-	cache []byte
-	cacheCap int
-	cacheLen int
+	buffer  *bufio.Reader
+	cache   []byte
 	caching bool
 }
 
 func newCachedReader(r *bufio.Reader) *cachedReader {
 	return &cachedReader{
-		buffer:   r,
-		cache:    make([]byte, 4096),
-		cacheCap: 4096,
-		cacheLen: 0,
-		caching:  false,
+		buffer:  r,
+		cache:   make([]byte, 0, 4096),
+		caching: false,
 	}
 }
 
 func (c *cachedReader) StartCaching() {
-	c.cacheLen = 0
+	c.cache = c.cache[:0]
 	c.caching = true
 }
 
-func (c *cachedReader) ReadByte() (byte, error) {
-	if !c.caching {
-		return c.buffer.ReadByte()
-	}
-	b, err := c.buffer.ReadByte()
+func (c *cachedReader) ReadByte() (b byte, err error) {
+	b, err = c.buffer.ReadByte()
 	if err != nil {
-		return b, err
+		return
 	}
-	if c.cacheLen < c.cacheCap {
-		c.cache[c.cacheLen] = b
-		c.cacheLen++
+	if c.caching {
+		c.cacheByte(b)
 	}
-	return b, err
+	return
 }
 
 func (c *cachedReader) Cache() []byte {
-	return c.cache[:c.cacheLen]
+	return c.cache
+}
+
+func (c *cachedReader) CacheWithLimit(n int) []byte {
+	if n < 1 {
+		return nil
+	}
+	l := len(c.cache)
+	if n > l {
+		n = l
+	}
+	return c.cache[:n]
 }
 
 func (c *cachedReader) StopCaching() {
@@ -55,15 +58,22 @@ func (c *cachedReader) Read(p []byte) (int, error) {
 	if err != nil {
 		return n, err
 	}
-	if c.caching && c.cacheLen < c.cacheCap {
+	if c.caching {
 		for i := 0; i < n; i++ {
-			c.cache[c.cacheLen] = p[i]
-			c.cacheLen++
-			if c.cacheLen >= c.cacheCap {
+			if !c.cacheByte(p[i]) {
 				break
 			}
 		}
 	}
 	return n, err
 }
 
+func (c *cachedReader) cacheByte(b byte) bool {
+	n := len(c.cache)
+	if n == cap(c.cache) {
+		return false
+	}
+	c.cache = c.cache[:n+1]
+	c.cache[n] = b
+	return true
+}

cached_reader_test.go

@@ -39,4 +39,19 @@ func TestCaching(t *testing.T) {
 	if !bytes.Equal(cached, []byte("BCDEF")) {
 		t.Fatalf("Incorrect cached buffer value")
 	}
+
+	cached = cachedReader.CacheWithLimit(-1)
+	if cached != nil {
+		t.Fatalf("Incorrect cached buffer value")
+	}
+
+	cached = cachedReader.CacheWithLimit(3)
+	if !bytes.Equal(cached, []byte("BCD")) {
+		t.Fatalf("Incorrect cached buffer value")
+	}
+
+	cached = cachedReader.CacheWithLimit(1000)
+	if !bytes.Equal(cached, []byte("BCDEF")) {
+		t.Fatalf("Incorrect cached buffer value")
+	}
 }

parse.go

@@ -2,6 +2,7 @@ package xmlquery
 
 import (
 	"bufio"
+	"bytes"
 	"encoding/xml"
 	"fmt"
 	"io"
@@ -184,7 +185,7 @@ func (p *parser) parse() (*Node, error) {
 
 			if node.NamespaceURI != "" {
 				if v, ok := p.space2prefix[node.NamespaceURI]; ok {
-					cached := string(p.reader.Cache())
+					cached := string(p.reader.CacheWithLimit(len(v.name) + len(node.Data) + 2))
 					if strings.HasPrefix(cached, fmt.Sprintf("%s:%s", v.name, node.Data)) || strings.HasPrefix(cached, fmt.Sprintf("<%s:%s", v.name, node.Data)) {
 						node.Prefix = v.name
 					}
@@ -244,9 +245,9 @@ func (p *parser) parse() (*Node, error) {
 			}
 		case xml.CharData:
 			// First, normalize the cache...
-			cached := strings.ToUpper(string(p.reader.Cache()))
+			cached := bytes.ToUpper(p.reader.CacheWithLimit(9))
 			nodeType := TextNode
-			if strings.HasPrefix(cached, "<![CDATA[") || strings.HasPrefix(cached, "![CDATA[") {
+			if bytes.HasPrefix(cached, []byte("<![CDATA[")) || bytes.HasPrefix(cached, []byte("![CDATA[")) {
 				nodeType = CharDataNode
 			}
 			node := &Node{Type: nodeType, Data: string(tok), level: p.level}