flatMap: segment nodes and the shared driver
map is a friendly contract: one element in, one element out, output length known before
the loop starts. flatMap promises nothing. Every element produces a whole collection,
of any size, including zero, so the output length is unknowable up front, and each inner collection
is an allocation somebody already paid for. Most libraries copy every inner into a growing
buffer and throw the inners away. FArray improves on that in two places: the copying, and how the
loop gets compiled.
Don't copy what you can keep
The inner collections f returns are already backed by arrays, FArray arrays. Copying them into one
big buffer discards perfectly good, already-materialized storage. So flatMap looks at the first
inner it produces and makes a bet. If inners are wide, it doesn't copy at all: it keeps each
inner's backing array as a segment and wraps them in a dedicated node, an array of arrays plus a
prefix-sum index, one more citizen of the lazy core.
Traversal walks the segments back to back; indexed access binary-searches the offsets. If inners are
narrow, the bet flips: a two-element segment costs more in pointer-chasing than it saves in
copying, so narrow inners are flattened into one contiguous leaf the classic way. The crossover is
measured, not guessed, and it differs per element kind (references keep the node from width 8;
primitives, whose flat arrays vectorize so well that locality is worth more, only from width 32). Tiny
inners don't even use System.arraycopy: below eight elements the stub call costs more than the copy,
so they move with a plain loop.
The 35× escape-analysis collapse
The first version of this op didn't look like that. It was fully inline: the whole node-or-flatten
loop, both paths, spliced into your method at every flatMap call site, the way .fuse (the
library's whole-pipeline fusion macro) does by design. And in isolation it was excellent: with the loop and your lambda in one compilation
unit, Graal's escape analysis deleted the per-element inner FArray(a, b) wrappers entirely.
The benchmarks were good until we appended one method call.
xs.flatMap(s => FArray(s, s)).take(n)
That chain ran 35× slower than the same flatMap alone: 74 ops/s where a bare array managed
1,400, reproducibly, and only when a structural call like take, drop or reverse trailed the
loop in the same method. Method size was innocent (a double filter, much bigger, was fine).
Deoptimization logs were empty. The tell was allocation: the healthy chain allocated 3.2 MB per op,
the poisoned one 7.2 MB. Escape analysis had silently given up on the spliced loop, and every
per-element wrapper it used to delete became real garbage on a slow path. The trigger was an
ordinary chain any user would write in their first five minutes.
The fix is the same shape map already uses: the loop is shared, not spliced. The inline
surface resolves kinds at compile time, wraps your lambda and the unboxed source read into one
specialized SAM, and calls a driver that is compiled once, in its own method, whose compilation the
caller cannot poison:
Because the inners are FArrays regardless of element type, one driver per output kind covers all
eighty-one kind combinations. The entry is small enough for the JIT to inline into a hot caller
(recovering the escape analysis when it's profitable), and when it doesn't, the loop still compiles
well instead of catastrophically. The poisoned chain went from 74 ops/s to beating the bare array;
plain flatMap got faster too, because the shared loop compiles well regardless of what the caller
does.
And the op on its own, Int beside String:
One caveat: a chain of flatMaps with tiny inners pays the SAM hand-off more than once, and
sits a few percent behind where the old spliced loop did on its best day. That's the same trade map
made, and it's the pipeline shape .fuse exists for: the whole chain becomes
one loop.