dada_codegen/cx/generate_expr/
wasm_place_repr.rs

1use std::sync::Arc;
2
3use dada_ir_sym::{
4    ir::classes::SymField,
5    ir::exprs::{SymPlaceExpr, SymPlaceExprKind},
6    ir::types::{SymTy, SymTyKind, SymTyName},
7    ir::variables::SymVariable,
8};
9use wasm_encoder::{Instruction, ValType};
10
11use crate::cx::wasm_repr::WasmRepr;
12
13use super::ExprCodegen;
14
15/// The WASM representation for a Dada place. Dada places can be
16/// spread across the WASM local variables and WASM memory.
17/// This type tells you which values are stored where.
18#[derive(Debug)]
19pub enum WasmPlaceRepr {
20    /// A primitive value stored in a WASM local variable.
21    Local(WasmLocal, ValType),
22    Heap(WasmPointer, ValType),
23    Struct(Vec<Arc<WasmPlaceRepr>>),
24    Class(WasmPointer, Vec<Arc<WasmPlaceRepr>>),
25    Nowhere,
26}
27
28impl<'db> ExprCodegen<'_, 'db> {
29    /// Returns a [`WasmPointer`] to the current start of a callee's stack frame.
30    /// This value is only valid until [`Self::insert_variable`] is next called.
31    pub(super) fn next_stack_frame(&self) -> WasmPointer {
32        WasmPointer {
33            base_variable: self.wasm_stack_pointer,
34            offset: self.wasm_stack_frame_size,
35        }
36    }
37
38    /// Introduce the variable `lv` into scope and create a place for it.
39    /// This can allocate more stack space in WASM memory.
40    /// You can find this place by invoking [`Self::place_for_local`] later on.
41    pub(super) fn insert_variable(&mut self, lv: SymVariable<'db>, ty: SymTy<'db>) {
42        let ty_repr = self.wasm_repr_of_type(ty);
43        let emplaced_repr = self.emplace_local(&ty_repr);
44        self.variables.insert(lv, emplaced_repr);
45    }
46
47    /// The representation of the place represented by `local_variable`.
48    pub(super) fn place_for_local(&self, local_variable: SymVariable<'db>) -> Arc<WasmPlaceRepr> {
49        self.variables[&local_variable].clone()
50    }
51
52    /// The representation of the given Dada place.
53    pub(super) fn place(&self, place: SymPlaceExpr<'db>) -> Arc<WasmPlaceRepr> {
54        let db = self.cx.db;
55        match *place.kind(db) {
56            SymPlaceExprKind::Var(v) => self.place_for_local(v),
57            SymPlaceExprKind::Error(_) => Arc::new(WasmPlaceRepr::Nowhere),
58            SymPlaceExprKind::Field(owner, field) => {
59                let owner_place = self.place(owner);
60                self.field_place(owner_place, owner.ty(db), field)
61            }
62        }
63    }
64
65    /// Push the value found in `place` onto the WASM stack.
66    pub(super) fn push_from(&mut self, place: &WasmPlaceRepr) {
67        match *place {
68            WasmPlaceRepr::Local(wasm_local, val_type) => {
69                self.push_from_local(val_type, wasm_local)
70            }
71            WasmPlaceRepr::Heap(wasm_pointer, val_type) => {
72                self.push_from_memory(val_type, wasm_pointer)
73            }
74            WasmPlaceRepr::Struct(ref fields) => {
75                fields.iter().for_each(|r| self.push_from(r));
76            }
77            WasmPlaceRepr::Class(flags, ref fields) => {
78                self.push_from_memory(ValType::I32, flags);
79                fields.iter().for_each(|r| self.push_from(r));
80            }
81            WasmPlaceRepr::Nowhere => (),
82        }
83    }
84
85    /// Push a shared copy the value found in `place` onto the WASM stack.
86    pub(super) fn push_shared_from(&mut self, place: &WasmPlaceRepr) {
87        match *place {
88            WasmPlaceRepr::Struct(ref fields) => {
89                fields.iter().for_each(|r| self.push_shared_from(r));
90            }
91            WasmPlaceRepr::Class(flags, ref fields) => {
92                // Push a 0 for the flags if shared
93                self.instructions.push(Instruction::I32Const(0));
94
95                self.push_from_memory(ValType::I32, flags);
96                fields.iter().for_each(|r| self.push_shared_from(r));
97            }
98            WasmPlaceRepr::Local(..) | WasmPlaceRepr::Heap(..) | WasmPlaceRepr::Nowhere => {
99                self.push_from(place);
100            }
101        }
102    }
103
104    /// Push a shared copy the value found in `place` onto the WASM stack.
105    pub(super) fn push_leased_from(&mut self, place: &WasmPlaceRepr) {
106        match *place {
107            WasmPlaceRepr::Class(flags, _) => {
108                self.push_pointer(flags);
109            }
110            _ => panic!("can only lease classes"),
111        }
112    }
113
114    /// Given that a value of type `value_ty` is on the wasm stack, pop it and store it into `to_place`.
115    pub(super) fn pop_and_store(&mut self, to_place: &WasmPlaceRepr) {
116        match *to_place {
117            WasmPlaceRepr::Local(wasm_local, val_type) => self.pop_to_local(val_type, wasm_local),
118            WasmPlaceRepr::Heap(wasm_pointer, val_type) => {
119                self.pop_to_memory(val_type, wasm_pointer)
120            }
121            WasmPlaceRepr::Struct(ref fields) => {
122                fields.iter().rev().for_each(|r| self.pop_and_store(r));
123            }
124            WasmPlaceRepr::Class(flags, ref fields) => {
125                fields.iter().rev().for_each(|r| self.pop_and_store(r));
126                self.pop_to_memory(ValType::I32, flags);
127            }
128            WasmPlaceRepr::Nowhere => (),
129        }
130    }
131
132    /// Representation for the place storing a given field found in
133    /// an owner of type `owner_ty` that is stored in `owner_place`.
134    fn field_place(
135        &self,
136        owner_place_repr: Arc<WasmPlaceRepr>,
137        owner_ty: SymTy<'db>,
138        field: SymField<'db>,
139    ) -> Arc<WasmPlaceRepr> {
140        let db = self.cx.db;
141        match owner_ty.kind(db) {
142            SymTyKind::Var(sym_variable) => self.field_place(
143                owner_place_repr,
144                self.generics[sym_variable].assert_type(db),
145                field,
146            ),
147            SymTyKind::Infer(_) => panic!("unresolved inference variable"),
148            SymTyKind::Never | SymTyKind::Error(_) => match &*owner_place_repr {
149                WasmPlaceRepr::Nowhere => owner_place_repr,
150                _ => panic!("unexpeced place for {owner_ty:?}: {owner_place_repr:?}"),
151            },
152            SymTyKind::Named(ty_name, _) => match *ty_name {
153                SymTyName::Future => match &*owner_place_repr {
154                    WasmPlaceRepr::Class(_, vec) => vec[0].clone(),
155                    WasmPlaceRepr::Nowhere => owner_place_repr,
156                    _ => panic!("unexpeced place for {owner_ty:?}: {owner_place_repr:?}"),
157                },
158                SymTyName::Primitive(_) => panic!("primitive types do not have fields"),
159                SymTyName::Tuple { arity: _ } => todo!(),
160                SymTyName::Aggregate(aggr) => {
161                    // Where is the owner's data stored?
162                    match &*owner_place_repr {
163                        WasmPlaceRepr::Struct(fields) | WasmPlaceRepr::Class(_, fields) => {
164                            let field_index = aggr
165                                .fields(db)
166                                .take_while(|f: &SymField<'_>| *f != field)
167                                .count();
168                            fields[field_index].clone()
169                        }
170                        WasmPlaceRepr::Nowhere => owner_place_repr,
171                        _ => panic!("unexpeced place for {owner_ty:?}: {owner_place_repr:?}"),
172                    }
173                }
174            },
175            #[expect(unused_variables)]
176            SymTyKind::Perm(sym_perm, sym_ty) => todo!(),
177        }
178    }
179
180    /// Returns the representation of a "local" storing a value of type `repr`.
181    /// A "local" place is one that uses WASM local variables as much as possible.
182    fn emplace_local(&mut self, repr: &WasmRepr) -> Arc<WasmPlaceRepr> {
183        match repr {
184            WasmRepr::Val(val_type) => {
185                let local = self.fresh_local_index(*val_type);
186                Arc::new(WasmPlaceRepr::Local(local, *val_type))
187            }
188            WasmRepr::Struct(vec) => Arc::new(WasmPlaceRepr::Struct(
189                vec.iter().map(|r| self.emplace_local(r)).collect(),
190            )),
191            WasmRepr::Class(_) => self.emplace_memory(repr),
192            WasmRepr::Nothing => Arc::new(WasmPlaceRepr::Nowhere),
193        }
194    }
195
196    /// The representation for a Dada place found in WASM memory
197    /// that stores values with representation `repr`.
198    fn emplace_memory(&mut self, repr: &WasmRepr) -> Arc<WasmPlaceRepr> {
199        match repr {
200            WasmRepr::Val(val_type) => {
201                let pointer = self.fresh_memory_slot(*val_type);
202                Arc::new(WasmPlaceRepr::Heap(pointer, *val_type))
203            }
204            WasmRepr::Struct(vec) => Arc::new(WasmPlaceRepr::Struct(
205                vec.iter().map(|r| self.emplace_memory(r)).collect(),
206            )),
207            WasmRepr::Class(vec) => {
208                let flag_word = self.fresh_memory_slot(ValType::I32);
209                Arc::new(WasmPlaceRepr::Class(
210                    flag_word,
211                    vec.iter().map(|r| self.emplace_memory(r)).collect(),
212                ))
213            }
214            WasmRepr::Nothing => Arc::new(WasmPlaceRepr::Nowhere),
215        }
216    }
217
218    /// Create a fresh local index storing a value of type `v`.
219    fn fresh_local_index(&mut self, v: ValType) -> WasmLocal {
220        let index = u32::try_from(self.wasm_locals.len()).expect("too many locals");
221        self.wasm_locals.push(v);
222        WasmLocal { index }
223    }
224
225    /// Create a fresh slot in memory storing a value of type `v`.
226    fn fresh_memory_slot(&mut self, v: ValType) -> WasmPointer {
227        let offset = self.wasm_stack_frame_size;
228        self.wasm_stack_frame_size += val_type_size_in_bytes(v);
229        WasmPointer {
230            base_variable: self.wasm_stack_pointer,
231            offset,
232        }
233    }
234
235    /// Push a value of type `val_type` found in `local`.
236    fn push_from_local(&mut self, val_type: wasm_encoder::ValType, local: WasmLocal) {
237        assert_eq!(self.wasm_locals[local.index as usize], val_type);
238        self.instructions.push(Instruction::LocalGet(local.index));
239    }
240
241    /// Pop a value of type `val_type` and store it in `local`.
242    fn pop_to_local(&mut self, v: ValType, local: WasmLocal) {
243        assert_eq!(self.wasm_locals[local.index as usize], v);
244        self.instructions.push(Instruction::LocalSet(local.index));
245    }
246
247    /// Push a value of type `val_type` found in the given memory slot.
248    fn push_from_memory(
249        &mut self,
250        v: ValType,
251        WasmPointer {
252            base_variable,
253            offset,
254        }: WasmPointer,
255    ) {
256        let offset = offset as u64;
257        self.instructions.push(match v {
258            ValType::I32 => Instruction::I32Load(wasm_encoder::MemArg {
259                offset,
260                align: 4,
261                memory_index: base_variable.index,
262            }),
263            ValType::I64 => Instruction::I64Load(wasm_encoder::MemArg {
264                offset,
265                align: 4,
266                memory_index: base_variable.index,
267            }),
268            ValType::F32 => Instruction::F32Load(wasm_encoder::MemArg {
269                offset,
270                align: 4,
271                memory_index: base_variable.index,
272            }),
273            ValType::F64 => Instruction::F64Load(wasm_encoder::MemArg {
274                offset,
275                align: 4,
276                memory_index: base_variable.index,
277            }),
278            ValType::V128 | ValType::Ref(_) => panic!("unexpected val type {v:?}"),
279        });
280    }
281
282    /// Push a pointer itself onto the WASM stack (not the data it refers to).
283    pub(super) fn push_pointer(
284        &mut self,
285        WasmPointer {
286            base_variable,
287            offset,
288        }: WasmPointer,
289    ) {
290        self.push_from_local(ValType::I32, base_variable);
291        self.instructions.push(Instruction::I32Const(offset as i32));
292        self.instructions.push(Instruction::I32Add);
293    }
294
295    /// Pop a value of type `val_type` and store it to the given memory slot.
296    fn pop_to_memory(
297        &mut self,
298        v: ValType,
299        WasmPointer {
300            base_variable,
301            offset,
302        }: WasmPointer,
303    ) {
304        let offset = offset as u64;
305        self.instructions.push(match v {
306            ValType::I32 => Instruction::I32Store(wasm_encoder::MemArg {
307                offset,
308                align: 4,
309                memory_index: base_variable.index,
310            }),
311            ValType::I64 => Instruction::I64Store(wasm_encoder::MemArg {
312                offset,
313                align: 4,
314                memory_index: base_variable.index,
315            }),
316            ValType::F32 => Instruction::F32Store(wasm_encoder::MemArg {
317                offset,
318                align: 4,
319                memory_index: base_variable.index,
320            }),
321            ValType::F64 => Instruction::F64Store(wasm_encoder::MemArg {
322                offset,
323                align: 4,
324                memory_index: base_variable.index,
325            }),
326            ValType::V128 | ValType::Ref(_) => panic!("unexpected val type {v:?}"),
327        });
328    }
329}
330
331impl WasmRepr {
332    /// Primitive WASM values needed for a value with this representation stored on the WASM stack or in memory.
333    pub fn flatten(&self) -> Vec<ValType> {
334        match self {
335            WasmRepr::Val(val_type) => vec![*val_type],
336
337            // Structs are just each field one after the other.
338            WasmRepr::Struct(fields) => fields.iter().flat_map(|r| r.local_val_tys()).collect(),
339
340            // Classes begin with an `I32` flag word.
341            WasmRepr::Class(fields) => std::iter::once(ValType::I32)
342                .chain(fields.iter().flat_map(|r| r.local_val_tys()))
343                .collect(),
344
345            WasmRepr::Nothing => vec![],
346        }
347    }
348
349    /// Returns the types of the WASM local variables that would be used to store a value with this representation.
350    /// Any data found inside of a class is stored in memory and hence not represented in the return type.
351    pub fn local_val_tys(&self) -> Vec<ValType> {
352        match self {
353            WasmRepr::Val(val_type) => vec![*val_type],
354            WasmRepr::Struct(fields) => fields.iter().flat_map(|r| r.local_val_tys()).collect(),
355            WasmRepr::Class(_) => vec![],
356            WasmRepr::Nothing => vec![],
357        }
358    }
359}
360
361fn val_type_size_in_bytes(v: ValType) -> u32 {
362    match v {
363        ValType::I32 => 4,
364        ValType::I64 => 8,
365        ValType::F32 => 4,
366        ValType::F64 => 8,
367        ValType::V128 => 16,
368        ValType::Ref(_) => panic!("ref values do not have a size in bytes"),
369    }
370}
371
372#[derive(Copy, Clone, Debug)]
373pub struct WasmLocal {
374    pub index: u32,
375}
376
377#[derive(Copy, Clone, Debug)]
378pub struct WasmPointer {
379    base_variable: WasmLocal,
380    offset: u32,
381}
dada_codegen/cx/generate_expr/wasm_place_repr.rs

dada_codegen/cx/generate_expr/
wasm_place_repr.rs
