Repository: ogen-go/ogen
Pull request #559 (2022-09-08 14:00:24)

refactor: use generic helpers to get sorted slice of map keys


Patch:
diff --git a/examples/go.mod b/examples/go.mod
index 85fe01e9..b269d52f 100644
--- a/examples/go.mod
+++ b/examples/go.mod
@@ -13,6 +13,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.9.0
 	go.uber.org/multierr v1.8.0
 	go.uber.org/zap v1.23.0
+	golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0
 )
 
 require (
diff --git a/examples/go.sum b/examples/go.sum
index 548b1372..d00ed5b7 100644
--- a/examples/go.sum
+++ b/examples/go.sum
@@ -52,6 +52,8 @@ go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8=
 go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/zap v1.23.0 h1:OjGQ5KQDEUawVHxNwQgPpiypGHOxo2mNZsOqTak4fFY=
 go.uber.org/zap v1.23.0/go.mod h1:D+nX8jyLsMHMYrln8A0rJjFt/T/9/bGgIhAqxv5URuY=
+golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0 h1:17k44ji3KFYG94XS5QEFC8pyuOlMh3IoR+vkmTZmJJs=
+golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 h1:6zppjxzCulZykYSLyVDYbneBfbaBIQPYMevg0bEwv2s=
 golang.org/x/net v0.0.0-20220822230855-b0a4917ee28c h1:JVAXQ10yGGVbSyoer5VILysz6YKjdNT2bsvlayjqhes=
 golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde h1:ejfdSekXMDxDLbRrJMwUk6KnSLZ2McaUCVcIKM+N6jc=
diff --git a/examples/hack.go b/examples/hack.go
index d07e17ed..b2d37e20 100644
--- a/examples/hack.go
+++ b/examples/hack.go
@@ -2,4 +2,9 @@
 
 package examples
 
-import _ "go.uber.org/zap"
+import (
+	_ "go.uber.org/zap"
+	_ "golang.org/x/exp/constraints"
+	_ "golang.org/x/exp/maps"
+	_ "golang.org/x/exp/slices"
+)
diff --git a/gen/fix.go b/gen/fix.go
index 402015c0..a0c5bb06 100644
--- a/gen/fix.go
+++ b/gen/fix.go
@@ -2,43 +2,44 @@ package gen
 
 import (
 	"reflect"
-	"sort"
 
 	"github.com/go-faster/errors"
+	"golang.org/x/exp/maps"
 
 	"github.com/ogen-go/ogen/gen/ir"
+	"github.com/ogen-go/ogen/internal/xmaps"
 )
 
 // Example:
-// ...
-// responses:
-//   200:
-//     contents:
-//       application/json:
-//         ref: #/components/schemas/Foo
-//   202:
-//     contents:
-//       application/json:
-//         ref: #/components/schemas/Foo
+//
+//	responses:
+//		200:
+//		  contents:
+//		    application/json:
+//		      ref: #/components/schemas/Foo
+//		202:
+//		  contents:
+//		    application/json:
+//		      ref: #/components/schemas/Foo
 //
 // This response refers to the same schema for different
 // status codes, and it will cause a collision:
 //
-// func encodeResponse(resp FooResponse) {
-//     switch resp.(type) {
-//	   case *Foo:
-//     case *Foo:
-//     }
-// }
+//	func encodeResponse(resp FooResponse) {
+//	    switch resp.(type) {
+//		case *Foo:
+//	    case *Foo:
+//	    }
+//	}
 //
 // To prevent collision we wrap referenced schema with aliases
 // and use them instead.
 //
-// type FooResponseOK Foo
-// func(*FooResponseOK) FooResponse() {}
+//	type FooResponseOK Foo
+//	func(*FooResponseOK) FooResponse() {}
 //
-// type FooResponseAccepted Foo
-// func(*FooResponseAccepted) FooResponse() {}
+//	type FooResponseAccepted Foo
+//	func(*FooResponseAccepted) FooResponse() {}
 //
 // Referring to the same schema in different content types
 // also can cause a collision and it will be fixed in the same way.
@@ -52,19 +53,14 @@ func fixEqualResponses(ctx *genctx, op *ir.Operation) error {
 	// (in case of referenced responses), we copy the response.
 	op.Responses = cloneResponse(op.Responses)
 
-	var statusCodes []int
-	for code := range op.Responses.StatusCode {
-		statusCodes = append(statusCodes, code)
-	}
-	sort.Ints(statusCodes)
-
+	statusCodes := xmaps.SortedKeys(op.Responses.StatusCode)
 	type candidate struct {
 		renameTo string
 		encoding ir.Encoding
 		typ      *ir.Type
 
 		replaceNoc bool
-		replaceCT  string
+		replaceCT  ir.ContentType
 		response   *ir.Response
 	}
 
@@ -101,30 +97,22 @@ func fixEqualResponses(ctx *genctx, op *ir.Operation) error {
 			}
 
 			var (
-				lcontents []string
-				rcontents []string
+				lcontents = xmaps.SortedKeys(lresp.Contents)
+				rcontents = xmaps.SortedKeys(rresp.Contents)
 			)
-			for ct := range lresp.Contents {
-				lcontents = append(lcontents, string(ct))
-			}
-			for ct := range rresp.Contents {
-				rcontents = append(rcontents, string(ct))
-			}
-			sort.Strings(lcontents)
-			sort.Strings(rcontents)
 			for _, lct := range lcontents {
 				for _, rct := range rcontents {
 					if lcode == rcode && lct == rct {
 						continue
 					}
-					lmedia, rmedia := lresp.Contents[ir.ContentType(lct)], rresp.Contents[ir.ContentType(rct)]
+					lmedia, rmedia := lresp.Contents[lct], rresp.Contents[rct]
 					ltype, rtype := lmedia.Type, rmedia.Type
 					if reflect.DeepEqual(ltype, rtype) {
-						lname, err := pascal(op.Name, lct, statusText(lcode))
+						lname, err := pascal(op.Name, string(lct), statusText(lcode))
 						if err != nil {
 							return errors.Wrap(err, "lname")
 						}
-						rname, err := pascal(op.Name, rct, statusText(rcode))
+						rname, err := pascal(op.Name, string(rct), statusText(rcode))
 						if err != nil {
 							return errors.Wrap(err, "rname")
 						}
@@ -162,7 +150,7 @@ func fixEqualResponses(ctx *genctx, op *ir.Operation) error {
 			continue
 		}
 
-		candidate.response.Contents[ir.ContentType(candidate.replaceCT)] = ir.Media{
+		candidate.response.Contents[candidate.replaceCT] = ir.Media{
 			Encoding: candidate.encoding,
 			Type:     alias,
 		}
@@ -179,7 +167,7 @@ func cloneResponse(r *ir.Responses) *ir.Responses {
 
 		return &ir.Response{
 			NoContent:      r.NoContent,
-			Contents:       cloneContents(r.Contents),
+			Contents:       maps.Clone(r.Contents),
 			Headers:        r.Headers,
 			WithStatusCode: r.WithStatusCode,
 			WithHeaders:    r.WithHeaders,
@@ -215,34 +203,31 @@ func fixEqualRequests(ctx *genctx, op *ir.Operation) error {
 
 	type candidate struct {
 		renameTo string
-		ctype    string
+		ctype    ir.ContentType
 		encoding ir.Encoding
 		t        *ir.Type
 	}
-	var candidates []candidate
-
-	var contents []string
-	for ct := range op.Request.Contents {
-		contents = append(contents, string(ct))
-	}
-	sort.Strings(contents)
+	var (
+		candidates []candidate
+		contents   = xmaps.SortedKeys(op.Request.Contents)
+	)
 
 	for _, lcontent := range contents {
-		lmedia := op.Request.Contents[ir.ContentType(lcontent)]
+		lmedia := op.Request.Contents[lcontent]
 		ltype := lmedia.Type
 		for _, rcontent := range contents {
 			if lcontent == rcontent {
 				continue
 			}
 
-			rmedia := op.Request.Contents[ir.ContentType(rcontent)]
+			rmedia := op.Request.Contents[rcontent]
 			rtype := rmedia.Type
 			if reflect.DeepEqual(ltype, rtype) {
-				lname, err := pascal(op.Name, lcontent)
+				lname, err := pascal(op.Name, string(lcontent))
 				if err != nil {
 					return errors.Wrap(err, "lname")
 				}
-				rname, err := pascal(op.Name, rcontent)
+				rname, err := pascal(op.Name, string(rcontent))
 				if err != nil {
 					return errors.Wrap(err, "rname")
 				}
@@ -270,7 +255,7 @@ func fixEqualRequests(ctx *genctx, op *ir.Operation) error {
 		// g.saveType(alias)
 		ctx.local.types[alias.Name] = alias
 
-		op.Request.Contents[ir.ContentType(candidate.ctype)] = ir.Media{
+		op.Request.Contents[candidate.ctype] = ir.Media{
 			Encoding: candidate.encoding,
 			Type:     alias,
 		}
@@ -282,19 +267,7 @@ func fixEqualRequests(ctx *genctx, op *ir.Operation) error {
 func cloneRequest(r *ir.Request) *ir.Request {
 	return &ir.Request{
 		Type:     r.Type,
-		Contents: cloneContents(r.Contents),
+		Contents: maps.Clone(r.Contents),
 		Spec:     r.Spec,
 	}
 }
-
-func cloneContents(c map[ir.ContentType]ir.Media) map[ir.ContentType]ir.Media {
-	if c == nil {
-		return nil
-	}
-
-	result := make(map[ir.ContentType]ir.Media, len(c))
-	for ct, m := range c {
-		result[ct] = m
-	}
-	return result
-}
diff --git a/gen/gen_contents.go b/gen/gen_contents.go
index dfbb921e..4cbd8cda 100644
--- a/gen/gen_contents.go
+++ b/gen/gen_contents.go
@@ -4,12 +4,12 @@ import (
 	"fmt"
 	"mime"
 	"path"
-	"sort"
 
 	"github.com/go-faster/errors"
 	"go.uber.org/zap"
 
 	"github.com/ogen-go/ogen/gen/ir"
+	"github.com/ogen-go/ogen/internal/xmaps"
 	"github.com/ogen-go/ogen/openapi"
 )
 
@@ -189,14 +189,10 @@ func (g *Generator) generateContents(
 		result = make(map[ir.ContentType]ir.Media, len(contents))
 		names  = make(map[ir.ContentType]string, len(contents))
 
-		keys        = make([]string, 0, len(contents))
+		keys        = xmaps.SortedKeys(contents)
 		unsupported []string
 		lastErr     error
 	)
-	for k := range contents {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
 
 	for _, contentType := range keys {
 		media := contents[contentType]
diff --git a/gen/gen_responses.go b/gen/gen_responses.go
index aad5565c..bdab19de 100644
--- a/gen/gen_responses.go
+++ b/gen/gen_responses.go
@@ -9,6 +9,7 @@ import (
 	"github.com/go-faster/errors"
 
 	"github.com/ogen-go/ogen/gen/ir"
+	"github.com/ogen-go/ogen/internal/xmaps"
 	"github.com/ogen-go/ogen/openapi"
 )
 
@@ -335,13 +336,7 @@ func injectHeaderFields(headers map[string]*ir.Parameter, t *ir.Type) {
 		panic(fmt.Sprintf("expected struct, got %q", t.Kind))
 	}
 
-	keys := make([]string, 0, len(headers))
-	for k := range headers {
-		keys = append(keys, k)
-	}
-
-	sort.Strings(keys)
-	for _, key := range keys {
+	for _, key := range xmaps.SortedKeys(headers) {
 		h := headers[key]
 		t.Fields = append(t.Fields, &ir.Field{
 			Name: h.Name,
diff --git a/gen/ir/type_iface.go b/gen/ir/type_iface.go
index 92f1ecdf..1399159a 100644
--- a/gen/ir/type_iface.go
+++ b/gen/ir/type_iface.go
@@ -3,6 +3,8 @@ package ir
 import (
 	"fmt"
 	"sort"
+
+	"github.com/ogen-go/ogen/internal/xmaps"
 )
 
 func (t *Type) CanHaveMethods() bool {
@@ -54,12 +56,7 @@ func (t *Type) Methods() []string {
 		panic(fmt.Sprintf("unexpected kind: %s", t.Kind))
 	}
 
-	var result []string
-	for m := range ms {
-		result = append(result, m)
-	}
-	sort.Strings(result)
-	return result
+	return xmaps.SortedKeys(ms)
 }
 
 func (t *Type) ListImplementations() []*Type {
diff --git a/gen/schema_gen_sum.go b/gen/schema_gen_sum.go
index cffedcb3..e86abb2c 100644
--- a/gen/schema_gen_sum.go
+++ b/gen/schema_gen_sum.go
@@ -11,6 +11,7 @@ import (
 	"github.com/go-faster/jx"
 
 	"github.com/ogen-go/ogen/gen/ir"
+	"github.com/ogen-go/ogen/internal/xmaps"
 	"github.com/ogen-go/ogen/jsonschema"
 )
 
@@ -344,14 +345,10 @@ func (g *schemaGen) oneOf(name string, schema *jsonschema.Schema) (*ir.Type, err
 	for _, s := range sum.SumOf {
 		k := s.Name
 		f := uniq[k]
-		v := sumVariant{
-			Name: k,
-		}
-		for fieldName := range f {
-			v.Unique = append(v.Unique, fieldName)
-		}
-		sort.Strings(v.Unique)
-		variants = append(variants, v)
+		variants = append(variants, sumVariant{
+			Name:   k,
+			Unique: xmaps.SortedKeys(f),
+		})
 	}
 	sort.SliceStable(variants, func(i, j int) bool {
 		a := variants[i]
diff --git a/gen/write.go b/gen/write.go
index 10172f21..f593d76c 100644
--- a/gen/write.go
+++ b/gen/write.go
@@ -5,13 +5,13 @@ import (
 	"fmt"
 	"os"
 	"regexp"
-	"sort"
 	"text/template"
 
 	"github.com/go-faster/errors"
 	"golang.org/x/tools/imports"
 
 	"github.com/ogen-go/ogen/gen/ir"
+	"github.com/ogen-go/ogen/internal/xmaps"
 )
 
 type TemplateConfig struct {
@@ -34,7 +34,7 @@ func (t TemplateConfig) SkipTest(typ *ir.Type) bool {
 	return t.skipTestRegex != nil && t.skipTestRegex.MatchString(typ.Name)
 }
 
-func (t TemplateConfig) collectStrings(cb func(typ *ir.Type) []string) (r []string) {
+func (t TemplateConfig) collectStrings(cb func(typ *ir.Type) []string) []string {
 	var (
 		add  func(typ *ir.Type)
 		m    = map[string]struct{}{}
@@ -81,11 +81,7 @@ func (t TemplateConfig) collectStrings(cb func(typ *ir.Type) []string) (r []stri
 		return nil
 	})
 
-	for exp := range m {
-		r = append(r, exp)
-	}
-	sort.Strings(r)
-	return r
+	return xmaps.SortedKeys(m)
 }
 
 // RegexStrings returns slice of all unique regex validators.
diff --git a/go.mod b/go.mod
index 2b1be1eb..7cdbc08a 100644
--- a/go.mod
+++ b/go.mod
@@ -17,6 +17,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.9.0
 	go.uber.org/multierr v1.8.0
 	go.uber.org/zap v1.23.0
+	golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0
 	golang.org/x/net v0.0.0-20220822230855-b0a4917ee28c
 	golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde
 	golang.org/x/tools v0.1.12
diff --git a/go.sum b/go.sum
index c6efaf1f..89ba579a 100644
--- a/go.sum
+++ b/go.sum
@@ -69,6 +69,8 @@ go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95a
 go.uber.org/zap v1.23.0 h1:OjGQ5KQDEUawVHxNwQgPpiypGHOxo2mNZsOqTak4fFY=
 go.uber.org/zap v1.23.0/go.mod h1:D+nX8jyLsMHMYrln8A0rJjFt/T/9/bGgIhAqxv5URuY=
 golang.org/x/crypto v0.0.0-20220214200702-86341886e292/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
+golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0 h1:17k44ji3KFYG94XS5QEFC8pyuOlMh3IoR+vkmTZmJJs=
+golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 h1:6zppjxzCulZykYSLyVDYbneBfbaBIQPYMevg0bEwv2s=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
diff --git a/internal/xmaps/sorted_keys.go b/internal/xmaps/sorted_keys.go
new file mode 100644
index 00000000..79bdb067
--- /dev/null
+++ b/internal/xmaps/sorted_keys.go
@@ -0,0 +1,14 @@
+package xmaps
+
+import (
+	"golang.org/x/exp/constraints"
+	"golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
+)
+
+// SortedKeys returns a sorted slice of keys in the map.
+func SortedKeys[M ~map[K]V, K constraints.Ordered, V any](m M) []K {
+	r := maps.Keys(m)
+	slices.Sort(r)
+	return r
+}
diff --git a/openapi/parser/parser.go b/openapi/parser/parser.go
index 3d98ade6..396946b4 100644
--- a/openapi/parser/parser.go
+++ b/openapi/parser/parser.go
@@ -2,14 +2,13 @@
 package parser
 
 import (
-	"sort"
-
 	"github.com/go-faster/errors"
 	yaml "github.com/go-faster/yamlx"
 
 	"github.com/ogen-go/ogen"
 	"github.com/ogen-go/ogen/internal/jsonpointer"
 	"github.com/ogen-go/ogen/internal/location"
+	"github.com/ogen-go/ogen/internal/xmaps"
 	"github.com/ogen-go/ogen/jsonschema"
 	"github.com/ogen-go/ogen/openapi"
 )
@@ -128,14 +127,9 @@ func (p *parser) parsePathItems() error {
 		paths = map[string]struct{}{}
 
 		pathsLoc = p.rootLoc.Field("paths")
-		keys     = make([]string, 0, len(p.spec.Paths))
 	)
-	for k := range p.spec.Paths {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
 
-	for _, path := range keys {
+	for _, path := range xmaps.SortedKeys(p.spec.Paths) {
 		item := p.spec.Paths[path]
 		if err := func() error {
 			id, err := pathID(path)


Test patch:
diff --git a/json/equal_test.go b/json/equal_test.go
index 0580cba4..25142c53 100644
--- a/json/equal_test.go
+++ b/json/equal_test.go
@@ -2,11 +2,12 @@ package json
 
 import (
 	"fmt"
-	"sort"
 	"testing"
 
 	"github.com/go-faster/jx"
 	"github.com/stretchr/testify/require"
+
+	"github.com/ogen-go/ogen/internal/xmaps"
 )
 
 func TestEqual(t *testing.T) {
@@ -86,12 +87,7 @@ func TestEqual(t *testing.T) {
 	}
 
 	sortedIter := func(cb func(k string, tts []testCase)) {
-		var keys []string
-		for k := range tests {
-			keys = append(keys, k)
-		}
-		sort.Strings(keys)
-		for _, k := range keys {
+		for _, k := range xmaps.SortedKeys(tests) {
 			cb(k, tests[k])
 		}
 	}