arrayiter.go

Documentation: github.com/lestrrat-go/iter/arrayiter

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