arrayiter.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  1. package arrayiter
  2. import (
  3. "context"
  4. "reflect"
  5. "sync"
  6. "github.com/pkg/errors"
  7. )
  8. func Iterate(ctx context.Context, a interface{}) (Iterator, error) {
  9. arv := reflect.ValueOf(a)
  10. switch arv.Kind() {
  11. case reflect.Array, reflect.Slice:
  12. default:
  13. return nil, errors.Errorf(`argument must be an array/slice (%s)`, arv.Type())
  14. }
  15. ch := make(chan *Pair)
  16. go func(ctx context.Context, ch chan *Pair, arv reflect.Value) {
  17. defer close(ch)
  18. for i := 0; i < arv.Len(); i++ {
  19. value := arv.Index(i)
  20. pair := &Pair{
  21. Index: i,
  22. Value: value.Interface(),
  23. }
  24. select {
  25. case <-ctx.Done():
  26. return
  27. case ch <- pair:
  28. }
  29. }
  30. }(ctx, ch, arv)
  31. return New(ch), nil
  32. }
  33. // Source represents a array that knows how to create an iterator
  34. type Source interface {
  35. Iterate(context.Context) Iterator
  36. }
  37. // Pair represents a single pair of key and value from a array
  38. type Pair struct {
  39. Index int
  40. Value interface{}
  41. }
  42. // Iterator iterates through keys and values of a array
  43. type Iterator interface {
  44. Next(context.Context) bool
  45. Pair() *Pair
  46. }
  47. type iter struct {
  48. ch chan *Pair
  49. mu sync.RWMutex
  50. next *Pair
  51. }
  52. // Visitor represents an object that handles each pair in a array
  53. type Visitor interface {
  54. Visit(int, interface{}) error
  55. }
  56. // VisitorFunc is a type of Visitor based on a function
  57. type VisitorFunc func(int, interface{}) error
  58. func (fn VisitorFunc) Visit(s int, v interface{}) error {
  59. return fn(s, v)
  60. }
  61. func New(ch chan *Pair) Iterator {
  62. return &iter{
  63. ch: ch,
  64. }
  65. }
  66. // Next returns true if there are more items to read from the iterator
  67. func (i *iter) Next(ctx context.Context) bool {
  68. i.mu.RLock()
  69. if i.ch == nil {
  70. i.mu.RUnlock()
  71. return false
  72. }
  73. i.mu.RUnlock()
  74. i.mu.Lock()
  75. defer i.mu.Unlock()
  76. select {
  77. case <-ctx.Done():
  78. i.ch = nil
  79. return false
  80. case v, ok := <-i.ch:
  81. if !ok {
  82. i.ch = nil
  83. return false
  84. }
  85. i.next = v
  86. return true
  87. }
  88. //nolint:govet
  89. return false // never reached
  90. }
  91. // Pair returns the currently buffered Pair. Calling Next() will reset its value
  92. func (i *iter) Pair() *Pair {
  93. i.mu.RLock()
  94. defer i.mu.RUnlock()
  95. return i.next
  96. }
  97. // Walk walks through each element in the array
  98. func Walk(ctx context.Context, s Source, v Visitor) error {
  99. for i := s.Iterate(ctx); i.Next(ctx); {
  100. pair := i.Pair()
  101. if err := v.Visit(pair.Index, pair.Value); err != nil {
  102. return errors.Wrapf(err, `failed to visit index %d`, pair.Index)
  103. }
  104. }
  105. return nil
  106. }
  107. func AsArray(ctx context.Context, s interface{}, v interface{}) error {
  108. var iter Iterator
  109. switch reflect.ValueOf(s).Kind() {
  110. case reflect.Array, reflect.Slice:
  111. x, err := Iterate(ctx, s)
  112. if err != nil {
  113. return errors.Wrap(err, `failed to iterate over array/slice type`)
  114. }
  115. iter = x
  116. default:
  117. ssrc, ok := s.(Source)
  118. if !ok {
  119. return errors.Errorf(`cannot iterate over %T: not a arrayiter.Source type`, s)
  120. }
  121. iter = ssrc.Iterate(ctx)
  122. }
  123. dst := reflect.ValueOf(v)
  124. // dst MUST be a pointer to a array type
  125. if kind := dst.Kind(); kind != reflect.Ptr {
  126. return errors.Errorf(`dst must be a pointer to a array (%s)`, dst.Type())
  127. }
  128. dst = dst.Elem()
  129. switch dst.Kind() {
  130. case reflect.Array, reflect.Slice:
  131. default:
  132. return errors.Errorf(`dst must be a pointer to an array or slice (%s)`, dst.Type())
  133. }
  134. var pairs []*Pair
  135. for iter.Next(ctx) {
  136. pair := iter.Pair()
  137. pairs = append(pairs, pair)
  138. }
  139. switch dst.Kind() {
  140. case reflect.Array:
  141. if len(pairs) < dst.Len() {
  142. return errors.Errorf(`dst array does not have enough space for elements (%d, want %d)`, dst.Len(), len(pairs))
  143. }
  144. case reflect.Slice:
  145. if dst.IsNil() {
  146. dst.Set(reflect.MakeSlice(dst.Type(), len(pairs), len(pairs)))
  147. }
  148. }
  149. // dst must be assignable
  150. if !dst.CanSet() {
  151. return errors.New(`dst is not writeable`)
  152. }
  153. elemtyp := dst.Type().Elem()
  154. for _, pair := range pairs {
  155. rvvalue := reflect.ValueOf(pair.Value)
  156. if !rvvalue.Type().AssignableTo(elemtyp) {
  157. return errors.Errorf(`cannot assign key of type %s to map key of type %s`, rvvalue.Type(), elemtyp)
  158. }
  159. dst.Index(pair.Index).Set(rvvalue)
  160. }
  161. return nil
  162. }