encoding/cbor: cleanup default temp allocator

4 files changed, 77 insertions, 35 deletions

diff --git a/core/encoding/cbor/cbor.odin b/core/encoding/cbor/cbor.odin
index ddbd53c8d..9df4dfa51 100644
--- a/core/encoding/cbor/cbor.odin
+++ b/core/encoding/cbor/cbor.odin
@@ -238,6 +238,7 @@ negative_u64_to_int :: #force_inline proc(u: Negative_U64) -> i128 {
 // Utility for converting between the different errors when they are subsets of the other.
 err_conv :: proc {
 	encode_to_marshal_err,
+	encode_to_marshal_err_p2,
 	decode_to_unmarshal_err,
 	decode_to_unmarshal_err_p,
 	decode_to_unmarshal_err_p2,
@@ -253,6 +254,10 @@ encode_to_marshal_err :: #force_inline proc(err: Encode_Error) -> Marshal_Error
 	}
 }
 
+encode_to_marshal_err_p2 :: #force_inline proc(v: $T, v2: $T2, err: Encode_Error) -> (T, T2, Marshal_Error) {
+	return v, v2, err_conv(err)
+}
+
 decode_to_unmarshal_err :: #force_inline proc(err: Decode_Error) -> Unmarshal_Error {
 	switch e in err {
 	case nil:                 return nil
diff --git a/core/encoding/cbor/coding.odin b/core/encoding/cbor/coding.odin
index 5d99aa6d2..1e77a35c8 100644
--- a/core/encoding/cbor/coding.odin
+++ b/core/encoding/cbor/coding.odin
@@ -4,6 +4,7 @@ import "core:bytes"
 import "core:encoding/endian"
 import "core:intrinsics"
 import "core:io"
+import "core:runtime"
 import "core:slice"
 import "core:strings"
 
@@ -54,6 +55,9 @@ Decoder_Flag :: enum {
 	
 	// Makes the decoder shrink of excess capacity from allocated buffers/containers before returning.
 	Shrink_Excess,
+
+	// Internal flag to do initialization.
+	_In_Progress,
 }
 
 Decoder_Flags :: bit_set[Decoder_Flag]
@@ -117,9 +121,8 @@ decode_from_decoder :: proc(d: Decoder, allocator := context.allocator) -> (v: V
 	context.allocator = allocator
 	
 	d := d
-	if d.max_pre_alloc <= 0 {
-		d.max_pre_alloc = DEFAULT_MAX_PRE_ALLOC
-	}
+
+	DECODE_PROGRESS_GUARD(&d)
 
 	v, err = _decode_from_decoder(d)
 	// Normal EOF does not exist here, we try to read the exact amount that is said to be provided.
@@ -225,21 +228,9 @@ encode_into_writer :: proc(w: io.Writer, v: Value, flags := ENCODE_SMALL) -> Enc
 // See the docs on the proc group `encode_into` for more info.
 encode_into_encoder :: proc(e: Encoder, v: Value) -> Encode_Error {
 	e := e
-	
-	outer: bool
-	defer if outer {
-		e.flags &~= {._In_Progress}
-	}
-
-	if ._In_Progress not_in e.flags {
-		outer = true
-		e.flags |= {._In_Progress}
-
-		if .Self_Described_CBOR in e.flags {
-			_encode_u64(e, TAG_SELF_DESCRIBED_CBOR, .Tag) or_return
-		}
-	}
 
+	ENCODE_PROGRESS_GUARD(&e) or_return
+	
 	switch v_spec in v {
 	case u8:           return _encode_u8(e.writer, v_spec, .Unsigned)
 	case u16:          return _encode_u16(e, v_spec, .Unsigned)
@@ -265,6 +256,66 @@ encode_into_encoder :: proc(e: Encoder, v: Value) -> Encode_Error {
 	}
 }
 
+@(deferred_in_out=_decode_progress_end)
+DECODE_PROGRESS_GUARD :: proc(d: ^Decoder) -> (is_begin: bool, tmp: runtime.Arena_Temp) {
+	if ._In_Progress in d.flags {
+		return
+	}
+	is_begin = true
+	
+	incl_elem(&d.flags, Decoder_Flag._In_Progress)
+
+	if context.allocator != context.temp_allocator {
+		tmp = runtime.default_temp_allocator_temp_begin()
+	}
+
+	if d.max_pre_alloc <= 0 {
+		d.max_pre_alloc = DEFAULT_MAX_PRE_ALLOC
+	}
+
+	return
+}
+
+_decode_progress_end :: proc(d: ^Decoder, is_begin: bool, tmp: runtime.Arena_Temp) {
+	if !is_begin {
+		return
+	}
+
+	excl_elem(&d.flags, Decoder_Flag._In_Progress)
+
+	runtime.default_temp_allocator_temp_end(tmp)
+}
+
+@(deferred_in_out=_encode_progress_end)
+ENCODE_PROGRESS_GUARD :: proc(e: ^Encoder) -> (is_begin: bool, tmp: runtime.Arena_Temp, err: Encode_Error) {
+	if ._In_Progress in e.flags {
+		return
+	}
+	is_begin = true
+
+	incl_elem(&e.flags, Encoder_Flag._In_Progress)
+
+	if context.allocator != context.temp_allocator {
+		tmp = runtime.default_temp_allocator_temp_begin()
+	}
+
+	if .Self_Described_CBOR in e.flags {
+		_encode_u64(e^, TAG_SELF_DESCRIBED_CBOR, .Tag) or_return
+	}
+
+	return
+}
+
+_encode_progress_end :: proc(e: ^Encoder, is_begin: bool, tmp: runtime.Arena_Temp, err: Encode_Error) {
+	if !is_begin || err != nil {
+		return
+	}
+
+	excl_elem(&e.flags, Encoder_Flag._In_Progress)
+
+	runtime.default_temp_allocator_temp_end(tmp)
+}
+
 _decode_header :: proc(r: io.Reader) -> (hdr: Header, err: io.Error) {
 	buf: [1]byte = ---
 	io.read_full(r, buf[:]) or_return
@@ -514,7 +565,7 @@ _decode_map :: proc(d: Decoder, add: Add) -> (v: Map, err: Decode_Error) {
 			return nil, kerr
 		} 
 
-		value := decode_from_decoder(d) or_return
+		value := _decode_from_decoder(d) or_return
 
 		append(&items, Map_Entry{
 			key   = key,
diff --git a/core/encoding/cbor/marshal.odin b/core/encoding/cbor/marshal.odin
index a5d5efb3e..898371adf 100644
--- a/core/encoding/cbor/marshal.odin
+++ b/core/encoding/cbor/marshal.odin
@@ -77,21 +77,8 @@ marshal_into_writer :: proc(w: io.Writer, v: any, flags := ENCODE_SMALL) -> Mars
 // See docs on the `marshal_into` proc group for more info.
 marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 	e := e
-	
-	init: bool
-	defer if init {
-		e.flags &~= {._In_Progress}
-	}
-	
-	// If not in progress we do initialization and set in progress.
-	if ._In_Progress not_in e.flags {
-		init = true
-		e.flags |= {._In_Progress}
 
-		if .Self_Described_CBOR in e.flags {
-			err_conv(_encode_u64(e, TAG_SELF_DESCRIBED_CBOR, .Tag)) or_return
-		}
-	}
+	err_conv(ENCODE_PROGRESS_GUARD(&e)) or_return
 
 	if v == nil {
 		return _encode_nil(e.writer)
diff --git a/core/encoding/cbor/unmarshal.odin b/core/encoding/cbor/unmarshal.odin
index dea4b749c..c7de2d87a 100644
--- a/core/encoding/cbor/unmarshal.odin
+++ b/core/encoding/cbor/unmarshal.odin
@@ -52,9 +52,8 @@ unmarshal_from_string :: proc(s: string, ptr: ^$T, flags := Decoder_Flags{}, all
 
 unmarshal_from_decoder :: proc(d: Decoder, ptr: ^$T, allocator := context.allocator) -> (err: Unmarshal_Error) {
 	d := d
-	if d.max_pre_alloc <= 0 {
-		d.max_pre_alloc = DEFAULT_MAX_PRE_ALLOC
-	}
+
+	DECODE_PROGRESS_GUARD(&d)
 
 	err = _unmarshal_any_ptr(d, ptr, allocator=allocator)