github.com/chipaca/snappy@v0.0.0-20210104084008-1f06296fe8ad/strutil/map.go (about) 1 // -*- Mode: Go; indent-tabs-mode: t -*- 2 3 /* 4 * Copyright (C) 2017 Canonical Ltd 5 * 6 * This program is free software: you can redistribute it and/or modify 7 * it under the terms of the GNU General Public License version 3 as 8 * published by the Free Software Foundation. 9 * 10 * This program is distributed in the hope that it will be useful, 11 * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 * GNU General Public License for more details. 14 * 15 * You should have received a copy of the GNU General Public License 16 * along with this program. If not, see <http://www.gnu.org/licenses/>. 17 * 18 */ 19 20 package strutil 21 22 import ( 23 "fmt" 24 25 "gopkg.in/yaml.v2" 26 ) 27 28 // OrderedMap is a map of strings to strings that preserves the 29 // insert order when calling "Keys()". 30 // 31 // Heavily based on the spread.Environment code (thanks for that!) 32 type OrderedMap struct { 33 keys []string 34 vals map[string]string 35 } 36 37 // NewOrderedMap creates a new ordered map initialized with the 38 // given pairs of strings. 39 func NewOrderedMap(pairs ...string) *OrderedMap { 40 o := &OrderedMap{vals: make(map[string]string), 41 keys: make([]string, len(pairs)/2), 42 } 43 for i := 0; i+1 < len(pairs); i += 2 { 44 o.vals[pairs[i]] = pairs[i+1] 45 o.keys[i/2] = pairs[i] 46 } 47 return o 48 } 49 50 // Keys returns a list of keys in the map sorted by insertion order 51 func (o *OrderedMap) Keys() []string { 52 return append([]string(nil), o.keys...) 53 } 54 55 // Get returns the value for the given key 56 func (o *OrderedMap) Get(k string) string { 57 return o.vals[k] 58 } 59 60 // Del removes the given key from the data structure 61 func (o *OrderedMap) Del(key string) { 62 l := len(o.vals) 63 delete(o.vals, key) 64 if len(o.vals) != l { 65 for i, k := range o.keys { 66 if k == key { 67 copy(o.keys[i:], o.keys[i+1:]) 68 o.keys = o.keys[:len(o.keys)-1] 69 } 70 } 71 } 72 } 73 74 // Set adds the given key, value to the map. If the key already 75 // exists it is removed and the new value is put on the end. 76 func (o *OrderedMap) Set(k, v string) { 77 o.Del(k) 78 o.keys = append(o.keys, k) 79 o.vals[k] = v 80 } 81 82 // Copy makes a copy of the map 83 func (o *OrderedMap) Copy() *OrderedMap { 84 copy := &OrderedMap{} 85 copy.keys = append([]string(nil), o.keys...) 86 copy.vals = make(map[string]string) 87 for k, v := range o.vals { 88 copy.vals[k] = v 89 } 90 return copy 91 } 92 93 // UnmarshalYAML unmarshals a yaml string map and preserves the order 94 func (o *OrderedMap) UnmarshalYAML(u func(interface{}) error) error { 95 var vals map[string]string 96 if err := u(&vals); err != nil { 97 return err 98 } 99 100 var seen = make(map[string]bool) 101 var keys = make([]string, len(vals)) 102 var order yaml.MapSlice 103 if err := u(&order); err != nil { 104 return err 105 } 106 for i, item := range order { 107 k, ok := item.Key.(string) 108 _, good := vals[k] 109 if !ok || !good { 110 return fmt.Errorf("cannot read %q", item.Key) 111 } 112 if seen[k] { 113 return fmt.Errorf("found duplicate key %q", k) 114 } 115 seen[k] = true 116 keys[i] = k 117 } 118 o.keys = keys 119 o.vals = vals 120 return nil 121 }