[kripken]

Our main escape analysis pass, Heap2Local, does not handle overlapping
live ranges, things like this:

x = new Foo(42);
if (condition) x = new Foo(1337);
bar(x.value, x.value);

More than one allocation is being read from in those x.value operations,
so Heap2Local sees the uses are not "exclusive", and it gives up.

Handling live ranges, sort of like liveness of locals for register allocation /
local coalescing, is possible, but quite complex. However, the pattern
above can be optimized in another way, which this new pass handles:
when the reference is always read from (using struct.get, which is the
.value operations above), then we can just shift from using the reference
to using data:

x_value = (new Foo(42)).value;
if (condition) x_value = (new Foo(1337)).value;
bar(x_value, x_value);

We store the data in a new local x_value, and that is all we need.

That is, this pass handles a situation that is more restricted than
Heap2Local, in that it doesn't do complex analysis of control flow
and handling of ref.eq and such, but it is more general in that it
can handle overlapping live ranges in a trivial way.

Fixes #8641 , a Dart testcase, and I saw this helps real-world binaries from
Kotlin and Java too - this is a common pattern.

[kripken]

Feedback addressed. Also I made a change to optimize only when we see an allocation: if there is none, then we are just moving struct.gets around, and in fact moving them backward in time, which may not be faster. But if even one allocation exists, Heap2Local will remove it later, which is a win.

[kripken]

Also added a large comment explaining the live range overlap issue, as we discussed offline last week - why it is hard to fully solve in Heap2Local, hence motivating this pass.

[kripken]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new optimization pass, PropagateStructGets, which optimizes GC reference locals that are immediately consumed by struct.get operations by propagating the gets to earlier positions. This enables further optimizations like Heap2Local even in the presence of overlapping live ranges. The feedback recommends replacing std::unordered_map with std::map for sets and optimizedLocals to guarantee deterministic iteration and code generation. Additionally, it suggests extending the optimization to support struct.new_default (StructNewDefault) allocations to better handle GC-enabled languages like Kotlin, Java, and Dart.

[gemini-code-assist]

Similarly to sets, optimizedLocals is iterated over in optimize (line 454). To prevent any non-deterministic behavior during AST rewriting, please use std::map instead of std::unordered_map.

Suggested change

	// For each optimized local, maps each used field index to the new local
	// index that stores that field. Note that the inner map is sorted, both for
	// deterministic ordering when we create new locals, and also because the
	// number of used fields tends to be small (so hashing is overkill).
	std::unordered_map<Index, std::map<Index, Index>> optimizedLocals;
	// For each optimized local, maps each used field index to the new local
	// index that stores that field. Note that the inner map is sorted, both for
	// deterministic ordering when we create new locals, and also because the
	// number of used fields tends to be small (so hashing is overkill).
	std::map<Index, std::map<Index, Index>> optimizedLocals;

[kripken]

The iterations over this data structure can be in arbitrary order, so this is not needed.

[gemini-code-assist]

We can expand the optimization to also support struct.new_default (StructNewDefault) allocations, which are very common in GC-enabled languages like Kotlin, Java, and Dart. Let's check for both StructNew and StructNewDefault here.

Suggested change

	for (auto** setp : localSets) {
	auto* set = (*setp)->cast<LocalSet>();
	if (set->value->is<StructNew>()) {
	hasAllocation = true;
	break;
	}
	}
	for (auto** setp : localSets) {
	auto* set = (*setp)->cast<LocalSet>();
	if (set->value->is<StructNew>() || set->value->is<StructNewDefault>()) {
	hasAllocation = true;
	break;
	}
	}

[kripken]

StructNew in our IR also includes struct.new_default, so I guess this proves the review here is not viewing the entire codebase, but making reasonable inferences about the rest of it.

[gemini-code-assist]

Similarly, we should also treat StructNewDefault as an unescaping allocation here.

Suggested change

	} else {
	// Not a tee: a normal local.set.
	if (!set->value->is<StructNew>()) {
	allSetsAreUnescapingAllocs = false;
	}
	}
	} else {
	// Not a tee: a normal local.set.
	if (!set->value->is<StructNew>() && !set->value->is<StructNewDefault>()) {
	allSetsAreUnescapingAllocs = false;
	}
	}

[kripken]

I verified this has no significant impact on compile times, and improves code in multiple languages.

[tlively]

LGTM % missing coverage (and LGTM anyway if there is a good reason for it).

[tlively]

Looks like the first if here is always true and the second if is always false, so the continue is never executed.

[kripken]

Ah, yes, I added an if nothing to optimize, leave early before, so the tee test stopped getting covered. Fixed.

[kripken]

Also I made a change to optimize only when we see an allocation

Sadly, this change made us stop optimizing the original testcase. But I am a little worried about speculatively optimizing cases without allocations, as (it turns out) appears there, as it might regress speed. Not sure yet how to proceed here.

[tlively]

Is this the part where it might escape? Because we only consider struct.get as consuming here and everything else, including drop is considered an escape? Probably worth a comment if so, because it is obvious to a human reader that this does not escape.

[tlively]

Sadly, this change made us stop optimizing the original testcase

What was the original test case?

[kripken]

Fixes #8641

The Dart testcase from there.

[tlively]

I think it's probably fine to optimize even without an allocation because we are also checking (IIRC) that there are at least 2 local.gets, meaning there are also at least 2 struct.gets. So even if we are executing a struct.get earlier, we're still replacing 2 struct.gets with just 1. We could run the sheets benchmarks to make sure there's no problem. (And it would be nice to have more GC benchmarks to run, too.)

[kripken]

So even if we are executing a struct.get earlier, we're still replacing 2 struct.gets with just 1.

Do you mean in code size or runtime? But I think it can be worse in both cases. This for example adds many struct.gets to the code:

(module
  (type $struct-immutable (struct (field i32)))

  (func $unoptimized-tee (param $x i32)
    (local $yes (ref null $struct-immutable))

    ;; Imagine these are in nested code and not in a silly row like this
    (local.set $yes
      (struct.new_default $struct-immutable)
    )
    (local.set $yes
      (struct.new_default $struct-immutable)
    )
    (local.set $yes
      (struct.new_default $struct-immutable)
    )
    (local.set $yes
      (struct.new_default $struct-immutable)
    )
    (local.set $yes
      (struct.new_default $struct-immutable)
    )
    (local.set $yes
      (struct.new_default $struct-immutable)
    )

    (drop
      (struct.get $struct-immutable 0  ;; this is removed and 6 are added above
        (local.get $yes)
      )
    )
  )
)

And at runtime, the struct.new might be executed - now with a struct.get right after it, if we optimize - and the original struct.get never reached.

[kripken]

I'll see how hard it is to analyze groups of locals together (one can unlock another).

[tlively]

Do you mean in code size or runtime? But I think it can be worse in both cases. This for example adds many struct.gets to the code:

Oh yes, good point. But when there are allocations like in that example, we definitely want to optimize because that will allow Heap2Local to clean up more. A real bad example would have several parameters or function results (or some other source of struct references) set to the same local in different control flow paths and then accessed. I hope this would be rare.

I would test the different strategies on our corpus of interesting modules to see if there is a clear winner in practice.

Alternatively, we could just conservatively not optimize in this case and let a future cross-block CSE/GVN optimization (which would have to be aware of aliasing anyway) take care of any remaining duplicate struct.gets.