Skip to content

Commit 0cfdf97

Browse files
inner-daemonsinner-daemons
committed
More stuff
1 parent 3b406b9 commit 0cfdf97

1 file changed

Lines changed: 65 additions & 56 deletions

File tree

‎naga/src/inline/mod.rs‎

Lines changed: 65 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ mod adjust;
33
use core::unreachable;
44

55
use crate::{
6-
valid::{FunctionInfo, ModuleInfo, TypeFlags},
7-
AddressSpace, Block, Function, Handle, Module, Span, Statement, Type, TypeInner,
6+
AddressSpace, Block, Expression, Function, Handle, LocalVariable, Module, Span, Statement,
7+
Type, TypeInner,
88
};
99
use nt::FastHashMap;
1010

@@ -26,7 +26,7 @@ struct InlineState<'a> {
2626
pub module: &'a mut Module,
2727
/// Collect all functions that may need to be inlined, but inline them lazily, keeping
2828
/// track of which have already been inlined.
29-
pub funcs_needing_inline: FastHashMap<Handle<Function>, bool>,
29+
pub funcs_needing_inline: FastHashMap<Handle<Function>, Option<Function>>,
3030
}
3131

3232
pub fn type_needs_unrestricted_pointer_params(module: &Module, ty: Handle<Type>) -> bool {
@@ -43,9 +43,9 @@ pub fn type_needs_unrestricted_pointer_params(module: &Module, ty: Handle<Type>)
4343
}
4444

