[beam] 01/01: [BEAM-3612] Support use of non map derived types.

[lo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lostluck pushed a commit to branch lostluck-gen-1
in repository https://gitbox.apache.org/repos/asf/beam.git

commit d77f13c24cdbc437602685bf88e4886c384d927a
Author: Robert Burke <lo...@users.noreply.github.com>
AuthorDate: Tue Sep 8 09:24:05 2020 -0700

    [BEAM-3612] Support use of non map derived types.
    
    Recursively unwrap slice and array types, in addition to pointer types when extracting base types to register.
---
 sdks/go/pkg/beam/util/starcgenx/starcgenx.go | 34 +++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/sdks/go/pkg/beam/util/starcgenx/starcgenx.go b/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
index 8961761..21deee8 100644
--- a/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
+++ b/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
@@ -328,18 +328,36 @@ func (e *Extractor) extractFromSignature(sig *types.Signature) {
 	e.extractFromTuple(sig.Results())
 }
 
+// extractFromContainer recurses through nested non-map container types to a non-derived
+// element type.
+func (e *Extractor) extractFromContainer(t types.Type) types.Type {
+	// Container types need to be iteratively unwrapped until we're at the base type,
+	// so we can get the import if necessary.
+	for {
+		if s, ok := t.(*types.Slice); ok {
+			t = s.Elem()
+			continue
+		}
+
+		if p, ok := t.(*types.Pointer); ok {
+			t = p.Elem()
+			continue
+		}
+
+		if a, ok := t.(*types.Array); ok {
+			t = a.Elem()
+			continue
+		}
+
+		return t
+	}
+}
+
 func (e *Extractor) extractFromTuple(tuple *types.Tuple) {
 	for i := 0; i < tuple.Len(); i++ {
 		s := tuple.At(i) // *types.Var
 
-		// Pointer types need to be iteratively unwrapped until we're at the base type,
-		// so we can get the import if necessary.
-		t := s.Type()
-		p, ok := t.(*types.Pointer)
-		for ok {
-			t = p.Elem()
-			p, ok = t.(*types.Pointer)
-		}
+		t := e.extractFromContainer(s.Type())
 		// Here's where we ensure we register new imports.
 		if t, ok := t.(*types.Named); ok {
 			if pkg := t.Obj().Pkg(); pkg != nil {