From afd2900439d84984d3f8f0e2f720258dbfbc2553 Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Thu, 31 Oct 2024 12:03:49 +0100
Subject: [PATCH] runtime: optimize findHead

This is similar to https://github.com/tinygo-org/tinygo/pull/3899, but
smaller and hopefully just as efficient.
---
 builder/sizes_test.go    |  6 +++---
 src/runtime/gc_blocks.go | 16 ++++++++++++++--
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/builder/sizes_test.go b/builder/sizes_test.go
index 8b669462a5..188af7e2b7 100644
--- a/builder/sizes_test.go
+++ b/builder/sizes_test.go
@@ -41,9 +41,9 @@ func TestBinarySize(t *testing.T) {
 	// This is a small number of very diverse targets that we want to test.
 	tests := []sizeTest{
 		// microcontrollers
-		{"hifive1b", "examples/echo", 4568, 280, 0, 2268},
-		{"microbit", "examples/serial", 2868, 388, 8, 2272},
-		{"wioterminal", "examples/pininterrupt", 6104, 1484, 116, 6832},
+		{"hifive1b", "examples/echo", 4588, 280, 0, 2268},
+		{"microbit", "examples/serial", 2888, 388, 8, 2272},
+		{"wioterminal", "examples/pininterrupt", 6124, 1484, 116, 6832},
 
 		// TODO: also check wasm. Right now this is difficult, because
 		// wasm binaries are run through wasm-opt and therefore the
diff --git a/src/runtime/gc_blocks.go b/src/runtime/gc_blocks.go
index 3c3862dbba..83bb2b7e65 100644
--- a/src/runtime/gc_blocks.go
+++ b/src/runtime/gc_blocks.go
@@ -123,8 +123,20 @@ func (b gcBlock) address() uintptr {
 // points to an allocated object. It returns the same block if this block
 // already points to the head.
 func (b gcBlock) findHead() gcBlock {
-	for b.state() == blockStateTail {
-		b--
+	const blockStateByteAllTails = uint8(blockStateTail<<(stateBits*3) | blockStateTail<<(stateBits*2) | blockStateTail<<(stateBits*1) | blockStateTail<<(stateBits*0))
+	for {
+		if b%blocksPerStateByte == blocksPerStateByte-1 {
+			stateByte := *(*uint8)(unsafe.Add(metadataStart, b/blocksPerStateByte))
+			if stateByte == blockStateByteAllTails {
+				b -= 4
+				continue
+			}
+		}
+		if b.state() == blockStateTail {
+			b--
+			continue
+		}
+		break
 	}
 	if gcAsserts {
 		if b.state() != blockStateHead && b.state() != blockStateMark {