4545
pub fn inline(module: &mut Module, strategy: InlineStrategy) {
46-
let mut funcs_needing_inline: FastHashMap<Handle<Function>, bool>;
46+
let mut funcs_needing_inline: FastHashMap<Handle<Function>, Option<Function>>;
4747
if strategy == InlineStrategy::All {
48-
funcs_needing_inline = module.functions.iter().map(|e| (e.0, false)).collect();
48+
funcs_needing_inline = module.functions.iter().map(|e| (e.0, None)).collect();
4949
} else {
5050
funcs_needing_inline = Default::default();
5151
for (handle, func) in module.functions.iter() {
@@ -54,7 +54,7 @@ pub fn inline(module: &mut Module, strategy: InlineStrategy) {
5454
.iter()
5555
.any(|arg| type_needs_unrestricted_pointer_params(module, r#arg.ty));
5656
if needs_inline {
57-
funcs_needing_inline.insert(handle, false);
57+
funcs_needing_inline.insert(handle, None);
5858
}
5959
}
6060
};
@@ -98,28 +98,27 @@ impl InlineState<'_> {
9898
let bool_type = self.module.types.insert(
9999
Type {
100100
name: None,
101-
inner: crate::TypeInner::Scalar(crate::Scalar::BOOL),
101+
inner: TypeInner::Scalar(crate::Scalar::BOOL),
102102
},
103103
Span::UNDEFINED,
104104
);
105-
let mut is_done_var = function.local_variables.append(
106-
crate::LocalVariable {
105+
let is_done_var = function.local_variables.append(
106+
LocalVariable {
107107
name: None,
108108
ty: bool_type,
109109
init: None,
110110
},
111111
Span::UNDEFINED,
112112
);
113-
let is_done_var_ptr = function.expressions.append(
114-
crate::Expression::LocalVariable(is_done_var),
115-
Span::UNDEFINED,
116-
);
113+
let is_done_var_ptr = function
114+
.expressions
115+
.append(Expression::LocalVariable(is_done_var), Span::UNDEFINED);
117116
let false_val = function.expressions.append(
118-
crate::Expression::Literal(crate::Literal::Bool(false)),
117+
Expression::Literal(crate::Literal::Bool(false)),
119118
Span::UNDEFINED,
120119
);
121120
let true_val = function.expressions.append(
122-
crate::Expression::Literal(crate::Literal::Bool(true)),
121+
Expression::Literal(crate::Literal::Bool(true)),
123122
Span::UNDEFINED,
124123
);
125124
for (st, span) in function
@@ -134,56 +133,66 @@ impl InlineState<'_> {
134133
ref result,
135134
ref arguments,
136135
} => {
137-
if self.funcs_needing_inline.get(&handle) == Some(&false) {
138-
self.funcs_needing_inline.insert(handle, true);
139-
140-
let function = core::mem::take(&mut self.module.functions[handle]);
136+
if matches!(self.funcs_needing_inline.get(&handle), Some(None)) {
137+
let function = self.module.functions[handle].clone();
141138

142139
let function = self.inline_all_calls(function, true);
143140

144-
self.module.functions[handle] = function;
141+
self.funcs_needing_inline.insert(handle, Some(function));
145142
}
146-
if self.funcs_needing_inline.contains_key(&handle) {
147-
new_block.push(
148-
Statement::Store {
149-
pointer: is_done_var_ptr,
150-
value: false_val,
143+
let Some(prepared) = self.funcs_needing_inline.get(&handle) else {
144+
new_block.push(st, span);
145+
continue;
146+
};
147+
let prepared = prepared.as_ref().unwrap();
148+
let call_result_var = prepared.result.as_ref().map(|r| {
149+
let var = function.local_variables.append(
150+
LocalVariable {
151+
name: None,
152+
ty: r.ty,
153+
init: None,
151154
},
152155
span,
153156
);
154-
let func = &self.module.functions[handle];
155-
let mut pasted_body = func.body.clone();
156-
let expr_offset = function.expressions.len() as u32;
157-
let local_variable_offset = function.local_variables.len() as u32;
158-
for (_, expr, span) in func.expressions.iter_span() {
159-
function.expressions.append(expr.clone(), *span);
160-
}
161-
for (_, var, span) in func.local_variables.iter_span() {
162-
function.local_variables.append(var.clone(), *span);
163-
}
157+
function.expressions[result.unwrap()] = Expression::LocalVariable(var);
158+
var
159+
});
160+
new_block.push(
161+
Statement::Store {
162+
pointer: is_done_var_ptr,
163+
value: false_val,
164+
},
165+
span,
166+
);
167+
let mut pasted_body = prepared.body.clone();
168+
let expr_offset = function.expressions.len() as u32;
169+
let local_variable_offset = function.local_variables.len() as u32;
170+
for (_, expr, span) in prepared.expressions.iter_span() {
171+
function.expressions.append(expr.clone(), *span);
172+
}
173+
for (_, var, span) in prepared.local_variables.iter_span() {
174+
function.local_variables.append(var.clone(), *span);
175+
}
164176

165-
let adjust_info = adjust::AdjustInfo {
166-
function_args: arguments,
167-
expressions: &mut function.expressions,
168-
statements: &mut pasted_body.body,
169-
expr_offset,
170-
local_variable_offset,
171-
};
172-
adjust_info.adjust_all();
177+
let adjust_info = adjust::AdjustInfo {
178+
function_args: arguments,
179+
expressions: &mut function.expressions,
180+
statements: &mut pasted_body.body,
181+
expr_offset,
182+
local_variable_offset,
183+
};
184+
adjust_info.adjust_all();
173185

174-
// Copy-paste the function body inside of a for-loop.
175-
// We will have to reuse adjust_body from the compact module to realign all local variable and expression indices.
176-
new_block.push(
177-
Statement::Loop {
178-
body: pasted_body,
179-
continuing: Block::default(),
180-
break_if: Some(true_val),
181-
},
182-
span,
183-
);
184-
} else {
185-
new_block.push(st, span);
186-
}
186+
// Copy-paste the function body inside of a for-loop.
187+
new_block.push(
188+
Statement::Loop {
189+
body: pasted_body,
190+
continuing: Block::default(),
191+
break_if: Some(true_val),
192+
},
193+
span,
194+
);
195+
// TODO: should we `Emit` the call result expression which is now a LocalVariable expression?
187196
}
188197
_ => new_block.push(st, span),
189198
}

0 commit comments

Comments
 (0)