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

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

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

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

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

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

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

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

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

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