impl: map size prediction

bytedance · Dec 14, 2023 · 66ef174 · 66ef174
1 parent 3490f68
commit 66ef174
Show file tree

Hide file tree

Showing 6 changed files with 117 additions and 28 deletions.
diff --git a/internal/decoder/assembler_regabi_amd64.go b/internal/decoder/assembler_regabi_amd64.go
@@ -20,18 +20,19 @@
 package decoder
 
 import (
-    `encoding/json`
-    `fmt`
-    `math`
-    `reflect`
-    `unsafe`
-
-    `github.com/bytedance/sonic/internal/caching`
-    `github.com/bytedance/sonic/internal/jit`
-    `github.com/bytedance/sonic/internal/native`
-    `github.com/bytedance/sonic/internal/native/types`
-    `github.com/bytedance/sonic/internal/rt`
-    `github.com/twitchyliquid64/golang-asm/obj`
+	"encoding/json"
+	"fmt"
+	"math"
+	"reflect"
+	"unsafe"
+
+	"github.com/bytedance/sonic/internal/caching"
+	"github.com/bytedance/sonic/internal/jit"
+	"github.com/bytedance/sonic/internal/native"
+	"github.com/bytedance/sonic/internal/native/types"
+	"github.com/bytedance/sonic/internal/rt"
+	"github.com/bytedance/sonic/option"
+	"github.com/twitchyliquid64/golang-asm/obj"
 )
 
 /** Register Allocations
@@ -868,6 +869,7 @@ func (self *_Assembler) range_unsigned_CX(i *rt.GoItab, t *rt.GoType, v uint64)
 
 var (
     _F_unquote = jit.Imm(int64(native.S_unquote))
+    _F_count_elems = jit.Imm(int64(native.S_count_elems))
 )
 
 func (self *_Assembler) slice_from(p obj.Addr, d int64) {
@@ -1149,6 +1151,7 @@ var (
     _F_growslice        = jit.Func(growslice)
     _F_makeslice        = jit.Func(makeslice)
     _F_makemap_small    = jit.Func(makemap_small)
+    _F_makemap          = jit.Func(makemap)
     _F_mapassign_fast64 = jit.Func(mapassign_fast64)
 )
 
@@ -1480,11 +1483,33 @@ func (self *_Assembler) _asm_OP_is_null_quote(p *_Instr) {
     self.Link("_not_null_quote_{n}")                                    // _not_null_quote_{n}:
 }
 
-func (self *_Assembler) _asm_OP_map_init(_ *_Instr) {
+func (self *_Assembler) _asm_OP_map_init(p *_Instr) {
     self.Emit("MOVQ" , jit.Ptr(_VP, 0), _AX)    // MOVQ    (VP), AX
     self.Emit("TESTQ", _AX, _AX)                // TESTQ   AX, AX
     self.Sjmp("JNZ"  , "_end_{n}")              // JNZ     _end_{n}
-    self.call_go(_F_makemap_small)              // CALL_GO makemap_small
+    if option.PredictContainerSize {         // OPT: estimate container size ahead
+        // fast-path: check '}' first
+        self.check_eof(1)
+        self.Emit("CMPB", jit.Sib(_IP, _IC, 1, 0), jit.Imm(int64('}')))   // CMPB    (IP)(IC), ${p.vb()}
+        self.Sjmp("JE"  , "_small_map_{n}")   
+        self.Emit("LEAQ", _ARG_s, _DI)      
+        self.Emit("SUBQ", jit.Imm(1), _IC)      
+        self.Emit("MOVQ", _IC, _ARG_ic)     
+        self.Emit("LEAQ", _ARG_ic, _SI)     
+        self.call_c(_F_count_elems)
+        self.Emit("ADDQ" , jit.Imm(1), _IC)     
+        self.Emit("TESTQ", _AX, _AX)
+        self.Sjmp("JS"   , _LB_parsing_error_v)    
+        self.Emit("MOVQ" , _AX, _BX)   
+        self.Emit("MOVQ" , jit.Type(p.vt()), _AX)
+        self.Emit("XORL" , _CX, _CX)   
+        self.call_go(_F_makemap)
+        self.Sjmp("JMP", "_end_{n}")
+        self.Link("_small_map_{n}")
+        self.call_go(_F_makemap_small)       
+    } else {
+        self.call_go(_F_makemap_small)              // CALL_GO makemap_small
+    }
     self.WritePtrAX(6, jit.Ptr(_VP, 0), false)    // MOVQ    AX, (VP)
     self.Link("_end_{n}")                       // _end_{n}:
     self.Emit("MOVQ" , _AX, _VP)                // MOVQ    AX, VP

diff --git a/internal/decoder/compiler.go b/internal/decoder/compiler.go
@@ -641,7 +641,7 @@ func (self *_Compiler) compileMapOp(p *_Program, sp int, vt reflect.Type, op _Op
     p.tag(sp + 1)
     skip := self.checkIfSkip(p, vt, '{')
     p.add(_OP_save)
-    p.add(_OP_map_init)
+    p.rtt(_OP_map_init, vt)
     p.add(_OP_save)
     p.add(_OP_lspace)
     j := p.pc()

diff --git a/internal/decoder/decoder_test.go b/internal/decoder/decoder_test.go
@@ -17,19 +17,22 @@
 package decoder
 
 import (
-    `encoding/json`
-    `runtime`
-    `runtime/debug`
-    `strings`
-    `sync`
-    `testing`
-    `time`
-    `reflect`
-
-    `github.com/bytedance/sonic/internal/rt`
-    `github.com/davecgh/go-spew/spew`
-    `github.com/stretchr/testify/assert`
-    `github.com/stretchr/testify/require`
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"runtime"
+	"runtime/debug"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/bytedance/sonic/internal/rt"
+	"github.com/bytedance/sonic/option"
+	"github.com/davecgh/go-spew/spew"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestMain(m *testing.M) {
@@ -86,6 +89,57 @@ func init() {
     _ = json.Unmarshal([]byte(TwitterJson), &_BindingValue)
 }
 
+func testMap(tt interface{}, limit int) {
+    var maker = func(N int) string{ 
+        var s = `{`
+        for i :=0; i<N; i++ {
+            s += fmt.Sprintf(`"%d":%d`, i, i)
+            if i < N-1 {
+                s += `,`
+            }
+        }
+        return s + `}`
+    }
+    for x:=0; x<limit; {
+        if b, ok := tt.(*testing.B); ok {
+            b.Run("N="+strconv.Itoa(x), func(b *testing.B) {
+                var s = maker(x)
+                for i:=0; i<b.N; i++ {
+                    var obj map[string]int
+                    _, _ = decode(s, &obj, false)
+                }
+            })
+        } else if t, ok := tt.(*testing.T); ok {
+            t.Run("N="+strconv.Itoa(x), func(t *testing.T) {
+                var obj map[string]int
+                var s = maker(x)
+                _, err := decode(s, &obj, false)
+                if err != nil {
+                    b.Fatal(err)
+                }
+            })
+        }
+        if x == 0 {
+            x = 1
+        } else {
+            x *= 10
+        }
+    }
+}
+
+func TestPredict(t *testing.T) {
+    option.PredictContainerSize = true
+    testMap(t, 10)
+    option.PredictContainerSize = false
+}
+
+func BenchmarkPredictContSize(b *testing.B) {
+    b.Run("map", func(b *testing.B) {
+        option.PredictContainerSize = true
+        testMap(b, 1001)
+    })
+    option.PredictContainerSize = false
+}
 
 func TestSkipMismatchTypeError(t *testing.T) {
     t.Run("struct", func(t *testing.T) {

diff --git a/internal/decoder/stubs_go116.go b/internal/decoder/stubs_go116.go
@@ -80,6 +80,9 @@ func growslice(et *rt.GoType, old rt.GoSlice, cap int) rt.GoSlice
 //go:linkname makemap_small runtime.makemap_small
 func makemap_small() unsafe.Pointer
 
+//go:linkname makemap runtime.makemap
+func makemap(t *rt.GoType, cap int, h unsafe.Pointer) unsafe.Pointer
+
 //go:linkname mapassign runtime.mapassign
 //goland:noinspection GoUnusedParameter
 func mapassign(t *rt.GoType, h unsafe.Pointer, k unsafe.Pointer) unsafe.Pointer

diff --git a/internal/decoder/stubs_go120.go b/internal/decoder/stubs_go120.go
@@ -80,6 +80,9 @@ func growslice(et *rt.GoType, old rt.GoSlice, cap int) rt.GoSlice
 //go:linkname makemap_small runtime.makemap_small
 func makemap_small() unsafe.Pointer
 
+//go:linkname makemap runtime.makemap
+func makemap(t *rt.GoType, cap int) unsafe.Pointer
+
 //go:linkname mapassign runtime.mapassign
 //goland:noinspection GoUnusedParameter
 func mapassign(t *rt.GoMapType, h unsafe.Pointer, k unsafe.Pointer) unsafe.Pointer

diff --git a/option/option.go b/option/option.go
@@ -42,6 +42,10 @@ var (
     // Default value(1) means `Pretouch()` will be recursively executed once,
     // if any nested struct is left (depth exceeds MaxInlineDepth)
     DefaultRecursiveDepth = 1
+
+    // enable predict the size of containers (array or object) for better speed of deserialization.
+    // WARNING: this may let sonic to consume more memory
+    PredictContainerSize bool
 )
 
 // DefaultCompileOptions set default compile options.