From 4adf4176fb573e3d21034c70b47098da74569044 Mon Sep 17 00:00:00 2001
From: liu <liuqiang.06@bytedance.com>
Date: Tue, 9 Jul 2024 12:44:05 +0800
Subject: [PATCH] feat: support marshal Inf/NaN as  with option (#669)

---
 api.go                                        |  3 +++
 encode_test.go                                | 26 +++++++++++++++++++
 encoder/encoder_native.go                     |  3 +++
 internal/encoder/alg/opts.go                  |  1 +
 internal/encoder/encoder.go                   |  6 ++++-
 internal/encoder/vm/vm.go                     |  8 ++++++
 .../encoder/x86/assembler_regabi_amd64.go     | 18 ++++++++++---
 sonic.go                                      |  3 +++
 8 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/api.go b/api.go
index ef9f8efaf..1dc79664b 100644
--- a/api.go
+++ b/api.go
@@ -77,6 +77,9 @@ type Config struct {
     
     // NoEncoderNewline indicates that the encoder should not add a newline after every message
     NoEncoderNewline bool
+
+    // Encode Infinity or Nan float into `null`, instead of returning an error.
+    EncodeNullForInfOrNan bool
 }
  
 var (
diff --git a/encode_test.go b/encode_test.go
index 09639cfaa..64a5e68b0 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -1198,4 +1198,30 @@ func TestEncoder_RandomInvalidUtf8(t *testing.T) {
         testEncodeInvalidUtf8(t, genRandJsonBytes(maxLen))
         testEncodeInvalidUtf8(t, genRandJsonRune(maxLen))
     }
+}
+
+func TestMarshalInfOrNan(t *testing.T) {
+    tests := [] interface{}{
+        math.Inf(1), math.Inf(-1), math.NaN(), float32(math.Inf(1)), float32(math.Inf(-1)), float32(math.NaN()), []interface{}{math.Inf(1), math.Inf(-1), math.NaN(), float32(math.Inf(1)), float32(math.Inf(-1)), float32(math.NaN())},
+
+		[]float64{math.Inf(1), math.Inf(-1), math.NaN()},
+		[]float32{float32(math.Inf(1)), float32(math.Inf(-1)), float32(math.NaN())},
+    }
+
+    allowNanInf := Config {
+        EncodeNullForInfOrNan: true,
+    }.Froze()
+    for _, tt := range tests {
+        b, err := allowNanInf.Marshal(tt)
+        assert.Nil(t, err)
+        if len(b) == 4 {
+            assert.Equal(t, string(b), "null")
+        } else {
+            println(string(b))
+        }
+
+        b, err = Marshal(tt)
+        assert.NotNil(t, err)
+        assert.True(t, strings.Contains(err.Error(), "json: unsupported value: NaN or ±Infinite"))
+    }
 }
\ No newline at end of file
diff --git a/encoder/encoder_native.go b/encoder/encoder_native.go
index 7067db66a..737ec001f 100644
--- a/encoder/encoder_native.go
+++ b/encoder/encoder_native.go
@@ -70,6 +70,9 @@ const (
 
     // CompatibleWithStd is used to be compatible with std encoder.
     CompatibleWithStd Options = encoder.CompatibleWithStd
+
+    // Encode Infinity or Nan float into `null`, instead of returning an error.
+    EncodeNullForInfOrNan Options = encoder.EncodeNullForInfOrNan
 )
 
 
diff --git a/internal/encoder/alg/opts.go b/internal/encoder/alg/opts.go
index 3dad938b5..c19e2de4e 100644
--- a/internal/encoder/alg/opts.go
+++ b/internal/encoder/alg/opts.go
@@ -25,6 +25,7 @@ const (
     BitValidateString
     BitNoValidateJSONMarshaler
     BitNoEncoderNewline 
+    BitEncodeNullForInfOrNan 
 	
     BitPointerValue = 63
 )
diff --git a/internal/encoder/encoder.go b/internal/encoder/encoder.go
index 4cd03f803..1efbcca5c 100644
--- a/internal/encoder/encoder.go
+++ b/internal/encoder/encoder.go
@@ -53,7 +53,8 @@ const (
     NoQuoteTextMarshaler Options = 1 << alg.BitNoQuoteTextMarshaler
 
     // NoNullSliceOrMap indicates all empty Array or Object are encoded as '[]' or '{}',
-    // instead of 'null'
+    // instead of 'null'. 
+    // NOTE: The priority of this option is lower than json tag `omitempty`.
     NoNullSliceOrMap     Options = 1 << alg.BitNoNullSliceOrMap
 
     // ValidateString indicates that encoder should validate the input string
@@ -69,6 +70,9 @@ const (
   
     // CompatibleWithStd is used to be compatible with std encoder.
     CompatibleWithStd Options = SortMapKeys | EscapeHTML | CompactMarshaler
+
+    // Encode Infinity or Nan float into `null`, instead of returning an error.
+    EncodeNullForInfOrNan Options = 1 << alg.BitEncodeNullForInfOrNan
 )
 
 // Encoder represents a specific set of encoder configurations.
diff --git a/internal/encoder/vm/vm.go b/internal/encoder/vm/vm.go
index bde62521c..b56145903 100644
--- a/internal/encoder/vm/vm.go
+++ b/internal/encoder/vm/vm.go
@@ -159,12 +159,20 @@ func Execute(b *[]byte, p unsafe.Pointer, s *vars.Stack, flags uint64, prog *ir.
 		case ir.OP_f32:
 			v := *(*float32)(p)
 			if math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) {
+				if flags&(1<<alg.BitEncodeNullForInfOrNan) != 0 {
+					buf = append(buf, 'n', 'u', 'l', 'l')
+					continue
+				}
 				return vars.ERR_nan_or_infinite
 			}
 			buf = alg.F32toa(buf, v)
 		case ir.OP_f64:
 			v := *(*float64)(p)
 			if math.IsNaN(v) || math.IsInf(v, 0) {
+				if flags&(1<<alg.BitEncodeNullForInfOrNan) != 0 {
+					buf = append(buf, 'n', 'u', 'l', 'l')
+					continue
+				}
 				return vars.ERR_nan_or_infinite
 			}
 			buf = alg.F64toa(buf, v)
diff --git a/internal/encoder/x86/assembler_regabi_amd64.go b/internal/encoder/x86/assembler_regabi_amd64.go
index 347818b18..470278e86 100644
--- a/internal/encoder/x86/assembler_regabi_amd64.go
+++ b/internal/encoder/x86/assembler_regabi_amd64.go
@@ -869,12 +869,18 @@ func (self *Assembler) _asm_OP_f32(_ *ir.Instr) {
 	self.Emit("MOVL", jit.Ptr(_SP_p, 0), _AX)  // MOVL     (SP.p), AX
 	self.Emit("ANDL", jit.Imm(_FM_exp32), _AX) // ANDL     $_FM_exp32, AX
 	self.Emit("XORL", jit.Imm(_FM_exp32), _AX) // XORL     $_FM_exp32, AX
-	self.Sjmp("JZ", _LB_error_nan_or_infinite) // JZ       _error_nan_or_infinite
+	self.Sjmp("JNZ",  "_encode_normal_f32_{n}")// JNZ      _encode_normal_f32_{n}
+	self.Emit("BTQ", jit.Imm(alg.BitEncodeNullForInfOrNan), _ARG_fv) // BTQ ${BitEncodeNullForInfOrNan}, fv
+	self.Sjmp("JNC", _LB_error_nan_or_infinite) // JNC     _error_nan_or_infinite
+	self._asm_OP_null(nil)
+	self.Sjmp("JMP", "_encode_f32_end_{n}")    // JMP      _encode_f32_end_{n}
+	self.Link("_encode_normal_f32_{n}")
 	self.save_c()                              // SAVE     $C_regs
 	self.rbuf_di()                             // MOVQ     RP, DI
 	self.Emit("MOVSS", jit.Ptr(_SP_p, 0), _X0) // MOVSS    (SP.p), X0
-	self.call_c(_F_f32toa)                     // CALL_C   f64toa
+	self.call_c(_F_f32toa)                     // CALL_C   f32toa
 	self.Emit("ADDQ", _AX, _RL)                // ADDQ     AX, RL
+	self.Link("_encode_f32_end_{n}")
 }
 
 func (self *Assembler) _asm_OP_f64(_ *ir.Instr) {
@@ -883,12 +889,18 @@ func (self *Assembler) _asm_OP_f64(_ *ir.Instr) {
 	self.Emit("MOVQ", jit.Imm(_FM_exp64), _CX) // MOVQ   $_FM_exp64, CX
 	self.Emit("ANDQ", _CX, _AX)                // ANDQ   CX, AX
 	self.Emit("XORQ", _CX, _AX)                // XORQ   CX, AX
-	self.Sjmp("JZ", _LB_error_nan_or_infinite) // JZ     _error_nan_or_infinite
+	self.Sjmp("JNZ",  "_encode_normal_f64_{n}")// JNZ    _encode_normal_f64_{n}
+	self.Emit("BTQ", jit.Imm(alg.BitEncodeNullForInfOrNan), _ARG_fv) // BTQ ${BitEncodeNullForInfOrNan}, fv
+	self.Sjmp("JNC", _LB_error_nan_or_infinite)// JNC    _error_nan_or_infinite
+	self._asm_OP_null(nil)
+	self.Sjmp("JMP", "_encode_f64_end_{n}")    // JMP    _encode_f64_end_{n}
+	self.Link("_encode_normal_f64_{n}")
 	self.save_c()                              // SAVE   $C_regs
 	self.rbuf_di()                             // MOVQ   RP, DI
 	self.Emit("MOVSD", jit.Ptr(_SP_p, 0), _X0) // MOVSD  (SP.p), X0
 	self.call_c(_F_f64toa)                     // CALL_C f64toa
 	self.Emit("ADDQ", _AX, _RL)                // ADDQ   AX, RL
+	self.Link("_encode_f64_end_{n}")
 }
 
 func (self *Assembler) _asm_OP_str(_ *ir.Instr) {
diff --git a/sonic.go b/sonic.go
index a92981c41..96b5cb1c4 100644
--- a/sonic.go
+++ b/sonic.go
@@ -65,6 +65,9 @@ func (cfg Config) Froze() API {
     if cfg.NoEncoderNewline {
         api.encoderOpts |= encoder.NoEncoderNewline
     }
+    if cfg.EncodeNullForInfOrNan {
+        api.encoderOpts |= encoder.EncodeNullForInfOrNan
+    }
 
     // configure decoder options:
     if cfg.UseInt64 {