github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/file/fsnode/iterators.go (about)

     1  package fsnode
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"os"
     7  )
     8  
     9  type sliceIterator struct {
    10  	remaining []T
    11  	closed    bool
    12  }
    13  
    14  var _ Iterator = (*sliceIterator)(nil)
    15  
    16  // NewIterator returns an iterator that yields the given nodes.
    17  func NewIterator(nodes ...T) Iterator {
    18  	// Copy the slice because we'll mutate to nil later.
    19  	nodes = append([]T{}, nodes...)
    20  	return &sliceIterator{remaining: nodes}
    21  }
    22  
    23  func (it *sliceIterator) Next(ctx context.Context) (T, error) {
    24  	if it.closed {
    25  		return nil, os.ErrClosed
    26  	}
    27  	if len(it.remaining) == 0 {
    28  		return nil, io.EOF
    29  	}
    30  	next := it.remaining[0]
    31  	it.remaining[0] = nil // TODO: Is this necessary to allow GC?
    32  	it.remaining = it.remaining[1:]
    33  	return next, nil
    34  }
    35  
    36  func (it *sliceIterator) Close(context.Context) error {
    37  	it.closed = true
    38  	it.remaining = nil
    39  	return nil
    40  }
    41  
    42  type lazyIterator struct {
    43  	make     func(context.Context) ([]T, error)
    44  	fetched  bool
    45  	delegate Iterator
    46  }
    47  
    48  var _ Iterator = (*lazyIterator)(nil)
    49  
    50  // NewLazyIterator uses the given make function upon the first call to Next to
    51  // make the nodes that it yields.
    52  func NewLazyIterator(make func(context.Context) ([]T, error)) Iterator {
    53  	return &lazyIterator{make: make}
    54  }
    55  
    56  func (it *lazyIterator) Next(ctx context.Context) (T, error) {
    57  	if err := it.ensureFetched(ctx); err != nil {
    58  		return nil, err
    59  	}
    60  	return it.delegate.Next(ctx)
    61  }
    62  
    63  func (it *lazyIterator) Close(ctx context.Context) error {
    64  	if it.delegate == nil {
    65  		return nil
    66  	}
    67  	err := it.delegate.Close(ctx)
    68  	it.delegate = nil
    69  	return err
    70  }
    71  
    72  func (it *lazyIterator) ensureFetched(ctx context.Context) error {
    73  	if it.fetched {
    74  		if it.delegate == nil {
    75  			return os.ErrClosed
    76  		}
    77  		return nil
    78  	}
    79  	nodes, err := it.make(ctx)
    80  	if err != nil {
    81  		return err
    82  	}
    83  	it.delegate = NewIterator(nodes...)
    84  	it.fetched = true
    85  	return nil
    86  }
    87  
    88  type concatIterator struct {
    89  	iters  []Iterator
    90  	closed bool
    91  }
    92  
    93  var _ Iterator = (*concatIterator)(nil)
    94  
    95  // NewConcatIterator returns the elements of the given iterators in order, reading each until EOF.
    96  // Manages calling Close on constituents (as it goes along and upon its own Close).
    97  func NewConcatIterator(iterators ...Iterator) Iterator {
    98  	return &concatIterator{iters: append([]Iterator{}, iterators...)}
    99  }
   100  
   101  func (it *concatIterator) Next(ctx context.Context) (T, error) {
   102  	if it.closed {
   103  		return nil, os.ErrClosed
   104  	}
   105  	for {
   106  		if len(it.iters) == 0 {
   107  			return nil, io.EOF
   108  		}
   109  		next, err := it.iters[0].Next(ctx)
   110  		if err == io.EOF {
   111  			err = nil
   112  			if next != nil {
   113  				err = iteratorEOFError(it.iters[0])
   114  			}
   115  			if closeErr := it.iters[0].Close(ctx); closeErr != nil && err == nil {
   116  				err = closeErr
   117  			}
   118  			it.iters[0] = nil // TODO: Is this actually necessary to allow GC?
   119  			it.iters = it.iters[1:]
   120  			if err != nil {
   121  				return nil, err
   122  			}
   123  			continue
   124  		}
   125  		return next, err
   126  	}
   127  }
   128  
   129  // Close attempts to Close remaining constituent iterators. Returns the first constituent Close
   130  // error (but attempts to close the rest anyway).
   131  func (it *concatIterator) Close(ctx context.Context) error {
   132  	it.closed = true
   133  	var err error
   134  	for _, iter := range it.iters {
   135  		if err2 := iter.Close(ctx); err2 != nil && err == nil {
   136  			err = err2
   137  		}
   138  	}
   139  	it.iters = nil
   140  	return err
   141  }
   142  
   143  type mapIterator struct {
   144  	iter Iterator
   145  	fn   func(context.Context, T) (T, error)
   146  }
   147  
   148  // MapIterator returns an Iterator that applies fn to each T yielded by iter.
   149  func MapIterator(iter Iterator, fn func(context.Context, T) (T, error)) Iterator {
   150  	return mapIterator{iter, fn}
   151  }
   152  func (it mapIterator) Next(ctx context.Context) (T, error) {
   153  	if it.fn == nil {
   154  		return nil, os.ErrClosed
   155  	}
   156  	node, err := it.iter.Next(ctx)
   157  	if err == nil && it.fn != nil {
   158  		node, err = it.fn(ctx, node)
   159  	}
   160  	return node, err
   161  }
   162  func (it mapIterator) Close(ctx context.Context) error {
   163  	it.fn = nil
   164  	return it.iter.Close(ctx)
   165  }