From 5c39787fbdac3196994e1b563a1b3d155bd14f92 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Sun, 6 Jun 2021 11:24:56 +0900
Subject: [PATCH] Enable FirstWin option for stream decoder

---
 benchmarks/decode_test.go  | 15 +++++++++++++++
 decode.go                  |  7 +++++++
 internal/decoder/stream.go |  2 ++
 internal/decoder/struct.go | 29 +++++++++++++++++++++++++++--
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/benchmarks/decode_test.go b/benchmarks/decode_test.go
index ae3439a..d074c1d 100644
--- a/benchmarks/decode_test.go
+++ b/benchmarks/decode_test.go
@@ -462,3 +462,18 @@ func Benchmark_Decode_LargeStruct_Stream_GoJson(b *testing.B) {
 		}
 	}
 }
+
+func Benchmark_Decode_LargeStruct_Stream_GoJsonFirstWinMode(b *testing.B) {
+	b.ReportAllocs()
+	reader := bytes.NewReader(LargeFixture)
+	for i := 0; i < b.N; i++ {
+		result := LargePayload{}
+		reader.Reset(LargeFixture)
+		if err := gojson.NewDecoder(reader).DecodeWithOption(
+			&result,
+			gojson.DecodeFieldPriorityFirstWin(),
+		); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
diff --git a/decode.go b/decode.go
index 31df2c3..e2cfcd4 100644
--- a/decode.go
+++ b/decode.go
@@ -134,6 +134,10 @@ func (d *Decoder) Buffered() io.Reader {
 // See the documentation for Unmarshal for details about
 // the conversion of JSON into a Go value.
 func (d *Decoder) Decode(v interface{}) error {
+	return d.DecodeWithOption(v)
+}
+
+func (d *Decoder) DecodeWithOption(v interface{}, optFuncs ...DecodeOptionFunc) error {
 	header := (*emptyInterface)(unsafe.Pointer(&v))
 	typ := header.typ
 	ptr := uintptr(header.ptr)
@@ -153,6 +157,9 @@ func (d *Decoder) Decode(v interface{}) error {
 		return err
 	}
 	s := d.s
+	for _, optFunc := range optFuncs {
+		optFunc(s.Option)
+	}
 	if err := dec.DecodeStream(s, 0, header.ptr); err != nil {
 		return err
 	}
diff --git a/internal/decoder/stream.go b/internal/decoder/stream.go
index 9171107..c274683 100644
--- a/internal/decoder/stream.go
+++ b/internal/decoder/stream.go
@@ -25,6 +25,7 @@ type Stream struct {
 	allRead               bool
 	UseNumber             bool
 	DisallowUnknownFields bool
+	Option                *Option
 }
 
 func NewStream(r io.Reader) *Stream {
@@ -32,6 +33,7 @@ func NewStream(r io.Reader) *Stream {
 		r:       r,
 		bufSize: initBufSize,
 		buf:     make([]byte, initBufSize),
+		Option:  &Option{},
 	}
 }
 
diff --git a/internal/decoder/struct.go b/internal/decoder/struct.go
index 70cf444..b99ab0a 100644
--- a/internal/decoder/struct.go
+++ b/internal/decoder/struct.go
@@ -661,6 +661,14 @@ func (d *structDecoder) DecodeStream(s *Stream, depth int64, p unsafe.Pointer) e
 		s.cursor++
 		return nil
 	}
+	var (
+		seenFields   map[int]struct{}
+		seenFieldNum int
+	)
+	firstWin := (s.Option.Flag & FirstWinOption) != 0
+	if firstWin {
+		seenFields = make(map[int]struct{}, d.fieldUniqueNameNum)
+	}
 	for {
 		s.reset()
 		field, key, err := d.keyStreamDecoder(d, s)
@@ -675,8 +683,25 @@ func (d *structDecoder) DecodeStream(s *Stream, depth int64, p unsafe.Pointer) e
 			if field.err != nil {
 				return field.err
 			}
-			if err := field.dec.DecodeStream(s, depth, unsafe.Pointer(uintptr(p)+field.offset)); err != nil {
-				return err
+			if firstWin {
+				if _, exists := seenFields[field.fieldIdx]; exists {
+					if err := s.skipValue(depth); err != nil {
+						return err
+					}
+				} else {
+					if err := field.dec.DecodeStream(s, depth, unsafe.Pointer(uintptr(p)+field.offset)); err != nil {
+						return err
+					}
+					seenFieldNum++
+					if d.fieldUniqueNameNum <= seenFieldNum {
+						return s.skipObject(depth)
+					}
+					seenFields[field.fieldIdx] = struct{}{}
+				}
+			} else {
+				if err := field.dec.DecodeStream(s, depth, unsafe.Pointer(uintptr(p)+field.offset)); err != nil {
+					return err
+				}
 			}
 		} else if s.DisallowUnknownFields {
 			return fmt.Errorf("json: unknown field %q", key)