Merge pull request #5969 from Kelimion/faster_big_itoa

Faster `big.itoa`.

Author: Kelimion

core/math/big/common.odin

@@ -223,9 +223,15 @@ when MATH_BIG_FORCE_64_BIT || (!MATH_BIG_FORCE_32_BIT && size_of(rawptr) == 8) {
 	*/
 	DIGIT        :: distinct u64
 	_WORD        :: distinct u128
+	// Base 10 extraction constants
+	ITOA_DIVISOR :: DIGIT(1_000_000_000_000_000_000)
+	ITOA_COUNT   :: 18
 } else {
 	DIGIT        :: distinct u32
 	_WORD        :: distinct u64
+	// Base 10 extraction constants
+	ITOA_DIVISOR :: DIGIT(100_000_000)
+	ITOA_COUNT   :: 8
 }
 #assert(size_of(_WORD) == 2 * size_of(DIGIT))
 

core/math/big/radix.odin

@@ -185,16 +185,15 @@ int_itoa_raw :: proc(a: ^Int, radix: i8, buffer: []u8, size := int(-1), zero_ter
 	/*
 		Fast path for radixes that are a power of two.
 	*/
+	count := count_bits(a) or_return
+
 	if is_power_of_two(int(radix)) {
 		if zero_terminate {
 			available -= 1
 			buffer[available] = 0
 		}
 
-		shift, count: int
-		// mask  := _WORD(radix - 1);
-		shift, err = log(DIGIT(radix), 2)
-		count, err = count_bits(a)
+		shift := log(DIGIT(radix), 2) or_return
 		digit: _WORD
 
 		for offset := 0; offset < count; offset += shift {
@@ -224,7 +223,17 @@ int_itoa_raw :: proc(a: ^Int, radix: i8, buffer: []u8, size := int(-1), zero_ter
 		return written, nil
 	}
 
-	return _itoa_raw_full(a, radix, buffer, zero_terminate)
+	// NOTE(Jeroen): The new method is faster for an `Int` up to ~32768 bits in size with optimizations.
+	//               At `.None` or `.Minimal`, it appears to always be faster.
+	//               If we optimize `itoa` further, this needs to be evaluated.
+	itoa_method := _itoa_raw_full
+
+	when ODIN_OPTIMIZATION_MODE >= .Size {
+		if count >= 32768 {
+			itoa_method = _itoa_raw_old
+		}
+	}
+	return itoa_method(a, radix, buffer, zero_terminate)
 }
 
 itoa :: proc{int_itoa_string, int_itoa_raw}
@@ -601,14 +610,89 @@ RADIX_TABLE_REVERSE_SIZE :: 80
 	Stores a bignum as a ASCII string in a given radix (2..64)
 	The buffer must be appropriately sized. This routine doesn't check.
 */
+
 _itoa_raw_full :: proc(a: ^Int, radix: i8, buffer: []u8, zero_terminate := false, allocator := context.allocator) -> (written: int, err: Error) {
 	assert_if_nil(a)
 	context.allocator = allocator
 
-	temp, denominator := &Int{}, &Int{}
+	// Calculate largest radix^n that fits within _DIGIT_BITS
+	divisor     := ITOA_DIVISOR
+	digit_count := ITOA_COUNT
+	_radix      := DIGIT(radix)
+
+	if radix != 10 {
+		i := _WORD(1)
+		digit_count = -1
+		for i < _WORD(1 << _DIGIT_BITS) {
+			divisor = DIGIT(i)
+			i *= _WORD(radix)
+			digit_count += 1
+		}
+	}
 
-	internal_copy(temp, a)           or_return
-	internal_set(denominator, radix) or_return
+	temp := &Int{}
+	internal_copy(temp, a) or_return
+	defer internal_destroy(temp)
+
+	available := len(buffer)
+	if zero_terminate {
+		available -= 1
+		buffer[available] = 0
+	}
+
+	if a.sign == .Negative {
+		temp.sign = .Zero_or_Positive
+	}
+
+	remainder: DIGIT
+	for {
+		if remainder, err = internal_divmod(temp, temp, divisor); err != nil {
+			return len(buffer) - available, err
+		}
+
+		count := digit_count
+		for available > 0 && count > 0 {
+			available -= 1
+			buffer[available] = RADIX_TABLE[remainder % _radix]
+			remainder /= _radix
+			count -= 1
+		}
+
+		if temp.used == 0 {
+			break
+		}
+	}
+
+	// Remove leading zero if we ended up with one.
+	if buffer[available] == '0' {
+		available += 1
+	}
+
+	if a.sign == .Negative {
+		available -= 1
+		buffer[available] = '-'
+	}
+
+	/*
+		If we overestimated the size, we need to move the buffer left.
+	*/
+	written = len(buffer) - available
+	if written < len(buffer) {
+		diff := len(buffer) - written
+		mem.copy(&buffer[0], &buffer[diff], written)
+	}
+	return written, nil
+}
+
+// Old internal digit extraction procedure.
+// We're keeping this around as ground truth for the tests.
+_itoa_raw_old :: proc(a: ^Int, radix: i8, buffer: []u8, zero_terminate := false, allocator := context.allocator) -> (written: int, err: Error) {
+	assert_if_nil(a)
+	context.allocator = allocator
+
+	temp := &Int{}
+	internal_copy(temp, a) or_return
+	defer internal_destroy(temp)
 
 	available := len(buffer)
 	if zero_terminate {
@@ -623,7 +707,6 @@ _itoa_raw_full :: proc(a: ^Int, radix: i8, buffer: []u8, zero_terminate := false
 	remainder: DIGIT
 	for {
 		if remainder, err = #force_inline internal_divmod(temp, temp, DIGIT(radix)); err != nil {
-			internal_destroy(temp, denominator)
 			return len(buffer) - available, err
 		}
 		available -= 1
@@ -638,8 +721,6 @@ _itoa_raw_full :: proc(a: ^Int, radix: i8, buffer: []u8, zero_terminate := false
 		buffer[available] = '-'
 	}
 
-	internal_destroy(temp, denominator)
-
 	/*
 		If we overestimated the size, we need to move the buffer left.
 	*/

tests/core/math/big/test_core_math_big.odin

@@ -287,4 +287,46 @@ atoi :: proc(t: ^testing.T, i: ^big.Int, a: string, loc := #caller_location) ->
 	err := big.atoi(i, a, 16)
 	testing.expect(t, err == nil, loc=loc)
 	return err == nil
+}
+
+@(test)
+test_itoa :: proc(t: ^testing.T) {
+	a := &big.Int{}
+	big.random(a, 2048)
+	defer big.destroy(a)
+
+	for radix in 2..=64 {
+		if big.is_power_of_two(radix) {
+			// Powers of two are trivial, and are handled before `_itoa_raw_*` is called.
+			continue
+		}
+
+		size, _ := big.radix_size(a, i8(radix), false)
+		buffer_old := make([]u8, size)
+		defer delete(buffer_old)
+		buffer_new := make([]u8, size)
+		defer delete(buffer_new)
+
+		written_old, _ := big._itoa_raw_old (a, i8(radix), buffer_old, false)
+		written_new, _ := big._itoa_raw_full(a, i8(radix), buffer_new, false)
+
+		str_old := string(buffer_old[:written_old])
+		str_new := string(buffer_new[:written_new])
+		testing.expect_value(t, str_new, str_old)
+	}
+
+	// Also test a number with a large number of zeroes
+	big.set(a, "2970714761494550000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000802525522395693895558562897961119110387707542077460459880227570865486047631557732177235787527971863645406120285117781450154113859156752194121206131440514109132606823127467068869589613665129498148285292867292641704871893467328665051712596763187306247339023362481")
+	size, _ := big.radix_size(a, 10, false)
+	buffer_old := make([]u8, size)
+	defer delete(buffer_old)
+	buffer_new := make([]u8, size)
+	defer delete(buffer_new)
+
+	written_old, _ := big._itoa_raw_old (a, 10, buffer_old, false)
+	written_new, _ := big._itoa_raw_full(a, 10, buffer_new, false)
+
+	str_old := string(buffer_old[:written_old])
+	str_new := string(buffer_new[:written_new])
+	testing.expect_value(t, str_new, str_old)
 }