swamp_script_code_gen/
alloc_util.rs

1use crate::alloc::{FrameMemoryRegion, ScopeAllocator};
2use seq_map::SeqMap;
3use swamp_script_types::{AnonymousStructType, EnumVariantType, Type};
4use swamp_vm_types::{
5    FLOAT_SIZE, INT_SIZE, MAP_REFERENCE_SIZE, MAP_SIZE, MemoryAlignment, MemoryOffset, MemorySize,
6    STR_SIZE, VEC_REFERENCE_SIZE,
7};
8use tracing::{error, info};
9
10pub fn layout_struct(anon_struct: &AnonymousStructType) -> (MemorySize, MemoryAlignment) {
11    let mut calculated_offset = MemoryOffset(0);
12    let mut largest_alignment = MemoryAlignment::U8;
13    for (_name, field) in &anon_struct.field_name_sorted_fields {
14        let (field_size, field_alignment) = type_size_and_alignment(&field.field_type);
15        if field_alignment.greater_than(largest_alignment) {
16            largest_alignment = field_alignment;
17        }
18        calculated_offset.space(field_size, field_alignment);
19    }
20
21    let total_offset = calculated_offset.space(MemorySize(0), largest_alignment);
22
23    (total_offset.as_size(), largest_alignment)
24}
25
26#[must_use]
27pub fn layout_tuple(types: &Vec<Type>) -> (MemorySize, MemoryAlignment) {
28    let mut calculated_offset = MemoryOffset(0);
29    let mut largest_alignment = MemoryAlignment::U8;
30    for ty in types {
31        let (field_size, field_alignment) = type_size_and_alignment(&ty);
32        if field_alignment.greater_than(largest_alignment) {
33            largest_alignment = field_alignment;
34        }
35        calculated_offset.space(field_size, field_alignment);
36    }
37    let total_offset = calculated_offset.space(MemorySize(0), largest_alignment);
38    (total_offset.as_size(), largest_alignment)
39}
40
41#[must_use]
42pub fn layout_tuple_elements(
43    types: &Vec<Type>,
44) -> (MemorySize, MemoryAlignment, Vec<(MemoryOffset, MemorySize)>) {
45    let mut calculated_offset = MemoryOffset(0);
46    let mut largest_alignment = MemoryAlignment::U8;
47
48    let mut elements = Vec::new();
49    for ty in types {
50        let (field_size, field_alignment) = type_size_and_alignment(&ty);
51        if field_alignment.greater_than(largest_alignment) {
52            largest_alignment = field_alignment;
53        }
54        elements.push((calculated_offset, field_size));
55        calculated_offset.space(field_size, field_alignment);
56    }
57    let total_offset = calculated_offset.space(MemorySize(0), largest_alignment);
58    (total_offset.as_size(), largest_alignment, elements)
59}
60
61pub fn layout_union(variants: &SeqMap<String, EnumVariantType>) -> (MemorySize, MemoryAlignment) {
62    let mut max_variant_alignment = MemoryAlignment::U8;
63    let mut max_variant_size = MemorySize(0);
64    let mut calculated_offset = MemoryOffset(0);
65    for (_name, variant) in variants {
66        let (variant_size, variant_alignment) = match variant {
67            EnumVariantType::Struct(anon_struct) => layout_struct(&anon_struct.anon_struct),
68            EnumVariantType::Tuple(types) => layout_tuple(&types.fields_in_order),
69            EnumVariantType::Nothing(_) => (MemorySize(0), MemoryAlignment::U8),
70        };
71
72        if variant_alignment.greater_than(max_variant_alignment) {
73            max_variant_alignment = variant_alignment;
74        }
75
76        if variant_size.0 > max_variant_size.0 {
77            max_variant_size = variant_size;
78        }
79    }
80
81    (max_variant_size, max_variant_alignment)
82}
83
84pub fn is_vec(ty: &Type) -> Option<(MemorySize, MemoryAlignment)> {
85    match ty {
86        Type::NamedStruct(named_struct) => {
87            if named_struct.module_path == vec!["core-0.0.0".to_string()]
88                && named_struct.assigned_name.starts_with("Vec<")
89            {
90                Some((MemorySize(VEC_REFERENCE_SIZE), MemoryAlignment::U16))
91            } else {
92                None
93            }
94        }
95        _ => None,
96    }
97}
98
99pub fn is_map(ty: &Type) -> Option<(MemorySize, MemoryAlignment)> {
100    match ty {
101        Type::NamedStruct(named_struct) => {
102            if named_struct.module_path == vec!["core-0.0.0".to_string()]
103                && named_struct.assigned_name.starts_with("Map<")
104            {
105                Some((MemorySize(MAP_SIZE), MemoryAlignment::U16))
106            } else {
107                None
108            }
109        }
110        _ => None,
111    }
112}
113
114pub fn type_size_and_alignment(ty: &Type) -> (MemorySize, MemoryAlignment) {
115    match ty {
116        Type::Int => (MemorySize(INT_SIZE), MemoryAlignment::U32),
117        Type::Float => (MemorySize(FLOAT_SIZE), MemoryAlignment::U32),
118        Type::String => (MemorySize(STR_SIZE), MemoryAlignment::U16),
119        Type::Bool => (MemorySize(1), MemoryAlignment::U8),
120        Type::Unit => (MemorySize(0), MemoryAlignment::U8),
121        Type::Never => (MemorySize(0), MemoryAlignment::U8),
122        Type::Tuple(types) => layout_tuple(types),
123        Type::NamedStruct(named_struct) => {
124            if named_struct.module_path == vec!["core-0.0.0".to_string()]
125                && named_struct.assigned_name.starts_with("Vec<")
126            {
127                (MemorySize(VEC_REFERENCE_SIZE), MemoryAlignment::U16)
128            } else if named_struct.module_path == vec!["core-0.0.0".to_string()]
129                && named_struct.assigned_name.starts_with("Map<")
130            {
131                (MemorySize(MAP_REFERENCE_SIZE), MemoryAlignment::U16)
132            } else {
133                type_size_and_alignment(&Type::AnonymousStruct(
134                    named_struct.anon_struct_type.clone(),
135                ))
136            }
137        }
138        Type::Slice(value_type) => type_size_and_alignment(value_type),
139        Type::SlicePair(key_type, value_type) => {
140            layout_tuple(&vec![*key_type.clone(), *value_type.clone()])
141        }
142        Type::AnonymousStruct(anon_struct) => layout_struct(anon_struct),
143        Type::Enum(enum_type) => {
144            let (offset, alignment) = layout_union(&enum_type.variants);
145
146            let alignment_octets: usize = alignment.into();
147
148            (MemorySize(offset.0 + alignment_octets as u16), alignment)
149        }
150        Type::Function(_) => (MemorySize(2), MemoryAlignment::U16),
151        Type::Optional(inner_type) => {
152            let (offset, alignment) = type_size_and_alignment(inner_type);
153
154            let alignment_octets: usize = alignment.into();
155
156            (MemorySize(offset.0 + alignment_octets as u16), alignment)
157        }
158        Type::Generic(a, b) => {
159            error!(?a, ?b, "generic can not be generated");
160            panic!("generic is not supported")
161        }
162        Type::MutableReference(referenced_type) => type_size_and_alignment(referenced_type),
163        Type::Blueprint(_) => panic!("not supported"),
164        Type::Variable(_) => panic!("not supported"),
165        Type::External(_) => todo!(),
166    }
167}
168
169pub fn reserve_space_for_type(ty: &Type, allocator: &mut ScopeAllocator) -> FrameMemoryRegion {
170    let (size, alignment) = type_size_and_alignment(ty);
171
172    allocator.reserve(size, alignment)
173}