Merge pull request #19 from itchyny/token-based-parsing

Unmarshal using token decoder

Author: iancoleman

orderedmap.go

@@ -3,23 +3,10 @@ package orderedmap
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
 	"sort"
 	"strings"
 )
 
-var NoValueError = errors.New("No value for this key")
-
-type KeyIndex struct {
-	Key   string
-	Index int
-}
-type ByIndex []KeyIndex
-
-func (a ByIndex) Len() int           { return len(a) }
-func (a ByIndex) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
-func (a ByIndex) Less(i, j int) bool { return a[i].Index < a[j].Index }
-
 type Pair struct {
 	key   string
 	value interface{}
@@ -117,201 +104,134 @@ func (o *OrderedMap) UnmarshalJSON(b []byte) error {
 	if o.values == nil {
 		o.values = map[string]interface{}{}
 	}
-	var err error
-	err = mapStringToOrderedMap(string(b), o)
+	err := json.Unmarshal(b, &o.values)
 	if err != nil {
 		return err
 	}
-	return nil
-}
-
-func mapStringToOrderedMap(s string, o *OrderedMap) error {
-	// parse string into map
-	m := map[string]interface{}{}
-	err := json.Unmarshal([]byte(s), &m)
-	if err != nil {
+	dec := json.NewDecoder(bytes.NewReader(b))
+	if _, err = dec.Token(); err != nil { // skip '{'
 		return err
 	}
-	// Get the order of the keys
-	orderedKeys := []KeyIndex{}
-	for k, _ := range m {
-		kEscaped := strings.Replace(k, `"`, `\"`, -1)
-		kQuoted := `"` + kEscaped + `"`
-		// Find how much content exists before this key.
-		// If all content from this key and after is replaced with a close
-		// brace, it should still form a valid json string.
-		sTrimmed := s
-		for len(sTrimmed) > 0 {
-			lastIndex := strings.LastIndex(sTrimmed, kQuoted)
-			if lastIndex == -1 {
-				break
-			}
-			sTrimmed = sTrimmed[0:lastIndex]
-			sTrimmed = strings.TrimSpace(sTrimmed)
-			if len(sTrimmed) > 0 && sTrimmed[len(sTrimmed)-1] == ',' {
-				sTrimmed = sTrimmed[0 : len(sTrimmed)-1]
-			}
-			maybeValidJson := sTrimmed + "}"
-			testMap := map[string]interface{}{}
-			err := json.Unmarshal([]byte(maybeValidJson), &testMap)
-			if err == nil {
-				// record the position of this key in s
-				ki := KeyIndex{
-					Key:   k,
-					Index: len(sTrimmed),
-				}
-				orderedKeys = append(orderedKeys, ki)
-				// shorten the string to get the next key
-				startOfValueIndex := lastIndex + len(kQuoted)
-				valueStr := s[startOfValueIndex : len(s)-1]
-				valueStr = strings.TrimSpace(valueStr)
-				if len(valueStr) > 0 && valueStr[0] == ':' {
-					valueStr = valueStr[1:len(valueStr)]
+	o.keys = make([]string, 0, len(o.values))
+	return decodeOrderedMap(dec, o)
+}
+
+func decodeOrderedMap(dec *json.Decoder, o *OrderedMap) error {
+	hasKey := make(map[string]bool, len(o.values))
+	for {
+		token, err := dec.Token()
+		if err != nil {
+			return err
+		}
+		if delim, ok := token.(json.Delim); ok && delim == '}' {
+			return nil
+		}
+		key := token.(string)
+		if hasKey[key] {
+			// duplicate key
+			for j, k := range o.keys {
+				if k == key {
+					copy(o.keys[j:], o.keys[j+1:])
+					break
 				}
-				valueStr = strings.TrimSpace(valueStr)
-				if valueStr[0] == '{' {
-					// if the value for this key is a map
-					// find end of valueStr by removing everything after last }
-					// until it forms valid json
-					hasValidJson := false
-					i := 1
-					for i < len(valueStr) && !hasValidJson {
-						if valueStr[i] != '}' {
-							i = i + 1
-							continue
-						}
-						subTestMap := map[string]interface{}{}
-						testValue := valueStr[0 : i+1]
-						err = json.Unmarshal([]byte(testValue), &subTestMap)
-						if err == nil {
-							hasValidJson = true
-							valueStr = testValue
-							break
-						}
-						i = i + 1
+			}
+			o.keys[len(o.keys)-1] = key
+		} else {
+			hasKey[key] = true
+			o.keys = append(o.keys, key)
+		}
+
+		token, err = dec.Token()
+		if err != nil {
+			return err
+		}
+		if delim, ok := token.(json.Delim); ok {
+			switch delim {
+			case '{':
+				if values, ok := o.values[key].(map[string]interface{}); ok {
+					newMap := OrderedMap{
+						keys:       make([]string, 0, len(values)),
+						values:     values,
+						escapeHTML: o.escapeHTML,
 					}
-					// convert to orderedmap
-					// this may be recursive it values in the map are also maps
-					if hasValidJson {
-						newMap := New()
-						err := mapStringToOrderedMap(valueStr, newMap)
-						if err != nil {
-							return err
-						}
-						m[k] = *newMap
+					if err = decodeOrderedMap(dec, &newMap); err != nil {
+						return err
 					}
-				} else if valueStr[0] == '[' {
-					// if the value for this key is a slice
-					// find end of valueStr by removing everything after last ]
-					// until it forms valid json
-					hasValidJson := false
-					i := 1
-					for i < len(valueStr) && !hasValidJson {
-						if valueStr[i] != ']' {
-							i = i + 1
-							continue
-						}
-						subTestSlice := []interface{}{}
-						testValue := valueStr[0 : i+1]
-						err = json.Unmarshal([]byte(testValue), &subTestSlice)
-						if err == nil {
-							hasValidJson = true
-							valueStr = testValue
-							break
-						}
-						i = i + 1
+					o.values[key] = newMap
+				} else if oldMap, ok := o.values[key].(OrderedMap); ok {
+					newMap := OrderedMap{
+						keys:       make([]string, 0, len(oldMap.values)),
+						values:     oldMap.values,
+						escapeHTML: o.escapeHTML,
 					}
-					// convert to slice with any map items converted to
-					// orderedmaps
-					// this may be recursive if values in the slice are slices
-					if hasValidJson {
-						newSlice := []interface{}{}
-						err := sliceStringToSliceWithOrderedMaps(valueStr, &newSlice)
-						if err != nil {
-							return err
-						}
-						m[k] = newSlice
+					if err = decodeOrderedMap(dec, &newMap); err != nil {
+						return err
 					}
-				} else {
-					o.Set(k, m[k])
+					o.values[key] = newMap
+				} else if err = decodeOrderedMap(dec, &OrderedMap{}); err != nil {
+					return err
+				}
+			case '[':
+				if values, ok := o.values[key].([]interface{}); ok {
+					if err = decodeSlice(dec, values, o.escapeHTML); err != nil {
+						return err
+					}
+				} else if err = decodeSlice(dec, []interface{}{}, o.escapeHTML); err != nil {
+					return err
 				}
-				break
 			}
 		}
 	}
-	// Sort the keys
-	sort.Sort(ByIndex(orderedKeys))
-	// Convert sorted keys to string slice
-	k := []string{}
-	for _, ki := range orderedKeys {
-		k = append(k, ki.Key)
-	}
-	// Set the OrderedMap values
-	o.values = m
-	o.keys = k
-	return nil
 }
 
-func sliceStringToSliceWithOrderedMaps(valueStr string, newSlice *[]interface{}) error {
-	// if the value for this key is a []interface, convert any map items to an orderedmap.
-	// find end of valueStr by removing everything after last ]
-	// until it forms valid json
-	itemsStr := strings.TrimSpace(valueStr)
-	itemsStr = itemsStr[1 : len(itemsStr)-1]
-	// get next item in the slice
-	itemIndex := 0
-	startItem := 0
-	endItem := 0
-	for endItem <= len(itemsStr) {
-		couldBeItemEnd := false
-		couldBeItemEnd = couldBeItemEnd || endItem == len(itemsStr)
-		couldBeItemEnd = couldBeItemEnd || (endItem < len(itemsStr) && itemsStr[endItem] == ',')
-		if !couldBeItemEnd {
-			endItem = endItem + 1
-			continue
-		}
-		// if this substring compiles to json, it's the next item
-		possibleItemStr := strings.TrimSpace(itemsStr[startItem:endItem])
-		var possibleItem interface{}
-		err := json.Unmarshal([]byte(possibleItemStr), &possibleItem)
+func decodeSlice(dec *json.Decoder, s []interface{}, escapeHTML bool) error {
+	for index := 0; ; index++ {
+		token, err := dec.Token()
 		if err != nil {
-			endItem = endItem + 1
-			continue
+			return err
 		}
-		if possibleItemStr[0] == '{' {
-			// if item is map, convert to orderedmap
-			oo := New()
-			err := mapStringToOrderedMap(possibleItemStr, oo)
-			if err != nil {
-				return err
-			}
-			// add new orderedmap item to new slice
-			slice := *newSlice
-			slice = append(slice, *oo)
-			*newSlice = slice
-		} else if possibleItemStr[0] == '[' {
-			// if item is slice, convert to slice with orderedmaps
-			newItem := []interface{}{}
-			err := sliceStringToSliceWithOrderedMaps(possibleItemStr, &newItem)
-			if err != nil {
-				return err
+		if delim, ok := token.(json.Delim); ok {
+			switch delim {
+			case '{':
+				if index < len(s) {
+					if values, ok := s[index].(map[string]interface{}); ok {
+						newMap := OrderedMap{
+							keys:       make([]string, 0, len(values)),
+							values:     values,
+							escapeHTML: escapeHTML,
+						}
+						if err = decodeOrderedMap(dec, &newMap); err != nil {
+							return err
+						}
+						s[index] = newMap
+					} else if oldMap, ok := s[index].(OrderedMap); ok {
+						newMap := OrderedMap{
+							keys:       make([]string, 0, len(oldMap.values)),
+							values:     oldMap.values,
+							escapeHTML: escapeHTML,
+						}
+						if err = decodeOrderedMap(dec, &newMap); err != nil {
+							return err
+						}
+						s[index] = newMap
+					}
+				} else if err = decodeOrderedMap(dec, &OrderedMap{}); err != nil {
+					return err
+				}
+			case '[':
+				if index < len(s) {
+					values := s[index].([]interface{})
+					if err = decodeSlice(dec, values, escapeHTML); err != nil {
+						return err
+					}
+				} else if err = decodeSlice(dec, []interface{}{}, escapeHTML); err != nil {
+					return err
+				}
+			case ']':
+				return nil
 			}
-			// replace original slice item with new slice
-			slice := *newSlice
-			slice = append(slice, newItem)
-			*newSlice = slice
-		} else {
-			// any non-slice and non-map item, just add json parsed item
-			slice := *newSlice
-			slice = append(slice, possibleItem)
-			*newSlice = slice
 		}
-		// remove this item from itemsStr
-		startItem = endItem + 1
-		endItem = endItem + 1
-		itemIndex = itemIndex + 1
 	}
-	return nil
 }
 
 func (o OrderedMap) MarshalJSON() ([]byte, error) {