claw_codegen/
statement.rs

1use crate::code::{CodeGenerator, ExpressionAllocator};
2
3use super::GenerationError;
4use ast::{ExpressionId, NameId, Statement};
5use claw_ast as ast;
6use claw_resolver::ItemId;
7
8use cranelift_entity::EntityRef;
9use wasm_encoder as enc;
10use wasm_encoder::Instruction;
11
12pub trait EncodeStatement {
13    fn alloc_expr_locals(&self, allocator: &mut ExpressionAllocator)
14        -> Result<(), GenerationError>;
15
16    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError>;
17}
18
19impl EncodeStatement for Statement {
20    fn alloc_expr_locals(
21        &self,
22        allocator: &mut ExpressionAllocator,
23    ) -> Result<(), GenerationError> {
24        let statement: &dyn EncodeStatement = match self {
25            Statement::Let(statement) => statement,
26            Statement::Assign(statement) => statement,
27            Statement::Call(statement) => statement,
28            Statement::If(statement) => statement,
29            Statement::Return(statement) => statement,
30        };
31        statement.alloc_expr_locals(allocator)
32    }
33
34    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
35        let statement: &dyn EncodeStatement = match self {
36            Statement::Let(statement) => statement,
37            Statement::Assign(statement) => statement,
38            Statement::Call(statement) => statement,
39            Statement::If(statement) => statement,
40            Statement::Return(statement) => statement,
41        };
42        statement.encode(code_gen)
43    }
44}
45
46impl EncodeStatement for ast::Let {
47    fn alloc_expr_locals(
48        &self,
49        allocator: &mut ExpressionAllocator,
50    ) -> Result<(), GenerationError> {
51        allocator.alloc_child(self.expression)
52    }
53
54    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
55        encode_assignment(self.ident, self.expression, code_gen)
56    }
57}
58
59impl EncodeStatement for ast::Assign {
60    fn alloc_expr_locals(
61        &self,
62        allocator: &mut ExpressionAllocator,
63    ) -> Result<(), GenerationError> {
64        allocator.alloc_child(self.expression)
65    }
66
67    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
68        encode_assignment(self.ident, self.expression, code_gen)
69    }
70}
71
72impl EncodeStatement for ast::Call {
73    fn alloc_expr_locals(
74        &self,
75        allocator: &mut ExpressionAllocator,
76    ) -> Result<(), GenerationError> {
77        for arg in self.args.iter() {
78            allocator.alloc_child(*arg)?;
79        }
80        Ok(())
81    }
82
83    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
84        for arg in self.args.iter() {
85            code_gen.encode_child(*arg)?;
86        }
87        let item = code_gen.lookup_name(self.ident);
88        code_gen.encode_call(item, &self.args, None)?;
89        Ok(())
90    }
91}
92
93impl EncodeStatement for ast::If {
94    fn alloc_expr_locals(
95        &self,
96        allocator: &mut ExpressionAllocator,
97    ) -> Result<(), GenerationError> {
98        allocator.alloc_child(self.condition)?;
99        for statement in self.block.iter() {
100            allocator.alloc_statement(*statement)?;
101        }
102        Ok(())
103    }
104
105    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
106        code_gen.encode_child(self.condition)?;
107        let fields = code_gen.fields(self.condition)?;
108        assert_eq!(fields.len(), 1);
109        code_gen.read_expr_field(self.condition, &fields[0]);
110        code_gen.instruction(&Instruction::If(enc::BlockType::Empty));
111        for statement in self.block.iter() {
112            code_gen.encode_statement(*statement)?;
113        }
114        code_gen.instruction(&Instruction::End);
115        Ok(())
116    }
117}
118
119impl EncodeStatement for ast::Return {
120    fn alloc_expr_locals(
121        &self,
122        allocator: &mut ExpressionAllocator,
123    ) -> Result<(), GenerationError> {
124        if let Some(expression) = self.expression {
125            allocator.alloc_child(expression)?;
126        }
127        Ok(())
128    }
129
130    fn encode(&self, code_gen: &mut CodeGenerator) -> Result<(), GenerationError> {
131        if let Some(expression) = self.expression {
132            code_gen.encode_child(expression)?;
133
134            let fields = code_gen.fields(expression)?;
135            if code_gen.spill_return() {
136                for field in fields.iter() {
137                    code_gen.read_return_ptr()?;
138                    code_gen.field_address(field);
139                    code_gen.read_expr_field(expression, field);
140                    code_gen.write_mem(field);
141                }
142                code_gen.read_return_ptr()?;
143            } else {
144                for field in fields.iter() {
145                    code_gen.read_expr_field(expression, field);
146                }
147            }
148        }
149        code_gen.instruction(&Instruction::Return);
150        Ok(())
151    }
152}
153
154fn encode_assignment(
155    ident: NameId,
156    expression: ExpressionId,
157    code_gen: &mut CodeGenerator,
158) -> Result<(), GenerationError> {
159    code_gen.encode_child(expression)?;
160    let fields = code_gen.fields(expression)?;
161    match code_gen.lookup_name(ident) {
162        ItemId::ImportFunc(_) => panic!("Assigning to imported function isn't allowed!!"),
163        ItemId::Type(_) => panic!("Assigning to imported type isn't allowed!!"),
164        ItemId::Global(global) => {
165            // TODO handle composite globals
166            for field in fields {
167                code_gen.read_expr_field(expression, &field);
168                code_gen.instruction(&Instruction::GlobalSet(global.index() as u32));
169            }
170        }
171        ItemId::Param(_) => panic!("Assigning to parameters isn't allowed!!"),
172        ItemId::Local(local) => {
173            for field in fields {
174                code_gen.read_expr_field(expression, &field);
175                code_gen.write_local_field(local, &field);
176            }
177        }
178        ItemId::Function(_) => panic!("Assigning to functions isn't allowed!!"),
179    }
180    Ok(())
181}