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

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

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

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

            
56
252751
    pub fn register_type(&mut self) -> Result<u32> {
57
252751
        let idx = bump(&mut self.type_count, "type")?;
58
252751
        self.types.ty().subtype(&SubType {
59
252751
            is_final: true,
60
252751
            supertype_idx: None,
61
252751
            composite_type: CompositeType {
62
252751
                inner: CompositeInnerType::Array(ArrayType(FieldType {
63
252751
                    element_type: StorageType::I8,
64
252751
                    mutable: true,
65
252751
                })),
66
252751
                shared: false,
67
252751
                describes: None,
68
252751
                descriptor: None,
69
252751
            },
70
252751
        });
71
252751
        Ok(idx)
72
252751
    }
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
252751
    pub fn register_i64_array_type(&mut self) -> Result<u32> {
77
252751
        let idx = bump(&mut self.type_count, "type")?;
78
252751
        self.types.ty().subtype(&SubType {
79
252751
            is_final: true,
80
252751
            supertype_idx: None,
81
252751
            composite_type: CompositeType {
82
252751
                inner: CompositeInnerType::Array(ArrayType(FieldType {
83
252751
                    element_type: StorageType::Val(ValType::I64),
84
252751
                    mutable: true,
85
252751
                })),
86
252751
                shared: false,
87
252751
                describes: None,
88
252751
                descriptor: None,
89
252751
            },
90
252751
        });
91
252751
        Ok(idx)
92
252751
    }
93

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

            
135
9427511
    pub fn register_import(
136
9427511
        &mut self,
137
9427511
        module: &str,
138
9427511
        name: &str,
139
9427511
        params: &[ValType],
140
9427511
        results: &[ValType],
141
9427511
    ) -> Result<u32> {
142
9427511
        self.reject_duplicate_func_name(name)?;
143
9427510
        let type_idx = self.get_or_create_func_type(params, results)?;
144
9427510
        self.imports
145
9427510
            .import(module, name, WasmEntityType::Function(type_idx));
146
9427510
        let func_idx = bump(&mut self.import_func_count, "imported function")?;
147
9427510
        self.func_names.insert(name.to_string(), func_idx);
148
9427510
        Ok(func_idx)
149
9427511
    }
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
9094747
    pub fn register_function(
160
9094747
        &mut self,
161
9094747
        name: &str,
162
9094747
        params: &[ValType],
163
9094747
        results: &[ValType],
164
9094747
    ) -> Result<u32> {
165
9094747
        self.reject_duplicate_func_name(name)?;
166
9094747
        let type_idx = self.get_or_create_func_type(params, results)?;
167
9094747
        self.functions.function(type_idx);
168
9094747
        let local_idx = bump(&mut self.local_func_count, "local function")?;
169
9094747
        let func_idx = self
170
9094747
            .import_func_count
171
9094747
            .checked_add(local_idx)
172
9094747
            .ok_or_else(|| Error::Compile("wasm function index space exhausted".to_string()))?;
173
9094747
        self.func_names.insert(name.to_string(), func_idx);
174
9094747
        Ok(func_idx)
175
9094747
    }
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
3276738
    pub fn export_func(&mut self, name: &str) -> Result<()> {
182
3276738
        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
3276738
        self.exports.export(name, ExportKind::Func, idx);
188
3276738
        Ok(())
189
3276738
    }
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
252750
    pub(super) fn declared_func_index(&self, name: &str) -> Result<u32> {
195
252750
        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
252750
    }
201

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

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

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