fix query

_example/draft-query.go

@@ -0,0 +1,87 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/goccy/go-json"
+)
+
+type User struct {
+	ID      int64
+	Name    string
+	Age     int
+	Address UserAddressResolver
+}
+
+type UserAddress struct {
+	UserID   int64
+	PostCode string
+	City     string
+	Address1 string
+	Address2 string
+}
+
+type UserRepository struct {
+	uaRepo *UserAddressRepository
+}
+
+func NewUserRepository() *UserRepository {
+	return &UserRepository{
+		uaRepo: NewUserAddressRepository(),
+	}
+}
+
+type UserAddressRepository struct{}
+
+func NewUserAddressRepository() *UserAddressRepository {
+	return &UserAddressRepository{}
+}
+
+type UserAddressResolver func(context.Context) (*UserAddress, error)
+
+func (resolver UserAddressResolver) MarshalJSON(ctx context.Context) ([]byte, error) {
+	address, err := resolver(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return json.MarshalWithQuery(address, json.QueryFromContext(ctx))
+}
+
+func (r *UserRepository) FindByID(ctx context.Context, id int64) (*User, error) {
+	v := User{ID: id, Name: "Ken", Age: 20}
+	// resolve relation from User to UserAddress
+	v.Address = func(ctx context.Context) (*UserAddress, error) {
+		return r.uaRepo.FindByUserID(ctx, v.ID)
+	}
+	return v, nil
+}
+
+func (*UserAddressRepository) FindByUserID(ctx context.Context, id int64) (*UserAddress, error) {
+	return &UserAddress{
+		UserID:   id,
+		City:     "A",
+		Address1: "hoge",
+		Address2: "fuga",
+	}, nil
+}
+
+func main() {
+	userRepo := NewUserRepository()
+	user, err := userRepo.FindByID(context.Background(), 1)
+	if err != nil {
+		panic(err)
+	}
+	q := json.NewQuery().Fields(
+		"Name",
+		"Age",
+		json.NewQuery("Address").Fields(
+			"City",
+		),
+	)
+	b, err := json.MarshalWithQuery(user, q)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Println(string(b))
+}

encode.go

@@ -148,6 +148,27 @@ func marshal(v interface{}, opt EncodeOption) ([]byte, error) {
 	return copied, nil
 }
 
+func marshalQuery(v interface{}, q *Query) ([]byte, error) {
+	ctx := takeEncodeRuntimeContext()
+
+	buf, err := encode(ctx, v, EncodeOptionHTMLEscape)
+	if err != nil {
+		releaseEncodeRuntimeContext(ctx)
+		return nil, err
+	}
+
+	// this line exists to escape call of `runtime.makeslicecopy` .
+	// if use `make([]byte, len(buf)-1)` and `copy(copied, buf)`,
+	// dst buffer size and src buffer size are differrent.
+	// in this case, compiler uses `runtime.makeslicecopy`, but it is slow.
+	buf = buf[:len(buf)-1]
+	copied := make([]byte, len(buf))
+	copy(copied, buf)
+
+	releaseEncodeRuntimeContext(ctx)
+	return copied, nil
+}
+
 func marshalNoEscape(v interface{}, opt EncodeOption) ([]byte, error) {
 	ctx := takeEncodeRuntimeContext()
 

json.go

@@ -13,6 +13,10 @@ type Marshaler interface {
 	MarshalJSON() ([]byte, error)
 }
 
+type QueryResolver interface {
+	ResolveQueryJSON(*Query) (interface{}, error)
+}
+
 // Unmarshaler is the interface implemented by types
 // that can unmarshal a JSON description of themselves.
 // The input can be assumed to be a valid encoding of
@@ -172,6 +176,10 @@ func MarshalWithOption(v interface{}, optFuncs ...EncodeOptionFunc) ([]byte, err
 	return marshal(v, opt)
 }
 
+func MarshalWithQuery(v interface{}, q *Query) ([]byte, error) {
+	return marshalQuery(v, q)
+}
+
 // MarshalIndent is like Marshal but applies Indent to format the output.
 // Each JSON element in the output will begin on a new line beginning with prefix
 // followed by one or more copies of indent according to the indentation nesting.

query.go

@@ -0,0 +1,84 @@
+package json
+
+import (
+	"context"
+	"fmt"
+)
+
+type Query struct {
+	name   string
+	fields []*Query
+	err    error
+}
+
+type queryKey struct{}
+
+func (q *Query) String() string {
+	if q.err != nil {
+		return ""
+	}
+	if q.fields == nil {
+		return ""
+	}
+	b, err := Marshal(q.dump())
+	if err != nil {
+		return ""
+	}
+	return string(b)
+}
+
+func (q *Query) Error() error {
+	return q.err
+}
+
+func (q *Query) Fields(fieldNameOrQueryList ...interface{}) *Query {
+	for _, fieldNameOrQuery := range fieldNameOrQueryList {
+		switch v := fieldNameOrQuery.(type) {
+		case string:
+			q.fields = append(q.fields, &Query{name: v})
+		case *Query:
+			q.fields = append(q.fields, v)
+			q.err = v.err
+		default:
+			q.err = fmt.Errorf("children types must be string or *Query but found %T", fieldNameOrQuery)
+		}
+		if q.err != nil {
+			break
+		}
+	}
+	return q
+}
+
+func (q *Query) dump() interface{} {
+	fields := []interface{}{}
+	for _, field := range q.fields {
+		fields = append(fields, field.dump())
+	}
+	if q.name != "" {
+		return map[string][]interface{}{
+			q.name: fields,
+		}
+	}
+	return interface{}(fields)
+}
+
+func NewQuery(name ...string) *Query {
+	if len(name) > 1 {
+		return &Query{err: fmt.Errorf(
+			"NewQuery's argument allow empty or single name only, but passed %v", name,
+		)}
+	}
+	return &Query{name: name}
+}
+
+func QueryFromContext(ctx context.Context) *Query {
+	query := ctx.Value(queryKey{})
+	if query == nil {
+		return nil
+	}
+	q, ok := query.(*Query)
+	if !ok {
+		return nil
+	}
+	return q
+}