1use std::collections::HashMap;
12
13use iridia::flow::{At, MathFunction, Node, NodeId, Op, ScalarKind, Values};
14
15use crate::compile::{Kind, binary_op, convert_op, kind_of, math_op, relational_op, unary_op};
16use crate::flat::MathOp;
17use crate::run::{bin, convert, math, unary};
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum Evaluated {
22 Word(u32),
24 Unknown(NodeId),
26}
27
28pub struct Evaluator<'v, L> {
30 values: &'v Values,
31 flush_denormals: bool,
32 leaf: L,
33 known: HashMap<NodeId, Evaluated>,
34}
35
36impl<L> std::fmt::Debug for Evaluator<'_, L> {
37 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38 f.debug_struct("Evaluator")
39 .field("flush_denormals", &self.flush_denormals)
40 .field("known", &self.known.len())
41 .finish_non_exhaustive()
42 }
43}
44
45enum Step {
47 Done(Evaluated),
48 Need(Vec<NodeId>),
49}
50
51impl<'v, L: FnMut(NodeId, &Node, &[u32]) -> Option<u32>> Evaluator<'v, L> {
52 pub fn new(values: &'v Values, flush_denormals: bool, leaf: L) -> Self {
57 Evaluator {
58 values,
59 flush_denormals,
60 leaf,
61 known: HashMap::new(),
62 }
63 }
64
65 pub fn evaluate(&mut self, id: NodeId) -> Result<Evaluated, String> {
67 let mut stack = vec![id];
68 while let Some(&top) = stack.last() {
69 if self.known.contains_key(&top) {
70 stack.pop();
71 continue;
72 }
73 match self.step(top)? {
74 Step::Done(v) => {
75 self.known.insert(top, v);
76 stack.pop();
77 }
78 Step::Need(more) => stack.extend(more),
82 }
83 }
84 self.known
85 .get(&id)
86 .copied()
87 .ok_or_else(|| "the walk ended without the node's value".into())
88 }
89
90 fn words(&self, ids: &[NodeId]) -> Option<Result<Vec<u32>, NodeId>> {
92 let mut out = Vec::with_capacity(ids.len());
93 for id in ids {
94 match self.known.get(id)? {
95 Evaluated::Word(w) => out.push(*w),
96 Evaluated::Unknown(leaf) => return Some(Err(*leaf)),
97 }
98 }
99 Some(Ok(out))
100 }
101
102 fn pending(&self, ids: &[NodeId]) -> Vec<NodeId> {
104 ids.iter()
105 .filter(|id| !self.known.contains_key(id))
106 .copied()
107 .collect()
108 }
109
110 fn reads(&self, id: NodeId) -> Result<Kind, String> {
112 kind_of(
113 self.values
114 .reads(id)
115 .ok_or_else(|| format!("an operation (node {id:?}) the graph states no kind for"))?,
116 )
117 }
118
119 fn leaf(&mut self, id: NodeId, node: &Node, at: &[NodeId]) -> Step {
121 match self.words(at) {
122 None => Step::Need(self.pending(at)),
123 Some(Err(_)) => Step::Done(Evaluated::Unknown(id)),
124 Some(Ok(words)) => Step::Done(match (self.leaf)(id, node, &words) {
125 Some(w) => Evaluated::Word(w),
126 None => Evaluated::Unknown(id),
127 }),
128 }
129 }
130
131 fn step(&mut self, id: NodeId) -> Result<Step, String> {
132 let node = self
133 .values
134 .get(id)
135 .ok_or_else(|| format!("node {id:?} is not of this graph"))?
136 .clone();
137 let ftz = self.flush_denormals;
138 Ok(match &node {
139 Node::Input { .. } | Node::Supplied { .. } | Node::Query { .. } => self.leaf(id, &node, &[]),
140 Node::Global { at, .. } => match at {
141 At::Component(_) => self.leaf(id, &node, &[]),
142 At::Indexed { index, .. } => self.leaf(id, &node, index),
143 },
144 Node::Texel {
145 coordinate, controls, ..
146 } => {
147 let at: Vec<NodeId> = coordinate.iter().chain(controls).copied().collect();
148 self.leaf(id, &node, &at)
149 }
150 Node::Literal(l) => Step::Done(Evaluated::Word(literal(l.0, self.values.kind(id))?)),
151 Node::Loop { .. } => return Err("a loop's value is not evaluated".into()),
152 Node::Merge { arms, .. } => match self.words(arms) {
153 None => Step::Need(self.pending(arms)),
154 Some(Err(leaf)) => Step::Done(Evaluated::Unknown(leaf)),
155 Some(Ok(words))
156 if let Some(&first) = words.first()
157 && words.iter().all(|w| *w == first) =>
158 {
159 Step::Done(Evaluated::Word(first))
160 }
161 Some(Ok(_)) => {
162 return Err(
163 "a merge whose arms differ: the graph does not state which condition picks which"
164 .into(),
165 );
166 }
167 },
168 Node::Op { op, operands } => match op {
169 Op::Select => {
170 let [condition, accept, reject] = operands[..] else {
171 return Err(format!("a select of {} operands", operands.len()));
172 };
173 match self.known.get(&condition) {
174 None => Step::Need(vec![condition]),
175 Some(Evaluated::Word(c)) => {
176 let chosen = if *c != 0 { accept } else { reject };
177 match self.known.get(&chosen) {
178 None => Step::Need(vec![chosen]),
179 Some(v) => Step::Done(*v),
180 }
181 }
182 Some(Evaluated::Unknown(leaf)) => {
184 let leaf = *leaf;
185 match self.words(&[accept, reject]) {
186 None => Step::Need(self.pending(&[accept, reject])),
187 Some(Ok(w))
188 if let &[a, b] = w.as_slice()
189 && a == b =>
190 {
191 Step::Done(Evaluated::Word(a))
192 }
193 Some(_) => Step::Done(Evaluated::Unknown(leaf)),
194 }
195 }
196 }
197 }
198 _ => match self.words(operands) {
199 None => Step::Need(self.pending(operands)),
200 Some(Err(leaf)) => Step::Done(Evaluated::Unknown(leaf)),
201 Some(Ok(w)) => Step::Done(Evaluated::Word(self.operation(id, *op, &w, ftz)?)),
202 },
203 },
204 })
205 }
206
207 fn operation(&self, id: NodeId, op: Op, w: &[u32], ftz: bool) -> Result<u32, String> {
209 let lane = |i: usize| -> Result<u32, String> {
210 w.get(i)
211 .copied()
212 .ok_or_else(|| format!("{op:?} with {} operands", w.len()))
213 };
214 Ok(match op {
215 Op::Unary(u) => unary(unary_op(u, self.reads(id)?), lane(0)?),
216 Op::Binary(b) => {
217 if w.len() != 2 {
218 return Err("a matrix product is not evaluated".into());
219 }
220 bin(binary_op(b, self.reads(id)?)?, lane(0)?, lane(1)?, ftz)?
221 }
222 Op::Convert(to) => match convert_op(self.reads(id)?, to)? {
223 Some(c) => convert(c, lane(0)?),
224 None => lane(0)?,
225 },
226 Op::Math(fun) => {
227 let m = math_op(fun, self.reads(id)?)?;
228 match m {
229 MathOp::Dot => {
230 let n = w.len() / 2;
231 let dot = || format!("a dot product of {} operands", w.len());
232 if !w.len().is_multiple_of(2) || !(1..=4).contains(&n) {
233 return Err(dot());
234 }
235 let (a, b) = w.split_at_checked(n).ok_or_else(dot)?;
236 math(
237 m,
238 reg(a),
239 reg(b),
240 [0; 4],
241 u8::try_from(n).map_err(|_| dot())?,
242 ftz,
243 )[0]
244 }
245 MathOp::Normalize | MathOp::PackHalf | MathOp::UnpackHalf => {
246 return Err(format!("{fun:?} mixes lanes: not evaluated"));
247 }
248 _ if lanes_mixed(fun) => return Err(format!("{fun:?} mixes lanes: not evaluated")),
249 _ => {
250 let at = |i: usize| [w.get(i).copied().unwrap_or(0), 0, 0, 0];
251 math(m, at(0), at(1), at(2), 1, ftz)[0]
252 }
253 }
254 }
255 Op::Relational(fun) => {
256 let m = relational_op(fun);
257 let lanes = u8::try_from(w.len())
258 .ok()
259 .filter(|n| (1..=4).contains(n))
260 .ok_or_else(|| format!("{fun:?} of {} lanes", w.len()))?;
261 math(m, reg(w), [0; 4], [0; 4], lanes, ftz)[0]
262 }
263 Op::Select => {
264 if lane(0)? != 0 {
265 lane(1)?
266 } else {
267 lane(2)?
268 }
269 }
270 Op::Index => {
271 return Err("an element picked by a run-time index: the graph does not state which candidate an index picks".into());
272 }
273 Op::Derivative(..) => return Err("a derivative is not evaluated".into()),
274 Op::Subgroup => return Err("a subgroup operation is not evaluated".into()),
275 })
276 }
277}
278
279fn reg(w: &[u32]) -> [u32; 4] {
281 let mut r = [0; 4];
282 for (out, v) in r.iter_mut().zip(w) {
283 *out = *v;
284 }
285 r
286}
287
288fn lanes_mixed(fun: MathFunction) -> bool {
291 matches!(
292 fun,
293 MathFunction::Cross
294 | MathFunction::Length
295 | MathFunction::Distance
296 | MathFunction::Outer
297 | MathFunction::Determinant
298 | MathFunction::Transpose
299 | MathFunction::Inverse
300 | MathFunction::Reflect
301 | MathFunction::Refract
302 | MathFunction::FaceForward
303 )
304}
305
306#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] fn literal(v: f64, kind: Option<ScalarKind>) -> Result<u32, String> {
309 let exact = |w: u32, back: f64| {
310 if back.to_bits() == v.to_bits() || (v.is_nan() && back.is_nan()) {
311 Ok(w)
312 } else {
313 Err(format!("literal {v} is no 32-bit {kind:?}"))
314 }
315 };
316 match kind {
317 Some(ScalarKind::Float) => {
318 let f = v as f32;
319 exact(f.to_bits(), f64::from(f))
320 }
321 Some(ScalarKind::Uint) => exact(v as u32, f64::from(v as u32)),
322 Some(ScalarKind::Sint) => exact((v as i32) as u32, f64::from(v as i32)),
323 Some(ScalarKind::Bool) => Ok(u32::from(v != 0.)),
324 other => Err(format!("a literal of kind {other:?}")),
325 }
326}