sops/stores/flatten.go

package stores

import (
	"fmt"
	"strconv"
	"strings"
)

const mapSeparator = "__map_"
const listSeparator = "__list_"

// flattenAndMerge flattens the provided value and merges into the
// into map using prefix
func flattenAndMerge(into map[string]interface{}, prefix string, value interface{}) {
	flattenedValue := flattenValue(value)
	if flattenedValue, ok := flattenedValue.(map[string]interface{}); ok {
		for flatK, flatV := range flattenedValue {
			into[prefix+flatK] = flatV
		}
	} else {
		into[prefix] = value
	}
}

func flattenValue(value interface{}) interface{} {
	var output interface{}
	switch value := value.(type) {
	case map[string]interface{}:
		newMap := make(map[string]interface{})
		for k, v := range value {
			flattenAndMerge(newMap, mapSeparator+k, v)
		}
		output = newMap
	case []interface{}:
		newMap := make(map[string]interface{})
		for i, v := range value {
			flattenAndMerge(newMap, listSeparator+fmt.Sprintf("%d", i), v)
		}
		output = newMap
	default:
		output = value
	}
	return output
}

// Flatten flattens a map with potentially nested maps into a flat
// map. Only string keys are allowed on both the top-level map and
// child maps.
func Flatten(in map[string]interface{}) map[string]interface{} {
	newMap := make(map[string]interface{})
	for k, v := range in {
		if flat, ok := flattenValue(v).(map[string]interface{}); ok {
			for flatK, flatV := range flat {
				newMap[k+flatK] = flatV
			}
		} else {
			newMap[k] = v
		}
	}
	return newMap
}

type token interface{}

type mapToken struct {
	key string
}

type listToken struct {
	position int
}

// tokenize converts a path generated by Flatten to be used as a key
// in the flattened map, and converts it to a slice of tokens
func tokenize(path string) []token {
	const (
		StateNormal = 0
		StateMap    = iota
		StateList   = iota
	)
	var tokens []token
	state := StateNormal
	lastTokenEnd := 0
	i := 0
	finishPrevToken := func() {
		var t token
		switch state {
		case StateNormal:
			t = mapToken{path[lastTokenEnd:i]}
		case StateMap:
			t = mapToken{path[lastTokenEnd+len(mapSeparator) : i]}
		case StateList:
			pos, _ := strconv.Atoi(path[lastTokenEnd+len(listSeparator) : i])
			t = listToken{pos}
		}
		lastTokenEnd = i
		tokens = append(tokens, t)
	}
	for i < len(path) {
		if strings.HasPrefix(path[i:], mapSeparator) {
			finishPrevToken()
			state = StateMap
			i += len(mapSeparator)
		} else if strings.HasPrefix(path[i:], listSeparator) {
			finishPrevToken()
			state = StateList
			i += len(listSeparator)
		} else {
			i++
		}
	}
	finishPrevToken()
	return tokens
}

// unflatten takes the currentNode, currentToken, nextToken and value
// and populates currentNode such that currentToken can be considered
// processed. It inspects nextToken to decide what type to allocate
// and assign under currentNode.
func unflatten(currentNode interface{}, currentToken, nextToken token, value interface{}) interface{} {
	switch currentToken := currentToken.(type) {
	case mapToken:
		currentNode := currentNode.(map[string]interface{})
		switch nextToken := nextToken.(type) {
		case mapToken:
			if _, ok := currentNode[currentToken.key]; !ok {
				currentNode[currentToken.key] = make(map[string]interface{})
			}
			next := currentNode[currentToken.key].(map[string]interface{})
			return next
		case listToken:
			if _, ok := currentNode[currentToken.key]; !ok {
				currentNode[currentToken.key] = make([]interface{}, nextToken.position+1)
			}
			next := currentNode[currentToken.key].([]interface{})
			if nextToken.position >= len(next) {
				// Grow the slice and reassign it
				newNext := make([]interface{}, nextToken.position+1)
				copy(newNext, next)
				next = newNext
				currentNode[currentToken.key] = next
			}
			return next
		default:
			currentNode[currentToken.key] = value
		}
	case listToken:
		currentNode := currentNode.([]interface{})
		switch nextToken := nextToken.(type) {
		case mapToken:
			if currentNode[currentToken.position] == nil {
				currentNode[currentToken.position] = make(map[string]interface{})
			}
			next := currentNode[currentToken.position].(map[string]interface{})
			return next
		case listToken:
			if currentNode[currentToken.position] == nil {
				currentNode[currentToken.position] = make([]interface{}, nextToken.position+1)
			}
			next := currentNode[currentToken.position].([]interface{})
			if nextToken.position >= len(next) {
				// Grow the slice and reassign it
				newNext := make([]interface{}, nextToken.position+1)
				copy(newNext, next)
				next = newNext
				currentNode[currentToken.position] = next
			}
			return next
		default:
			currentNode[currentToken.position] = value
		}
	}
	return nil
}

// Unflatten unflattens a map flattened by Flatten
func Unflatten(in map[string]interface{}) map[string]interface{} {
	newMap := make(map[string]interface{})
	for k, v := range in {
		var current interface{} = newMap
		tokens := append(tokenize(k), nil)
		for i := 0; i < len(tokens)-1; i++ {
			current = unflatten(current, tokens[i], tokens[i+1], v)
		}
	}
	return newMap
}