1
//! Wasm module section-level registration: types, imports, functions,
2
//! exports, data segments, plus the local-pool allocator.
3
//!
4
//! Every helper that mutates the wasm encoder sections lives here so
5
//! the other context submodules (types, pair, ratio, commodity, ...)
6
//! are pure consumers of stable indices.
7

            
8
use super::{CompileContext, LOCAL_POOL_BASE};
9
use crate::ast::WasmType;
10
use crate::error::{Error, Result};
11
use tracing::debug;
12
use wasm_encoder::{
13
    ArrayType, CompositeInnerType, CompositeType, EntityType as WasmEntityType, ExportKind,
14
    FieldType, StorageType, StructType, SubType, ValType,
15
};
16

            
17
28110228
fn bump(counter: &mut u32, kind: &'static str) -> Result<u32> {
18
28110228
    let idx = *counter;
19
28110228
    *counter = counter
20
28110228
        .checked_add(1)
21
28110228
        .ok_or_else(|| Error::Compile(format!("wasm {kind} index space exhausted")))?;
22
28110228
    Ok(idx)
23
28110228
}
24

            
25
impl CompileContext {
26
359629
    pub fn alloc_local(&mut self, ty: WasmType) -> Result<u32> {
27
359629
        let idx = bump(&mut self.next_local, "local")?;
28
359629
        self.local_types.push((ty, idx));
29
359629
        Ok(idx)
30
359629
    }
31

            
32
208716
    pub fn reset_locals(&mut self) {
33
208716
        self.next_local = LOCAL_POOL_BASE;
34
208716
        self.local_types.clear();
35
208716
        self.closure_bodies.clear();
36
208716
        self.serializer =
37
208716
            super::super::layout::OutputSerializer::new(super::super::expr::LOCAL_OUTPUT_BASE);
38
208716
    }
39

            
40
208715
    pub fn build_locals_declaration(&self) -> Vec<(u32, ValType)> {
41
208715
        let preallocated: [(u32, ValType); 6] = [
42
208715
            (1, self.string_ref()),
43
208715
            (1, ValType::I32),
44
208715
            (1, ValType::I32),
45
208715
            (1, self.ratio_ref()),
46
208715
            (1, ValType::I32),
47
208715
            (1, ValType::I32),
48
208715
        ];
49
208715
        let mut locals: Vec<(u32, ValType)> = preallocated.to_vec();
50
336977
        for &(ty, _) in &self.local_types {
51
336977
            locals.push((1, self.wasm_val_type(ty)));
52
336977
        }
53
208715
        locals
54
208715
    }
55

            
56
219162
    pub fn register_type(&mut self) -> Result<u32> {
57
219162
        let idx = bump(&mut self.type_count, "type")?;
58
219162
        self.types.ty().subtype(&SubType {
59
219162
            is_final: true,
60
219162
            supertype_idx: None,
61
219162
            composite_type: CompositeType {
62
219162
                inner: CompositeInnerType::Array(ArrayType(FieldType {
63
219162
                    element_type: StorageType::I8,
64
219162
                    mutable: true,
65
219162
                })),
66
219162
                shared: false,
67
219162
                describes: None,
68
219162
                descriptor: None,
69
219162
            },
70
219162
        });
71
219162
        Ok(idx)
72
219162
    }
73

            
74
    /// Registers a mutable `(array (mut i64))` type. Used for the commodity
75
    /// unit-term (ADR-0028): a flat array of sorted `(hi, lo, exp)` i64 triples.
76
219162
    pub fn register_i64_array_type(&mut self) -> Result<u32> {
77
219162
        let idx = bump(&mut self.type_count, "type")?;
78
219162
        self.types.ty().subtype(&SubType {
79
219162
            is_final: true,
80
219162
            supertype_idx: None,
81
219162
            composite_type: CompositeType {
82
219162
                inner: CompositeInnerType::Array(ArrayType(FieldType {
83
219162
                    element_type: StorageType::Val(ValType::I64),
84
219162
                    mutable: true,
85
219162
                })),
86
219162
                shared: false,
87
219162
                describes: None,
88
219162
                descriptor: None,
89
219162
            },
90
219162
        });
91
219162
        Ok(idx)
92
219162
    }
93

            
94
2410784
    pub fn register_struct_type(&mut self, fields: &[ValType]) -> Result<u32> {
95
2410784
        let idx = bump(&mut self.type_count, "type")?;
96
2410784
        let struct_fields: Vec<FieldType> = fields
97
2410784
            .iter()
98
2410784
            .map(|vt| FieldType {
99
8328158
                element_type: StorageType::Val(*vt),
100
                mutable: false,
101
8328158
            })
102
2410784
            .collect();
103
2410784
        self.types.ty().subtype(&SubType {
104
2410784
            is_final: true,
105
2410784
            supertype_idx: None,
106
2410784
            composite_type: CompositeType {
107
2410784
                inner: CompositeInnerType::Struct(StructType {
108
2410784
                    fields: struct_fields.into_boxed_slice(),
109
2410784
                }),
110
2410784
                shared: false,
111
2410784
                describes: None,
112
2410784
                descriptor: None,
113
2410784
            },
114
2410784
        });
115
2410784
        Ok(idx)
116
2410784
    }
117

            
118
    /// Errors if `name` is already bound in `func_names`. The name space is a
119
    /// closed universe (built-in helpers + reserved imports `__nomi_raise` /
120
    /// `log` / `__nomi_catch_each` + user host fns + unique monomorph/lambda
121
    /// helper names), and a silent overwrite would re-point an already-resolved
122
    /// caller at the wrong index. Notably this rejects a user `HostFnSpec` whose
123
    /// `import_name` collides with a reserved built-in (e.g. `"log"`), which
124
    /// would otherwise misroute PRINT/DISPLAY to the user fn or emit invalid
125
    /// wasm on a signature mismatch — surfaced as a structured `Error::Compile`.
126
15151836
    fn reject_duplicate_func_name(&self, name: &str) -> Result<()> {
127
15151836
        if self.func_names.contains_key(name) {
128
1
            return Err(Error::Compile(format!(
129
1
                "wasm function name '{name}' is already registered (reserved built-in or duplicate host fn)"
130
1
            )));
131
15151835
        }
132
15151835
        Ok(())
133
15151836
    }
134

            
135
7174944
    pub fn register_import(
136
7174944
        &mut self,
137
7174944
        module: &str,
138
7174944
        name: &str,
139
7174944
        params: &[ValType],
140
7174944
        results: &[ValType],
141
7174944
    ) -> Result<u32> {
142
7174944
        self.reject_duplicate_func_name(name)?;
143
7174943
        let type_idx = self.get_or_create_func_type(params, results)?;
144
7174943
        self.imports
145
7174943
            .import(module, name, WasmEntityType::Function(type_idx));
146
7174943
        let func_idx = bump(&mut self.import_func_count, "imported function")?;
147
7174943
        self.func_names.insert(name.to_string(), func_idx);
148
7174943
        Ok(func_idx)
149
7174944
    }
150

            
151
    /// Reserves a stable wasm function index for `name` *before* its
152
    /// body is emitted, and returns the new index. The function
153
    /// section already accounts for the slot at this point, so any
154
    /// body emitted afterward can `call $name` (including itself for
155
    /// recursion) by looking the index up via [`Self::declared_func_index`].
156
    /// The body must land on `pending_helpers` in the order matching the
157
    /// reservation sequence; later phases drain that queue into the code
158
    /// section.
159
7976892
    pub fn register_function(
160
7976892
        &mut self,
161
7976892
        name: &str,
162
7976892
        params: &[ValType],
163
7976892
        results: &[ValType],
164
7976892
    ) -> Result<u32> {
165
7976892
        self.reject_duplicate_func_name(name)?;
166
7976892
        let type_idx = self.get_or_create_func_type(params, results)?;
167
7976892
        self.functions.function(type_idx);
168
7976892
        let local_idx = bump(&mut self.local_func_count, "local function")?;
169
7976892
        let func_idx = self
170
7976892
            .import_func_count
171
7976892
            .checked_add(local_idx)
172
7976892
            .ok_or_else(|| Error::Compile("wasm function index space exhausted".to_string()))?;
173
7976892
        self.func_names.insert(name.to_string(), func_idx);
174
7976892
        Ok(func_idx)
175
7976892
    }
176

            
177
    /// Exports a previously-declared function by name. The name set is the
178
    /// closed, static universe registered by the constructor, so a miss is a
179
    /// compiler bug — surfaced as a structured `Error::Compile` rather than a
180
    /// `HashMap[key]` panic (CLAUDE.md).
181
2931631
    pub fn export_func(&mut self, name: &str) -> Result<()> {
182
2931631
        let idx = self.func_names.get(name).copied().ok_or_else(|| {
183
            Error::Compile(format!(
184
                "internal: cannot export unregistered wasm function '{name}'"
185
            ))
186
        })?;
187
2931631
        self.exports.export(name, ExportKind::Func, idx);
188
2931631
        Ok(())
189
2931631
    }
190

            
191
    /// Looks up a declared function index by name, erroring (never panicking)
192
    /// on a miss. For the rare construction-time helper not covered by a typed
193
    /// `WasmIds` field (e.g. `unit_div`); emit-time code uses `ctx.ids.*`.
194
219161
    pub(super) fn declared_func_index(&self, name: &str) -> Result<u32> {
195
219161
        self.func_names.get(name).copied().ok_or_else(|| {
196
            Error::Compile(format!(
197
                "internal: wasm function '{name}' was not registered during context construction"
198
            ))
199
        })
200
219161
    }
201

            
202
151872
    pub fn serializer(&mut self) -> &mut super::super::layout::OutputSerializer {
203
151872
        &mut self.serializer
204
151872
    }
205

            
206
15374407
    pub(super) fn get_or_create_func_type(
207
15374407
        &mut self,
208
15374407
        params: &[ValType],
209
15374407
        results: &[ValType],
210
15374407
    ) -> Result<u32> {
211
15374407
        if let Some(inner) = self.type_cache.get(params)
212
10128369
            && let Some(&idx) = inner.get(results)
213
        {
214
7000804
            return Ok(idx);
215
8373603
        }
216
8373603
        let idx = bump(&mut self.type_count, "type")?;
217
8373603
        self.types
218
8373603
            .ty()
219
8373603
            .function(params.iter().copied(), results.iter().copied());
220
8373603
        self.type_cache
221
8373603
            .entry(params.to_vec())
222
8373603
            .or_default()
223
8373603
            .insert(results.to_vec(), idx);
224
8373603
        Ok(idx)
225
15374407
    }
226

            
227
1376053
    pub fn add_data(&mut self, bytes: &[u8]) -> Result<u32> {
228
1376053
        let idx = bump(&mut self.data_count, "data segment")?;
229
1376053
        debug!(idx, len = bytes.len(), "adding data segment");
230
1376053
        self.data.passive(bytes.iter().copied());
231
1376053
        Ok(idx)
232
1376053
    }
233
}