github.com/bugraaydogar/snapd@v0.0.0-20210315170335-8c70bb858939/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  }