github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/geo/iterator.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package geo
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroachdb-parser/pkg/sql/pgwire/pgcode"
    15  	"github.com/cockroachdb/cockroachdb-parser/pkg/sql/pgwire/pgerror"
    16  	"github.com/cockroachdb/errors"
    17  	"github.com/twpayne/go-geom"
    18  )
    19  
    20  // GeomTIterator decomposes geom.T  objects into individual components
    21  // (i.e. either a POINT, LINESTRING or POLYGON) and presents them as an
    22  // iterator. It automatically decomposes MULTI-* and GEOMETRYCOLLECTION
    23  // objects. This prevents an allocation compared to decomposing the objects
    24  // into a geom.T array.
    25  type GeomTIterator struct {
    26  	g             geom.T
    27  	emptyBehavior EmptyBehavior
    28  	// idx is the index into the MULTI-* or GEOMETRYCOLLECTION item.
    29  	idx int
    30  	// subIt is the iterator inside a GeometryCollection.
    31  	// Note GeometryCollections can be nested.
    32  	subIt *GeomTIterator
    33  }
    34  
    35  // NewGeomTIterator returns a new GeomTIterator.
    36  func NewGeomTIterator(g geom.T, emptyBehavior EmptyBehavior) GeomTIterator {
    37  	return GeomTIterator{g: g, emptyBehavior: emptyBehavior}
    38  }
    39  
    40  // Next returns the next geom.T object, a bool as to whether there is an
    41  // entry and an error if any.
    42  func (it *GeomTIterator) Next() (geom.T, bool, error) {
    43  	next, hasNext, err := it.next()
    44  	if err != nil || !hasNext {
    45  		return nil, hasNext, err
    46  	}
    47  	for {
    48  		if !next.Empty() {
    49  			return next, hasNext, nil
    50  		}
    51  		switch it.emptyBehavior {
    52  		case EmptyBehaviorOmit:
    53  			next, hasNext, err = it.next()
    54  			if err != nil || !hasNext {
    55  				return nil, hasNext, err
    56  			}
    57  		case EmptyBehaviorError:
    58  			return nil, false, NewEmptyGeometryError()
    59  		default:
    60  			return nil, false, errors.AssertionFailedf("programmer error: unknown behavior: %T", it.emptyBehavior)
    61  		}
    62  	}
    63  }
    64  
    65  // next() is the internal method for Next.
    66  // It handles fetching the next item in the iterator, recursing down structures
    67  // if necessary. It does not check for emptiness.
    68  func (it *GeomTIterator) next() (geom.T, bool, error) {
    69  	switch t := it.g.(type) {
    70  	case *geom.Point, *geom.LineString, *geom.Polygon:
    71  		if it.idx == 1 {
    72  			return nil, false, nil
    73  		}
    74  		it.idx++
    75  		return t, true, nil
    76  	case *geom.MultiPoint:
    77  		if it.idx == t.NumPoints() {
    78  			return nil, false, nil
    79  		}
    80  		p := t.Point(it.idx)
    81  		it.idx++
    82  		return p, true, nil
    83  	case *geom.MultiLineString:
    84  		if it.idx == t.NumLineStrings() {
    85  			return nil, false, nil
    86  		}
    87  		p := t.LineString(it.idx)
    88  		it.idx++
    89  		return p, true, nil
    90  	case *geom.MultiPolygon:
    91  		if it.idx == t.NumPolygons() {
    92  			return nil, false, nil
    93  		}
    94  		p := t.Polygon(it.idx)
    95  		it.idx++
    96  		return p, true, nil
    97  	case *geom.GeometryCollection:
    98  		for {
    99  			if it.idx == t.NumGeoms() {
   100  				return nil, false, nil
   101  			}
   102  			if it.subIt == nil {
   103  				it.subIt = &GeomTIterator{g: t.Geom(it.idx), emptyBehavior: it.emptyBehavior}
   104  			}
   105  			ret, next, err := it.subIt.next()
   106  			if err != nil {
   107  				return nil, false, err
   108  			}
   109  			if next {
   110  				return ret, next, nil
   111  			}
   112  			// Reset and move to the next item in the collection.
   113  			it.idx++
   114  			it.subIt = nil
   115  		}
   116  	default:
   117  		return nil, false, pgerror.Newf(pgcode.InvalidParameterValue, "unknown type: %T", t)
   118  	}
   119  }
   120  
   121  // Reset resets an iterator back to the first element.
   122  func (it *GeomTIterator) Reset() {
   123  	it.idx = 0
   124  	it.subIt = nil
   125  }