1
//! Closure-signature registry.
2
//!
3
//! Tier 1.5 promotes lambdas from compile-time stringification to real
4
//! wasm functions. To pass a lambda around as a runtime value we need
5
//! a wasm GC type that bundles a typed funcref with its captured-env
6
//! ref — that's `$closure_<sig>`. Closures sharing the same
7
//! `(arg-types) -> ret-type` signature share the wasm closure-struct
8
//! type and the funcref signature; their env-struct types vary
9
//! independently per scope so two closures with disjoint captures
10
//! don't collide.
11
//!
12
//! This slice ships the per-signature interner that the
13
//! `WasmType::Closure(sig_id)` arm of `wasm_val_type` resolves through;
14
//! the env-struct allocator and monomorph cache (used to lower
15
//! `(funcall closure-arg ...)` and `(defun ...)` recursion through a
16
//! real fn) land alongside their first call-site so the no-dead-code
17
//! rule stays clean. See `doc/adr/0027-lambda-real-wasm-fn.org` for the
18
//! full design.
19

            
20
#[cfg(test)]
21
mod tests;
22

            
23
use std::collections::HashMap;
24

            
25
use wasm_encoder::{
26
    CompositeInnerType, CompositeType, FieldType, Function, HeapType, RefType, StorageType,
27
    StructType, SubType, ValType,
28
};
29

            
30
use super::CompileContext;
31
use crate::ast::{ClosureSigId, Expr, LambdaParams, WasmType};
32
use crate::error::{Error, Result};
33

            
34
/// Per-signature closure type metadata. Held in `CompileContext` so
35
/// the compiler can look up `$closure_<sig>` and `$fn_<sig>` indices
36
/// from a [`ClosureSigId`] without rehashing the signature on each
37
/// reference. `fn_type_idx` is the typed funcref's wasm type — needed
38
/// at every `call_ref` site so the engine can validate the funcref's
39
/// shape statically.
40
#[derive(Debug, Clone)]
41
pub(crate) struct ClosureSig {
42
    pub closure_type_idx: u32,
43
    pub fn_type_idx: u32,
44
    pub params: Vec<WasmType>,
45
    pub result: WasmType,
46
}
47

            
48
/// Layout of a per-scope `$env_<id>` struct holding values captured by
49
/// a lambda body. One entry per `CaptureSet` name, in insertion order
50
/// so the outer site (which fills the struct) and the helper-fn body
51
/// (which reads it) walk the same field indices. Field types are
52
/// captured at the point [`CompileContext::intern_env_struct`] is
53
/// called from the lookup of each name in the caller's symbol table.
54
#[derive(Debug, Clone)]
55
pub(crate) struct EnvStructLayout {
56
    pub type_idx: u32,
57
    pub fields: Vec<EnvField>,
58
}
59

            
60
#[derive(Debug, Clone)]
61
pub(crate) struct EnvField {
62
    pub name: String,
63
    pub ty: WasmType,
64
}
65

            
66
#[derive(Debug, Default, Clone)]
67
pub(crate) struct ClosureRegistry {
68
    sigs: HashMap<(Vec<WasmType>, WasmType), ClosureSigId>,
69
    by_id: Vec<ClosureSig>,
70
    next_sig_id: u32,
71
    helper_count: u32,
72
    env_struct_count: u32,
73
    closure_literal_data_idx: Option<u32>,
74
}
75

            
76
/// Snapshot of `CompileContext`'s local-pool fields so a helper-fn
77
/// emit can carve out a fresh allocator without clobbering the caller's
78
/// in-flight state. Rebound by [`CompileContext::restore_local_pool`].
79
pub(crate) struct LocalPoolSnapshot {
80
    next_local: u32,
81
    local_types: Vec<(WasmType, u32)>,
82
    closure_bodies: HashMap<u32, (LambdaParams, Expr)>,
83
    serializer: super::super::layout::OutputSerializer,
84
}
85

            
86
impl CompileContext {
87
    /// Returns a `(ref null $closure_<sig>)` valtype. Helper for the
88
    /// `WasmType::Closure(sig)` arm of `wasm_val_type`.
89
5987
    pub(crate) fn closure_ref(&self, sig: ClosureSigId) -> ValType {
90
5987
        let entry = &self.closures.by_id[sig.0 as usize];
91
5987
        ValType::Ref(RefType {
92
5987
            nullable: true,
93
5987
            heap_type: HeapType::Concrete(entry.closure_type_idx),
94
5987
        })
95
5987
    }
96

            
97
10881
    pub(crate) fn closure_sig(&self, sig: ClosureSigId) -> &ClosureSig {
98
10881
        &self.closures.by_id[sig.0 as usize]
99
10881
    }
100

            
101
    /// Records a value-position lambda's source, keyed by the wasm local its
102
    /// closure value was bound to (see [`Self::closure_bodies`]).
103
1975
    pub(in crate::compiler) fn record_closure_body(
104
1975
        &mut self,
105
1975
        local_idx: u32,
106
1975
        params: LambdaParams,
107
1975
        body: Expr,
108
1975
    ) {
109
1975
        self.closure_bodies.insert(local_idx, (params, body));
110
1975
    }
111

            
112
    /// The source `(params, body)` of the closure value bound to `local_idx`,
113
    /// if one was recorded — lets a higher-order native inline the body per
114
    /// element with the actual element type instead of `call_ref`ing the
115
    /// fixed-signature closure.
116
822
    pub(in crate::compiler) fn closure_body(
117
822
        &self,
118
822
        local_idx: u32,
119
822
    ) -> Option<&(LambdaParams, Expr)> {
120
822
        self.closure_bodies.get(&local_idx)
121
822
    }
122

            
123
    /// Drops any recorded body for `local_idx`. Called when a `setf` overwrites
124
    /// a closure local: the recorded source no longer describes the value now in
125
    /// that local, so an inline would apply the WRONG closure. Forgetting it
126
    /// falls the higher-order native back to the (always-correct) `call_ref`.
127
136
    pub(in crate::compiler) fn forget_closure_body(&mut self, local_idx: u32) {
128
136
        self.closure_bodies.remove(&local_idx);
129
136
    }
130

            
131
    /// Look up or register the wasm types behind a closure signature.
132
    /// Returns the same id for repeat calls with the same signature so
133
    /// distinct lambdas with the same shape share `$fn_<sig>` and
134
    /// `$closure_<sig>`.
135
3479
    pub(crate) fn intern_closure_signature(
136
3479
        &mut self,
137
3479
        params: &[WasmType],
138
3479
        result: WasmType,
139
3479
    ) -> Result<ClosureSigId> {
140
3479
        let key = (params.to_vec(), result);
141
3479
        if let Some(id) = self.closures.sigs.get(&key) {
142
69
            return Ok(*id);
143
3410
        }
144
3410
        let env_anyref = self.anyref();
145
3410
        let mut fn_params = Vec::with_capacity(params.len() + 1);
146
3410
        fn_params.push(env_anyref);
147
4362
        fn_params.extend(params.iter().map(|ty| self.wasm_val_type(*ty)));
148
3410
        let result_vt = self.wasm_val_type(result);
149
3410
        let fn_type_idx = self.get_or_create_func_type(&fn_params, &[result_vt])?;
150

            
151
3410
        let closure_type_idx = self.type_count;
152
3410
        let closure_fields = vec![
153
3410
            FieldType {
154
3410
                element_type: StorageType::Val(ValType::Ref(RefType {
155
3410
                    nullable: false,
156
3410
                    heap_type: HeapType::Concrete(fn_type_idx),
157
3410
                })),
158
3410
                mutable: false,
159
3410
            },
160
3410
            FieldType {
161
3410
                element_type: StorageType::Val(env_anyref),
162
3410
                mutable: false,
163
3410
            },
164
        ];
165
3410
        self.types.ty().subtype(&SubType {
166
3410
            is_final: true,
167
3410
            supertype_idx: None,
168
3410
            composite_type: CompositeType {
169
3410
                inner: CompositeInnerType::Struct(StructType {
170
3410
                    fields: closure_fields.into_boxed_slice(),
171
3410
                }),
172
3410
                shared: false,
173
3410
                describes: None,
174
3410
                descriptor: None,
175
3410
            },
176
3410
        });
177
3410
        self.type_count = self
178
3410
            .type_count
179
3410
            .checked_add(1)
180
3410
            .ok_or_else(|| Error::Compile("wasm type index space exhausted".to_string()))?;
181
3410
        let id = ClosureSigId(self.closures.next_sig_id);
182
        self.closures.next_sig_id =
183
3410
            self.closures.next_sig_id.checked_add(1).ok_or_else(|| {
184
                Error::Compile("closure signature id space exhausted".to_string())
185
            })?;
186

            
187
3410
        self.closures.by_id.push(ClosureSig {
188
3410
            closure_type_idx,
189
3410
            fn_type_idx,
190
3410
            params: params.to_vec(),
191
3410
            result,
192
3410
        });
193
3410
        self.closures.sigs.insert(key, id);
194
3410
        Ok(id)
195
3479
    }
196

            
197
    /// Registers a fresh `$env_<id>` struct type for one captured-name
198
    /// scope. Each call allocates a new wasm type (no dedup) — different
199
    /// lambdas captures different names even when shapes coincide, and
200
    /// the per-call cost is one type-section entry. Caller's
201
    /// responsibility to populate the layout in declared order.
202
340
    pub(crate) fn intern_env_struct(&mut self, fields: &[EnvField]) -> Result<EnvStructLayout> {
203
340
        let env_type_idx = self.type_count;
204
340
        let struct_fields: Vec<FieldType> = fields
205
340
            .iter()
206
340
            .map(|f| FieldType {
207
340
                element_type: StorageType::Val(self.wasm_val_type(f.ty)),
208
                mutable: false,
209
340
            })
210
340
            .collect();
211
340
        self.types.ty().subtype(&SubType {
212
340
            is_final: true,
213
340
            supertype_idx: None,
214
340
            composite_type: CompositeType {
215
340
                inner: CompositeInnerType::Struct(StructType {
216
340
                    fields: struct_fields.into_boxed_slice(),
217
340
                }),
218
340
                shared: false,
219
340
                describes: None,
220
340
                descriptor: None,
221
340
            },
222
340
        });
223
340
        self.type_count = self
224
340
            .type_count
225
340
            .checked_add(1)
226
340
            .ok_or_else(|| Error::Compile("wasm type index space exhausted".to_string()))?;
227
340
        self.closures.env_struct_count = self
228
340
            .closures
229
340
            .env_struct_count
230
340
            .checked_add(1)
231
340
            .ok_or_else(|| Error::Compile("env-struct id space exhausted".to_string()))?;
232
340
        Ok(EnvStructLayout {
233
340
            type_idx: env_type_idx,
234
340
            fields: fields.to_vec(),
235
340
        })
236
340
    }
237

            
238
    /// Allocates a unique helper-fn name for a lambda body. Used by the
239
    /// emit path so distinct lambda bodies get stable wasm names even
240
    /// when sharing a `ClosureSigId`.
241
3468
    pub(crate) fn next_lambda_helper_name(&mut self) -> Result<String> {
242
3468
        let id = self.closures.helper_count;
243
3468
        self.closures.helper_count = self
244
3468
            .closures
245
3468
            .helper_count
246
3468
            .checked_add(1)
247
3468
            .ok_or_else(|| Error::Compile("lambda helper id space exhausted".to_string()))?;
248
3468
        Ok(format!("__lambda_{id}"))
249
3468
    }
250

            
251
    /// Saves the in-flight local-pool state (local allocator + types +
252
    /// serializer cursor) so a helper-fn body can be emitted into a
253
    /// fresh pool. The new pool numbers locals from `next_local_base`,
254
    /// which the caller picks to match the helper fn's wasm-level
255
    /// layout — script-mode entry points use `LOCAL_POOL_BASE` (5
256
    /// preallocated slots), lambda helpers use `param_count` (no
257
    /// preallocated slots, locals begin immediately after params).
258
    /// Pair with [`Self::restore_local_pool`] once the helper's
259
    /// `Function` has been finished and queued.
260
5033
    pub(crate) fn take_local_pool(&mut self, next_local_base: u32) -> LocalPoolSnapshot {
261
5033
        let saved = LocalPoolSnapshot {
262
5033
            next_local: self.next_local,
263
5033
            local_types: std::mem::take(&mut self.local_types),
264
5033
            closure_bodies: std::mem::take(&mut self.closure_bodies),
265
5033
            serializer: std::mem::replace(
266
5033
                &mut self.serializer,
267
5033
                super::super::layout::OutputSerializer::new(super::super::expr::LOCAL_OUTPUT_BASE),
268
5033
            ),
269
5033
        };
270
5033
        self.next_local = next_local_base;
271
5033
        saved
272
5033
    }
273

            
274
4897
    pub(crate) fn restore_local_pool(&mut self, snapshot: LocalPoolSnapshot) {
275
4897
        self.next_local = snapshot.next_local;
276
4897
        self.local_types = snapshot.local_types;
277
4897
        self.closure_bodies = snapshot.closure_bodies;
278
4897
        self.serializer = snapshot.serializer;
279
4897
    }
280

            
281
    /// Builds the locals declaration for a lambda helper fn whose
282
    /// `param_count` slots are reserved (env + user params, all func
283
    /// args). Mirrors [`Self::build_locals_declaration`] but starts
284
    /// numbering after the params and skips the script-mode preallocated
285
    /// pool — helper bodies use the local-pool allocator only.
286
4556
    pub(crate) fn build_helper_locals(&self, param_count: u32) -> Vec<(u32, ValType)> {
287
4556
        let mut locals: Vec<(u32, ValType)> = Vec::with_capacity(self.local_types.len());
288
4556
        let _ = param_count;
289
7616
        for &(ty, _) in &self.local_types {
290
7616
            locals.push((1, self.wasm_val_type(ty)));
291
7616
        }
292
4556
        locals
293
4556
    }
294

            
295
4354
    pub(crate) fn queue_helper(&mut self, helper: Function) {
296
4354
        self.pending_helpers.push(helper);
297
4354
    }
298

            
299
    /// Returns the passive-data-segment idx for the literal `<closure>`
300
    /// rendering used by [`OutputSerializer::write_debug_closure_from_stack`].
301
    /// The bytes are interned on first use so a module that constructs
302
    /// many closures shares one data segment.
303
408
    pub(crate) fn closure_literal_data_idx(&mut self) -> Result<u32> {
304
408
        if let Some(idx) = self.closures.closure_literal_data_idx {
305
            return Ok(idx);
306
408
        }
307
408
        let idx = self.add_data(CLOSURE_LITERAL_BYTES)?;
308
408
        self.closures.closure_literal_data_idx = Some(idx);
309
408
        Ok(idx)
310
408
    }
311
}
312

            
313
pub(crate) const CLOSURE_LITERAL_BYTES: &[u8] = b"<closure>";
314
pub(crate) const CLOSURE_LITERAL_LEN: u32 = CLOSURE_LITERAL_BYTES.len() as u32;