encoding/cbor: make temp allocations more explicit

core/encoding/cbor/coding.odin

@@ -26,9 +26,6 @@ Encoder_Flag :: enum {
 	// NOTE: In order to do this, all keys of a map have to be pre-computed, sorted, and
 	// then written, this involves temporary allocations for the keys and a copy of the map itself.
 	Deterministic_Map_Sorting, 
-	
-	// Internal flag to do initialization.
-	_In_Progress,
 }
 
 Encoder_Flags :: bit_set[Encoder_Flag]
@@ -40,8 +37,9 @@ ENCODE_FULLY_DETERMINISTIC :: Encoder_Flags{.Deterministic_Int_Size, .Determinis
 ENCODE_SMALL :: Encoder_Flags{.Deterministic_Int_Size, .Deterministic_Float_Size}
 
 Encoder :: struct {
-	flags:  Encoder_Flags,
-	writer: io.Writer,
+	flags:          Encoder_Flags,
+	writer:         io.Writer,
+	temp_allocator: runtime.Allocator,
 }
 
 Decoder_Flag :: enum {
@@ -56,9 +54,6 @@ 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]
@@ -122,7 +117,9 @@ decode_from_decoder :: proc(d: Decoder, allocator := context.allocator) -> (v: V
 	
 	d := d
 
-	_DECODE_PROGRESS_GUARD(&d)
+	if d.max_pre_alloc <= 0 {
+		d.max_pre_alloc = DEFAULT_MAX_PRE_ALLOC
+	}
 
 	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.
@@ -191,7 +188,7 @@ have to be precomputed, sorted and only then written to the output.
 
 Empty flags will do nothing extra to the value.
 
-The allocations for the `.Deterministic_Map_Sorting` flag are done using the `context.temp_allocator`
+The allocations for the `.Deterministic_Map_Sorting` flag are done using the given temp_allocator.
 but are followed by the necessary `delete` and `free` calls if the allocator supports them.
 This is helpful when the CBOR size is so big that you don't want to collect all the temporary
 allocations until the end.
@@ -206,22 +203,22 @@ encode :: encode_into
 
 // Encodes the CBOR value into binary CBOR allocated on the given allocator.
 // See the docs on the proc group `encode_into` for more info.
-encode_into_bytes :: proc(v: Value, flags := ENCODE_SMALL, allocator := context.allocator) -> (data: []byte, err: Encode_Error) {
+encode_into_bytes :: proc(v: Value, flags := ENCODE_SMALL, allocator := context.allocator, temp_allocator := context.temp_allocator) -> (data: []byte, err: Encode_Error) {
 	b := strings.builder_make(allocator) or_return
-	encode_into_builder(&b, v, flags) or_return
+	encode_into_builder(&b, v, flags, temp_allocator) or_return
 	return b.buf[:], nil
 }
 
 // Encodes the CBOR value into binary CBOR written to the given builder.
 // See the docs on the proc group `encode_into` for more info.
-encode_into_builder :: proc(b: ^strings.Builder, v: Value, flags := ENCODE_SMALL) -> Encode_Error {
-	return encode_into_writer(strings.to_stream(b), v, flags)
+encode_into_builder :: proc(b: ^strings.Builder, v: Value, flags := ENCODE_SMALL, temp_allocator := context.temp_allocator) -> Encode_Error {
+	return encode_into_writer(strings.to_stream(b), v, flags, temp_allocator)
 }
 
 // Encodes the CBOR value into binary CBOR written to the given writer.
 // See the docs on the proc group `encode_into` for more info.
-encode_into_writer :: proc(w: io.Writer, v: Value, flags := ENCODE_SMALL) -> Encode_Error {
-	return encode_into_encoder(Encoder{flags, w}, v)
+encode_into_writer :: proc(w: io.Writer, v: Value, flags := ENCODE_SMALL, temp_allocator := context.temp_allocator) -> Encode_Error {
+	return encode_into_encoder(Encoder{flags, w, temp_allocator}, v)
 }
 
 // Encodes the CBOR value into binary CBOR written to the given encoder.
@@ -229,8 +226,15 @@ encode_into_writer :: proc(w: io.Writer, v: Value, flags := ENCODE_SMALL) -> Enc
 encode_into_encoder :: proc(e: Encoder, v: Value) -> Encode_Error {
 	e := e
 
-	_ENCODE_PROGRESS_GUARD(&e) or_return
-	
+	if e.temp_allocator.procedure == nil {
+		e.temp_allocator = context.temp_allocator
+	}
+
+	if .Self_Described_CBOR in e.flags {
+		_encode_u64(e, TAG_SELF_DESCRIBED_CBOR, .Tag) or_return
+		e.flags &~= { .Self_Described_CBOR }
+	}
+
 	switch v_spec in v {
 	case u8:           return _encode_u8(e.writer, v_spec, .Unsigned)
 	case u16:          return _encode_u16(e, v_spec, .Unsigned)
@@ -256,66 +260,6 @@ 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
-	
-	d.flags |= { ._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
-	}
-
-	d.flags &~= { ._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
-
-	e.flags |= { ._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
-	}
-
-	e.flags &~= { ._In_Progress }
-
-	runtime.default_temp_allocator_temp_end(tmp)
-}
-
 _decode_header :: proc(r: io.Reader) -> (hdr: Header, err: io.Error) {
 	hdr = Header(_decode_u8(r) or_return)
 	return
@@ -602,13 +546,13 @@ _encode_map :: proc(e: Encoder, m: Map) -> (err: Encode_Error) {
 		entry:       Map_Entry,
 	}
 
-	entries := make([]Map_Entry_With_Key, len(m), context.temp_allocator) or_return
-	defer delete(entries, context.temp_allocator)
+	entries := make([]Map_Entry_With_Key, len(m), e.temp_allocator) or_return
+	defer delete(entries, e.temp_allocator)
 
 	for &entry, i in entries {
 		entry.entry = m[i]
 
-		buf := strings.builder_make(context.temp_allocator) or_return
+		buf := strings.builder_make(e.temp_allocator) or_return
 		
 		ke := e
 		ke.writer = strings.to_stream(&buf)
@@ -624,7 +568,7 @@ _encode_map :: proc(e: Encoder, m: Map) -> (err: Encode_Error) {
 
 	for entry in entries {
 		io.write_full(e.writer, entry.encoded_key) or_return
-		delete(entry.encoded_key, context.temp_allocator)
+		delete(entry.encoded_key, e.temp_allocator)
 
 		encode(e, entry.entry.value) or_return
 	}

core/encoding/cbor/doc.odin

@@ -4,23 +4,21 @@ Also provided are conversion to and from JSON and the CBOR diagnostic format.
 
 **Allocations:**
 
-In general, when in the following table it says allocations are done on the `context.temp_allocator`, these allocations
+In general, when in the following table it says allocations are done on the `temp_allocator`, these allocations
 are still attempted to be deallocated.
-This allows you to use an allocator with freeing implemented as the `context.temp_allocator` which is handy with big CBOR.
+This allows you to use an allocator with freeing implemented as the `temp_allocator` which is handy with big CBOR.
 
-If you use the default `context.temp_allocator` it will be returned back to its state when the process (en/decoding, (un)marshal) started.
-
-- *Encoding*:  If the `.Deterministic_Map_Sorting` flag is set on the encoder, this allocates on `context.temp_allocator`
+- *Encoding*:  If the `.Deterministic_Map_Sorting` flag is set on the encoder, this allocates on the given `temp_allocator`
                some space for the keys of maps in order to sort them and then write them.
                Other than that there are no allocations (only for the final bytes if you use `cbor.encode_into_bytes`.
 
 - *Decoding*:  Allocates everything on the given allocator and input given can be deleted after decoding.
-               *No* allocations are done on the `context.temp_allocator`.
+               *No* temporary allocations are done.
 
 - *Marshal*:   Same allocation strategy as encoding.
 
 - *Unmarshal*: Allocates everything on the given allocator and input given can be deleted after unmarshalling.
-               Some temporary allocations are done on the `context.temp_allocator`.
+               Some temporary allocations are done on the given `temp_allocator`.
 
 **Determinism:**
 

core/encoding/cbor/marshal.odin

@@ -29,7 +29,7 @@ have to be precomputed, sorted and only then written to the output.
 
 Empty flags will do nothing extra to the value.
 
-The allocations for the `.Deterministic_Map_Sorting` flag are done using the `context.temp_allocator`
+The allocations for the `.Deterministic_Map_Sorting` flag are done using the given `temp_allocator`.
 but are followed by the necessary `delete` and `free` calls if the allocator supports them.
 This is helpful when the CBOR size is so big that you don't want to collect all the temporary
 allocations until the end.
@@ -45,7 +45,7 @@ marshal :: marshal_into
 
 // Marshals the given value into a CBOR byte stream (allocated using the given allocator).
 // See docs on the `marshal_into` proc group for more info.
-marshal_into_bytes :: proc(v: any, flags := ENCODE_SMALL, allocator := context.allocator) -> (bytes: []byte, err: Marshal_Error) {
+marshal_into_bytes :: proc(v: any, flags := ENCODE_SMALL, allocator := context.allocator, temp_allocator := context.temp_allocator) -> (bytes: []byte, err: Marshal_Error) {
 	b, alloc_err := strings.builder_make(allocator)
  	// The builder as a stream also returns .EOF if it ran out of memory so this is consistent.
 	if alloc_err != nil {
@@ -54,7 +54,7 @@ marshal_into_bytes :: proc(v: any, flags := ENCODE_SMALL, allocator := context.a
 
 	defer if err != nil { strings.builder_destroy(&b) }
 
-	if err = marshal_into_builder(&b, v, flags); err != nil {
+	if err = marshal_into_builder(&b, v, flags, temp_allocator); err != nil {
 		return
 	}
 
@@ -63,14 +63,14 @@ marshal_into_bytes :: proc(v: any, flags := ENCODE_SMALL, allocator := context.a
 
 // Marshals the given value into a CBOR byte stream written to the given builder.
 // See docs on the `marshal_into` proc group for more info.
-marshal_into_builder :: proc(b: ^strings.Builder, v: any, flags := ENCODE_SMALL) -> Marshal_Error {
-	return marshal_into_writer(strings.to_writer(b), v, flags)
+marshal_into_builder :: proc(b: ^strings.Builder, v: any, flags := ENCODE_SMALL, temp_allocator := context.temp_allocator) -> Marshal_Error {
+	return marshal_into_writer(strings.to_writer(b), v, flags, temp_allocator)
 }
 
 // Marshals the given value into a CBOR byte stream written to the given writer.
 // See docs on the `marshal_into` proc group for more info.
-marshal_into_writer :: proc(w: io.Writer, v: any, flags := ENCODE_SMALL) -> Marshal_Error {
-	encoder := Encoder{flags, w}
+marshal_into_writer :: proc(w: io.Writer, v: any, flags := ENCODE_SMALL, temp_allocator := context.temp_allocator) -> Marshal_Error {
+	encoder := Encoder{flags, w, temp_allocator}
 	return marshal_into_encoder(encoder, v)
 }
 
@@ -79,7 +79,14 @@ marshal_into_writer :: proc(w: io.Writer, v: any, flags := ENCODE_SMALL) -> Mars
 marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 	e := e
 
-	err_conv(_ENCODE_PROGRESS_GUARD(&e)) or_return
+	if e.temp_allocator.procedure == nil {
+		e.temp_allocator = context.temp_allocator
+	}
+
+	if .Self_Described_CBOR in e.flags {
+		err_conv(_encode_u64(e, TAG_SELF_DESCRIBED_CBOR, .Tag)) or_return
+		e.flags &~= { .Self_Described_CBOR }
+	}
 
 	if v == nil {
 		return _encode_nil(e.writer)
@@ -321,7 +328,7 @@ marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 
 			switch info.key.id {
 			case string:
-				entries := make([dynamic]Encoded_Entry_Fast(^[]byte), 0, map_cap, context.temp_allocator) or_return
+				entries := make([dynamic]Encoded_Entry_Fast(^[]byte), 0, map_cap, e.temp_allocator) or_return
 				defer delete(entries)
 
 				for bucket_index in 0..<map_cap {
@@ -355,7 +362,7 @@ marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 				return
 
 			case cstring:
-				entries := make([dynamic]Encoded_Entry_Fast(^cstring), 0, map_cap, context.temp_allocator) or_return
+				entries := make([dynamic]Encoded_Entry_Fast(^cstring), 0, map_cap, e.temp_allocator) or_return
 				defer delete(entries)
 
 				for bucket_index in 0..<map_cap {
@@ -391,15 +398,15 @@ marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 				return
 
 			case:
-				entries := make([dynamic]Encoded_Entry, 0, map_cap, context.temp_allocator) or_return
+				entries := make([dynamic]Encoded_Entry, 0, map_cap, e.temp_allocator) or_return
 				defer delete(entries)
 
 				for bucket_index in 0..<map_cap {
 					runtime.map_hash_is_valid(hs[bucket_index]) or_continue
 
 					key := rawptr(runtime.map_cell_index_dynamic(ks, info.map_info.ks, bucket_index))
-					key_builder := strings.builder_make(0, 8, context.temp_allocator) or_return
-					marshal_into(Encoder{e.flags, strings.to_stream(&key_builder)}, any{ key, info.key.id }) or_return
+					key_builder := strings.builder_make(0, 8, e.temp_allocator) or_return
+					marshal_into(Encoder{e.flags, strings.to_stream(&key_builder), e.temp_allocator}, any{ key, info.key.id }) or_return
 					append(&entries, Encoded_Entry{ &key_builder.buf, bucket_index }) or_return
 				}
 
@@ -470,7 +477,7 @@ marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 				name:  string,
 				field: int,
 			}
-			entries := make([dynamic]Name, 0, n, context.temp_allocator) or_return
+			entries := make([dynamic]Name, 0, n, e.temp_allocator) or_return
 			defer delete(entries)
 
 			for name, i in info.names {
@@ -530,7 +537,7 @@ marshal_into_encoder :: proc(e: Encoder, v: any) -> (err: Marshal_Error) {
 		case reflect.Type_Info_Named:
 			err_conv(_encode_text(e, vt.name)) or_return
 		case:
-			builder := strings.builder_make(context.temp_allocator) or_return
+			builder := strings.builder_make(e.temp_allocator) or_return
 			defer strings.builder_destroy(&builder)
 			reflect.write_type(&builder, vti)
 			err_conv(_encode_text(e, strings.to_string(builder))) or_return

core/encoding/cbor/unmarshal.odin

@@ -13,7 +13,7 @@ import "core:unicode/utf8"
 Unmarshals the given CBOR into the given pointer using reflection.
 Types that require allocation are allocated using the given allocator.
 
-Some temporary allocations are done on the `context.temp_allocator`, but, if you want to,
+Some temporary allocations are done on the given `temp_allocator`, but, if you want to,
 this can be set to a "normal" allocator, because the necessary `delete` and `free` calls are still made.
 This is helpful when the CBOR size is so big that you don't want to collect all the temporary allocations until the end.
 
@@ -31,8 +31,8 @@ unmarshal :: proc {
 	unmarshal_from_string,
 }
 
-unmarshal_from_reader :: proc(r: io.Reader, ptr: ^$T, flags := Decoder_Flags{}, allocator := context.allocator) -> (err: Unmarshal_Error) {
-	err = unmarshal_from_decoder(Decoder{ DEFAULT_MAX_PRE_ALLOC, flags, r }, ptr, allocator=allocator)
+unmarshal_from_reader :: proc(r: io.Reader, ptr: ^$T, flags := Decoder_Flags{}, allocator := context.allocator, temp_allocator := context.temp_allocator) -> (err: Unmarshal_Error) {
+	err = unmarshal_from_decoder(Decoder{ DEFAULT_MAX_PRE_ALLOC, flags, r }, ptr, allocator, temp_allocator)
 
 	// Normal EOF does not exist here, we try to read the exact amount that is said to be provided.
 	if err == .EOF { err = .Unexpected_EOF }
@@ -40,23 +40,21 @@ unmarshal_from_reader :: proc(r: io.Reader, ptr: ^$T, flags := Decoder_Flags{},
 }
 
 // Unmarshals from a string, see docs on the proc group `Unmarshal` for more info.
-unmarshal_from_string :: proc(s: string, ptr: ^$T, flags := Decoder_Flags{}, allocator := context.allocator) -> (err: Unmarshal_Error) {
+unmarshal_from_string :: proc(s: string, ptr: ^$T, flags := Decoder_Flags{}, allocator := context.allocator, temp_allocator := context.temp_allocator) -> (err: Unmarshal_Error) {
 	sr: strings.Reader
 	r := strings.to_reader(&sr, s)
 
-	err = unmarshal_from_reader(r, ptr, flags, allocator)
+	err = unmarshal_from_reader(r, ptr, flags, allocator, temp_allocator)
 
 	// Normal EOF does not exist here, we try to read the exact amount that is said to be provided.
 	if err == .EOF { err = .Unexpected_EOF }
 	return
 }
 
-unmarshal_from_decoder :: proc(d: Decoder, ptr: ^$T, allocator := context.allocator) -> (err: Unmarshal_Error) {
+unmarshal_from_decoder :: proc(d: Decoder, ptr: ^$T, allocator := context.allocator, temp_allocator := context.temp_allocator) -> (err: Unmarshal_Error) {
 	d := d
 
-	_DECODE_PROGRESS_GUARD(&d)
-
-	err = _unmarshal_any_ptr(d, ptr, allocator=allocator)
+	err = _unmarshal_any_ptr(d, ptr, nil, allocator, temp_allocator)
 
 	// Normal EOF does not exist here, we try to read the exact amount that is said to be provided.
 	if err == .EOF { err = .Unexpected_EOF }
@@ -64,8 +62,9 @@ unmarshal_from_decoder :: proc(d: Decoder, ptr: ^$T, allocator := context.alloca
 
 }
 
-_unmarshal_any_ptr :: proc(d: Decoder, v: any, hdr: Maybe(Header) = nil, allocator := context.allocator) -> Unmarshal_Error {
+_unmarshal_any_ptr :: proc(d: Decoder, v: any, hdr: Maybe(Header) = nil, allocator := context.allocator, temp_allocator := context.temp_allocator) -> Unmarshal_Error {
 	context.allocator = allocator
+	context.temp_allocator = temp_allocator
 	v := v
 
 	if v == nil || v.id == nil {

tests/core/encoding/cbor/test_core_cbor.odin

@@ -855,7 +855,7 @@ expect_float :: proc(t: ^testing.T, encoded: string, expected: $T, loc := #calle
 
 buf: bytes.Buffer
 stream  := bytes.buffer_to_stream(&buf)
-encoder := cbor.Encoder{cbor.ENCODE_FULLY_DETERMINISTIC, stream}
+encoder := cbor.Encoder{cbor.ENCODE_FULLY_DETERMINISTIC, stream, {}}
 
 expect_encoding :: proc(t: ^testing.T, val: cbor.Value, encoded: string, loc := #caller_location) {
 	bytes.buffer_reset(&buf)