File: mapiter.go

package info (click to toggle)
golang-github-lestrrat-go-iter 1.0.2-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 108 kB
  • sloc: makefile: 2
file content (195 lines) | stat: -rw-r--r-- 4,320 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
package mapiter

import (
	"context"
	"fmt"
	"reflect"
	"sync"
)

// Iterate creates an iterator from arbitrary map types. This is not
// the most efficient tool, but it's the quickest way to create an
// iterator for maps.
// Also, note that you cannot make any assumptions on the order of
// pairs being returned.
func Iterate(ctx context.Context, m interface{}) (Iterator, error) {
	mrv := reflect.ValueOf(m)

	if mrv.Kind() != reflect.Map {
		return nil, fmt.Errorf(`argument must be a map (%s)`, mrv.Type())
	}

	ch := make(chan *Pair)
	go func(ctx context.Context, ch chan *Pair, mrv reflect.Value) {
		defer close(ch)
		for _, key := range mrv.MapKeys() {
			value := mrv.MapIndex(key)
			pair := &Pair{
				Key:   key.Interface(),
				Value: value.Interface(),
			}
			select {
			case <-ctx.Done():
				return
			case ch <- pair:
			}
		}
	}(ctx, ch, mrv)

	return New(ch), nil
}

// Source represents a map that knows how to create an iterator
type Source interface {
	Iterate(context.Context) Iterator
}

// Pair represents a single pair of key and value from a map
type Pair struct {
	Key   interface{}
	Value interface{}
}

// Iterator iterates through keys and values of a map
type Iterator interface {
	Next(context.Context) bool
	Pair() *Pair
}

type iter struct {
	ch   chan *Pair
	mu   sync.RWMutex
	next *Pair
}

// Visitor represents an object that handles each pair in a map
type Visitor interface {
	Visit(interface{}, interface{}) error
}

// VisitorFunc is a type of Visitor based on a function
type VisitorFunc func(interface{}, interface{}) error

func (fn VisitorFunc) Visit(s interface{}, v interface{}) error {
	return fn(s, v)
}

func New(ch chan *Pair) Iterator {
	return &iter{
		ch: ch,
	}
}

// Next returns true if there are more items to read from the iterator
func (i *iter) Next(ctx context.Context) bool {
	i.mu.RLock()
	if i.ch == nil {
		i.mu.RUnlock()
		return false
	}
	i.mu.RUnlock()

	i.mu.Lock()
	defer i.mu.Unlock()
	select {
	case <-ctx.Done():
		i.ch = nil
		return false
	case v, ok := <-i.ch:
		if !ok {
			i.ch = nil
			return false
		}
		i.next = v
		return true
	}

	//nolint:govet
	return false // never reached
}

// Pair returns the currently buffered Pair. Calling Next() will reset its value
func (i *iter) Pair() *Pair {
	i.mu.RLock()
	defer i.mu.RUnlock()
	return i.next
}

// Walk walks through each element in the map
func Walk(ctx context.Context, s Source, v Visitor) error {
	for i := s.Iterate(ctx); i.Next(ctx); {
		pair := i.Pair()
		if err := v.Visit(pair.Key, pair.Value); err != nil {
			return fmt.Errorf(`failed to visit key %s: %w`, pair.Key, err)
		}
	}
	return nil
}

// AsMap returns the values obtained from the source as a map
func AsMap(ctx context.Context, s interface{}, v interface{}) error {
	var iter Iterator
	switch reflect.ValueOf(s).Kind() {
	case reflect.Map:
		x, err := Iterate(ctx, s)
		if err != nil {
			return fmt.Errorf(`failed to iterate over map type: %w`, err)
		}
		iter = x
	default:
		ssrc, ok := s.(Source)
		if !ok {
			return fmt.Errorf(`cannot iterate over %T: not a mapiter.Source type`, s)
		}
		iter = ssrc.Iterate(ctx)
	}

	dst := reflect.ValueOf(v)

	// dst MUST be a pointer to a map type
	if kind := dst.Kind(); kind != reflect.Ptr {
		return fmt.Errorf(`dst must be a pointer to a map (%s)`, dst.Type())
	}

	dst = dst.Elem()
	if dst.Kind() != reflect.Map {
		return fmt.Errorf(`dst must be a pointer to a map (%s)`, dst.Type())
	}

	if dst.IsNil() {
		dst.Set(reflect.MakeMap(dst.Type()))
	}

	// dst must be assignable
	if !dst.CanSet() {
		return fmt.Errorf(`dst is not writeable`)
	}

	keytyp := dst.Type().Key()
	valtyp := dst.Type().Elem()

	for iter.Next(ctx) {
		pair := iter.Pair()

		rvkey := reflect.ValueOf(pair.Key)
		rvvalue := reflect.ValueOf(pair.Value)

		if !rvkey.Type().AssignableTo(keytyp) {
			return fmt.Errorf(`cannot assign key of type %s to map key of type %s`, rvkey.Type(), keytyp)
		}

		switch rvvalue.Kind() {
		// we can only check if we can assign to rvvalue to valtyp if it's non-nil
		case reflect.Invalid:
			rvvalue = reflect.New(valtyp).Elem()
		default:
			if !rvvalue.Type().AssignableTo(valtyp) {
				return fmt.Errorf(`cannot assign value of type %s to map value of type %s`, rvvalue.Type(), valtyp)
			}
		}

		dst.SetMapIndex(rvkey, rvvalue)
	}

	return nil
}