data_map.go

package mmdbwriter

import "github.com/maxmind/mmdbwriter/mmdbtype"

type dataMapKey string

type dataMapValue struct {
	data mmdbtype.DataType
	// This exists to reduce allocations by interning the key.
	key dataMapKey
}

// dataMap is used to deduplicate data inserted into the tree to reduce
// memory usage using keys generated by keyWriter.
type dataMap struct {
	data      map[dataMapKey]dataMapValue
	keyWriter *keyWriter
}

const noDataMapKey dataMapKey = ""

func newDataMap() *dataMap {
	return &dataMap{
		data:      map[dataMapKey]dataMapValue{},
		keyWriter: newKeyWriter(),
	}
}

func (dm *dataMap) get(k dataMapKey) mmdbtype.DataType {
	if k == noDataMapKey {
		return nil
	}
	return dm.data[k].data
}

// store stores the value in the dataMap and returns a key for it.
func (dm *dataMap) store(v mmdbtype.DataType) (dataMapKey, error) {
	key, err := dm.keyWriter.key(v)
	if err != nil {
		return "", err
	}

	dmv, ok := dm.data[dataMapKey(key)]
	if !ok {
		dmKey := dataMapKey(key)
		dmv = dataMapValue{
			key:  dmKey,
			data: v,
		}
		dm.data[dmKey] = dmv
	}

	return dmv.key, nil
}