cabi: memoize the flattened param/result count, not per call (#261)

runtime/src/cabi/flatten.ts

@@ -73,6 +73,46 @@ export function flattenTypes(ts: ValType[], opts: LiftOptions): CoreType[] {
   return ts.flatMap((t) => flattenType(t, opts));
 }
 
+/**
+ * `flattenTypes(ts, opts).length`, memoized on (ts array identity, ptrType).
+ *
+ * Every arm of `flattenType` reaches `opts` only through
+ * `requireMemory(opts).ptrType()` (string/list-without-length read the
+ * pointer width; every other arm is opts-free or recurses structurally), so
+ * the flattened element count is a pure function of `(ts, ptrType)` — one
+ * map per pointer width is the whole of the cache key. `ts` is always a
+ * plan-owned `ft.params`/`ft.results` array, the same stability argument
+ * `spillTupleType` (values.ts) relies on, so its identity is a valid key.
+ *
+ * This caches the COUNT, not the flattened array: `values.ts`'s only use of
+ * `flattenTypes` on the per-call path is `.length`, and a cached number has
+ * no aliasing/mutation hazard to guard (a cached array would need freezing
+ * plus a lossy `readonly`-to-mutable cast at every read site — the shape
+ * this replaced). Callers that need the actual flat types (instantiate-time
+ * `flattenFunctype`) still call `flattenTypes` directly, uncached; that path
+ * runs once per function, not once per call, so it doesn't need this.
+ *
+ * The null-memory path is deliberately NOT cached: `requireMemory` throws
+ * when `opts.memory` is null (for any `ts` containing a string/unbounded
+ * list), and that throw must still surface on every call, not just the
+ * first.
+ */
+const flatCountCacheByPtrType = {
+  i32: new WeakMap<ValType[], number>(),
+  i64: new WeakMap<ValType[], number>(),
+};
+
+export function flatCount(ts: ValType[], opts: LiftOptions): number {
+  const mem = opts.memory;
+  if (mem === null) return flattenTypes(ts, opts).length;
+  const cache = flatCountCacheByPtrType[mem.ptrType()];
+  const hit = cache.get(ts);
+  if (hit !== undefined) return hit;
+  const count = flattenTypes(ts, opts).length;
+  cache.set(ts, count);
+  return count;
+}
+
 export function flattenType(t: ValType, opts: LiftOptions): CoreType[] {
   const d = despecialize(t);
   switch (d.kind) {

runtime/src/cabi/values.ts

@@ -8,7 +8,7 @@ import { load } from "./load.ts";
 import { store } from "./store.ts";
 import { type CoreValueIter, liftFlat, type ValueIter } from "./lift.ts";
 import { lowerFlat } from "./lower.ts";
-import { flattenTypes } from "./flatten.ts";
+import { flatCount } from "./flatten.ts";
 import { type LiftLowerContext, requireMemory } from "./context.ts";
 import { asIndex } from "./memory.ts";
 import type { ComponentValue, CoreValue, TupleType, ValType } from "./types.ts";
@@ -40,8 +40,7 @@ export function liftFlatValues(
   vi: CoreValueIter,
   ts: ValType[],
 ): ComponentValue[] {
-  const flatTypes = flattenTypes(ts, cx.opts);
-  if (flatTypes.length > maxFlat) {
+  if (flatCount(ts, cx.opts) > maxFlat) {
     const mem = requireMemory(cx.opts);
     const ptrRaw = vi.next(mem.ptrType());
     const tupleType = spillTupleType(ts);
@@ -67,8 +66,7 @@ export function lowerFlatValues(
   ts: ValType[],
   outParam: ValueIter | null = null,
 ): CoreValue[] {
-  const flatTypes = flattenTypes(ts, cx.opts);
-  if (flatTypes.length > maxFlat) {
+  if (flatCount(ts, cx.opts) > maxFlat) {
     const mem = requireMemory(cx.opts);
     const tupleType = spillTupleType(ts);
     const tupleValue: Record<string, ComponentValue> = {};