From 7c4be53b0731c70d4cc8b97e2c092e126bcc5c70 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 03:31:00 +0000
Subject: [PATCH 1/7] feat: Add identified embeddable structs

---
 internal/codegen/golang/gen.go    |  6 +-
 internal/codegen/golang/result.go | 92 +++++++++++++++++++++++--------
 internal/codegen/golang/struct.go |  1 +
 3 files changed, 73 insertions(+), 26 deletions(-)

diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
index 138a9f41cd..a3f0b4dc96 100644
--- a/internal/codegen/golang/gen.go
+++ b/internal/codegen/golang/gen.go
@@ -243,13 +243,15 @@ const {{.ConstantName}} = {{$.Q}}-- name: {{.MethodName}} {{.Cmd}}
 
 {{if .Arg.EmitStruct}}
 type {{.Arg.Type}} struct { {{- range .Arg.Struct.Fields}}
-  {{.Name}} {{.Type}} {{if or ($.EmitJSONTags) ($.EmitDBTags)}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
+    {{.Name}} {{.Type}} {{if or ($.EmitJSONTags) ($.EmitDBTags)}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
 }
 {{end}}
 
 {{if .Ret.EmitStruct}}
-type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
+type {{.Ret.Type}} struct { {{- range .Ret.Struct.Structs}}
+  {{.Name}}
+  {{- end}} {{- range .Ret.Struct.Fields}}
   {{.Name}} {{.Type}} {{if or ($.EmitJSONTags) ($.EmitDBTags)}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
 }
diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index cbce7d4a2e..0e0d316426 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -103,7 +103,8 @@ func buildStructs(r *compiler.Result, settings config.CombinedSettings) []Struct
 }
 
 type goColumn struct {
-	id int
+	id    int
+	Embed *Struct
 	*compiler.Column
 }
 
@@ -183,37 +184,80 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 				Typ:  goType(r, c, settings),
 			}
 		} else if len(query.Columns) > 1 {
-			var gs *Struct
-			var emit bool
+			var columns []goColumn
 
-			for _, s := range structs {
-				if len(s.Fields) != len(query.Columns) {
-					continue
-				}
-				same := true
-				for i, f := range s.Fields {
-					c := query.Columns[i]
-					sameName := f.Name == StructName(columnName(c, i), settings)
-					sameType := f.Type == goType(r, c, settings)
-					sameTable := sameTableName(c.Table, s.Table, r.Catalog.DefaultSchema)
-					if !sameName || !sameType || !sameTable {
-						same = false
+			for ci := 0; ci < len(query.Columns); {
+				c := query.Columns[ci]
+				var embed *Struct
+
+				// Checks for matching structs.
+				for _, s := range structs {
+					// Ensuring tables match the column selector.
+					if c.Table == nil || c.Table.Name != s.Table.Rel {
+						continue
+					}
+					// If the query doesn't have enough fields, it cannot
+					// fufill the struct.
+					if len(query.Columns) < len(s.Fields) {
+						continue
 					}
+					same := true
+					for fi, f := range s.Fields {
+						fieldOffset := ci + fi
+						// If the location of this field doesn't fit into our columns,
+						// we know the struct can't fit either.
+						if fieldOffset > len(query.Columns)-1 {
+							break
+						}
+						c := query.Columns[fieldOffset]
+						sameName := f.Name == StructName(columnName(c, fieldOffset), settings)
+						sameType := f.Type == goType(r, c, settings)
+						sameTable := sameTableName(c.Table, s.Table, r.Catalog.DefaultSchema)
+						if !sameName || !sameType || !sameTable {
+							same = false
+						}
+					}
+					if same {
+						embed = &s
+						break
+					}
+				}
+
+				// Used to track the amount of columns matched.
+				// A struct could be embedded, and in that case
+				// for performance we want to skip over those
+				// matched columns.
+				colsMatched := 1
+				if embed != nil {
+					colsMatched = len(embed.Fields)
 				}
-				if same {
-					gs = &s
+				for colID := ci; colID < ci+colsMatched; colID++ {
+					columns = append(columns, goColumn{
+						id:     colID,
+						Embed:  embed,
+						Column: query.Columns[colID],
+					})
+				}
+				ci += colsMatched
+			}
+
+			var emit bool
+			var gs *Struct
+			// Check if all columns match a consistent embedded struct.
+			// If they do, we don't need to generate a new struct for the row.
+			for _, c := range columns {
+				if gs == nil {
+					gs = c.Embed
+					continue
+				}
+
+				if gs != c.Embed {
+					gs = nil
 					break
 				}
 			}
 
 			if gs == nil {
-				var columns []goColumn
-				for i, c := range query.Columns {
-					columns = append(columns, goColumn{
-						id:     i,
-						Column: c,
-					})
-				}
 				gs = columnsToStruct(r, gq.MethodName+"Row", columns, settings)
 				emit = true
 			}
diff --git a/internal/codegen/golang/struct.go b/internal/codegen/golang/struct.go
index 74be48e3f3..f2c6c14ab7 100644
--- a/internal/codegen/golang/struct.go
+++ b/internal/codegen/golang/struct.go
@@ -10,6 +10,7 @@ import (
 type Struct struct {
 	Table   core.FQN
 	Name    string
+	Structs []Struct
 	Fields  []Field
 	Comment string
 }

From 404538788147313d9271dfcc2e6b2f03c3b2f42a Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 03:52:14 +0000
Subject: [PATCH 2/7] Embedded structs now fully work

---
 internal/codegen/golang/field.go  |  1 +
 internal/codegen/golang/gen.go    |  5 +++--
 internal/codegen/golang/query.go  |  9 +++++++--
 internal/codegen/golang/result.go | 18 ++++++++++++++++--
 internal/codegen/golang/struct.go | 10 +++++-----
 5 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/internal/codegen/golang/field.go b/internal/codegen/golang/field.go
index e036b0e041..9cce884372 100644
--- a/internal/codegen/golang/field.go
+++ b/internal/codegen/golang/field.go
@@ -10,6 +10,7 @@ import (
 
 type Field struct {
 	Name    string
+	Struct  string
 	Type    string
 	Tags    map[string]string
 	Comment string
diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
index a3f0b4dc96..5884420f92 100644
--- a/internal/codegen/golang/gen.go
+++ b/internal/codegen/golang/gen.go
@@ -249,11 +249,12 @@ type {{.Arg.Type}} struct { {{- range .Arg.Struct.Fields}}
 {{end}}
 
 {{if .Ret.EmitStruct}}
-type {{.Ret.Type}} struct { {{- range .Ret.Struct.Structs}}
+type {{.Ret.Type}} struct { {{- range .Ret.Struct.Embedded}}
   {{.Name}}
-  {{- end}} {{- range .Ret.Struct.Fields}}
+  {{- end}} {{- range .Ret.Struct.Fields}} {{if not .Struct}}
   {{.Name}} {{.Type}} {{if or ($.EmitJSONTags) ($.EmitDBTags)}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
+  {{- end}}
 }
 {{end}}
 
diff --git a/internal/codegen/golang/query.go b/internal/codegen/golang/query.go
index afade092d4..42f5229d2e 100644
--- a/internal/codegen/golang/query.go
+++ b/internal/codegen/golang/query.go
@@ -1,6 +1,7 @@
 package golang
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/kyleconroy/sqlc/internal/metadata"
@@ -79,10 +80,14 @@ func (v QueryValue) Scan() string {
 		}
 	} else {
 		for _, f := range v.Struct.Fields {
+			ref := fmt.Sprintf("%s.%s", v.Name, f.Name)
+			if f.Struct != "" {
+				ref = fmt.Sprintf("%s.%s.%s", v.Name, f.Struct, f.Name)
+			}
 			if strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" {
-				out = append(out, "pq.Array(&"+v.Name+"."+f.Name+")")
+				out = append(out, "pq.Array(&"+ref+")")
 			} else {
-				out = append(out, "&"+v.Name+"."+f.Name)
+				out = append(out, "&"+ref)
 			}
 		}
 	}
diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index 0e0d316426..e42128b84f 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -251,6 +251,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 					continue
 				}
 
+				// Cheaper to compare the pointer instead of the name.
 				if gs != c.Embed {
 					gs = nil
 					break
@@ -285,9 +286,18 @@ func columnsToStruct(r *compiler.Result, name string, columns []goColumn, settin
 	gs := Struct{
 		Name: name,
 	}
+	embedded := map[string]interface{}{}
 	seen := map[string]int{}
 	suffixes := map[int]int{}
 	for i, c := range columns {
+		if c.Embed != nil {
+			if _, ok := embedded[c.Embed.Name]; !ok {
+				// We only want to include each embedded struct once.
+				gs.Embedded = append(gs.Embedded, *c.Embed)
+				embedded[c.Embed.Name] = nil
+			}
+		}
+
 		colName := columnName(c.Column, i)
 		tagName := colName
 		fieldName := StructName(colName, settings)
@@ -311,11 +321,15 @@ func columnsToStruct(r *compiler.Result, name string, columns []goColumn, settin
 		if settings.Go.EmitJSONTags {
 			tags["json:"] = JSONTagName(tagName, settings)
 		}
-		gs.Fields = append(gs.Fields, Field{
+		f := Field{
 			Name: fieldName,
 			Type: goType(r, c.Column, settings),
 			Tags: tags,
-		})
+		}
+		if c.Embed != nil {
+			f.Struct = c.Embed.Name
+		}
+		gs.Fields = append(gs.Fields, f)
 		seen[colName]++
 	}
 	return &gs
diff --git a/internal/codegen/golang/struct.go b/internal/codegen/golang/struct.go
index f2c6c14ab7..2f643bc3d0 100644
--- a/internal/codegen/golang/struct.go
+++ b/internal/codegen/golang/struct.go
@@ -8,11 +8,11 @@ import (
 )
 
 type Struct struct {
-	Table   core.FQN
-	Name    string
-	Structs []Struct
-	Fields  []Field
-	Comment string
+	Table    core.FQN
+	Name     string
+	Embedded []Struct
+	Fields   []Field
+	Comment  string
 }
 
 func StructName(name string, settings config.CombinedSettings) string {

From a37d35a85c2467686d6bed7cda686dad3d44fe84 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 04:08:26 +0000
Subject: [PATCH 3/7] Fix embedding of multiple structs

---
 internal/codegen/golang/result.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index e42128b84f..877adb2708 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -184,7 +184,10 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 				Typ:  goType(r, c, settings),
 			}
 		} else if len(query.Columns) > 1 {
-			var columns []goColumn
+			var (
+				columns  []goColumn
+				embedded = map[string]interface{}{}
+			)
 
 			for ci := 0; ci < len(query.Columns); {
 				c := query.Columns[ci]
@@ -201,6 +204,10 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 					if len(query.Columns) < len(s.Fields) {
 						continue
 					}
+					// We can only embed one struct of each type.
+					if _, ok := embedded[s.Name]; ok {
+						continue
+					}
 					same := true
 					for fi, f := range s.Fields {
 						fieldOffset := ci + fi
@@ -230,6 +237,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 				colsMatched := 1
 				if embed != nil {
 					colsMatched = len(embed.Fields)
+					embedded[embed.Name] = nil
 				}
 				for colID := ci; colID < ci+colsMatched; colID++ {
 					columns = append(columns, goColumn{

From 76b5622c2175755989d7e447284ad22a02a344af Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 04:22:06 +0000
Subject: [PATCH 4/7] Fix duplicate identifier incrementing

---
 internal/codegen/golang/result.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index 877adb2708..aad94075b7 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -338,7 +338,9 @@ func columnsToStruct(r *compiler.Result, name string, columns []goColumn, settin
 			f.Struct = c.Embed.Name
 		}
 		gs.Fields = append(gs.Fields, f)
-		seen[colName]++
+		if c.Embed == nil {
+			seen[colName]++
+		}
 	}
 	return &gs
 }

From 28cc3bcb907cf9d37b0aae86e6e74678a9db1749 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 04:26:37 +0000
Subject: [PATCH 5/7] Fix single columns marshal

---
 internal/codegen/golang/result.go | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index aad94075b7..4dfe3c046e 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -250,20 +250,16 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 			}
 
 			var emit bool
-			var gs *Struct
+			gs := columns[0].Embed
 			// Check if all columns match a consistent embedded struct.
 			// If they do, we don't need to generate a new struct for the row.
 			for _, c := range columns {
-				if gs == nil {
-					gs = c.Embed
-					continue
-				}
-
 				// Cheaper to compare the pointer instead of the name.
 				if gs != c.Embed {
 					gs = nil
 					break
 				}
+				gs = c.Embed
 			}
 
 			if gs == nil {

From 52f49961e87bdda5e7c9c2ff2e72dfdbe02c0060 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 04:34:45 +0000
Subject: [PATCH 6/7] Directly access embedded struct when referenced

---
 internal/codegen/golang/result.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index 4dfe3c046e..6b8fed4f1b 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -314,6 +314,9 @@ func columnsToStruct(r *compiler.Result, name string, columns []goColumn, settin
 			suffix = v + 1
 		}
 		suffixes[c.id] = suffix
+		if c.Embed != nil {
+			suffix = 0
+		}
 		if suffix > 0 {
 			tagName = fmt.Sprintf("%s_%d", tagName, suffix)
 			fieldName = fmt.Sprintf("%s_%d", fieldName, suffix)

From 31f95b0cae10ec6501a280f196a350176688d3d8 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 26 Apr 2021 04:36:58 +0000
Subject: [PATCH 7/7] Fix length access when too many fields

---
 internal/codegen/golang/result.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index 6b8fed4f1b..949a687406 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -201,7 +201,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 					}
 					// If the query doesn't have enough fields, it cannot
 					// fufill the struct.
-					if len(query.Columns) < len(s.Fields) {
+					if len(query.Columns)-ci < len(s.Fields) {
 						continue
 					}
 					// We can only embed one struct of each type.