Differenciate between local and external function calls at IR level

src/back/msl.rs

@@ -618,7 +618,7 @@ impl<W: Write> Writer<W> {
                 }))
             }
             crate::Expression::Call {
-                ref name,
+                origin: crate::FunctionOrigin::External(ref name),
                 ref arguments,
             } => match name.as_str() {
                 "cos" | "normalize" | "sin" => {

src/front/glsl/mod.rs

@@ -85,7 +85,6 @@ struct Parser<'a> {
     globals_lookup: FastHashMap<String, Global>,
     constants: Arena<Constant>,
     functions: Arena<Function>,
-    function_lookup: FastHashMap<String, Handle<Function>>,
     shader_stage: ShaderStage,
 }
 
@@ -98,7 +97,6 @@ impl<'a> Parser<'a> {
             globals_lookup: FastHashMap::default(),
             constants: Arena::new(),
             functions: Arena::new(),
-            function_lookup: FastHashMap::default(),
             shader_stage,
         }
     }
@@ -347,7 +345,7 @@ impl<'a> Parser<'a> {
         }
 
         let handle = self.functions.append(Function {
-            name: Some(name.clone()),
+            name: Some(name),
             parameter_types,
             return_type: ty,
             global_usage: vec![],
@@ -355,9 +353,6 @@ impl<'a> Parser<'a> {
             expressions,
             body,
         });
-
-        self.function_lookup.insert(name, handle);
-
         Ok(handle)
     }
 
@@ -859,7 +854,7 @@ impl<'a> Parser<'a> {
                         })
                     }
                     _ => Ok(Expression::Call {
-                        name,
+                        origin: crate::FunctionOrigin::External(name),
                         arguments: args
                             .into_iter()
                             .map(|arg| {
@@ -901,7 +896,6 @@ impl<'a> Parser<'a> {
                         &self.globals,
                         locals,
                         &self.functions,
-                        &self.function_lookup,
                     )
                     .map_err(|e| Error { kind: e.into() })?;
                 let base_type = &self.types[type_handle];

src/front/spv.rs

@@ -330,6 +330,12 @@ struct LookupSampledImage {
     sampler: Handle<crate::Expression>,
 }
 
+struct DeferredFunctionCall {
+    source_handle: Handle<crate::Function>,
+    expr_handle: Handle<crate::Expression>,
+    dst_id: spirv::Word,
+}
+
 pub struct Parser<I> {
     data: I,
     state: ModuleState,
@@ -348,6 +354,7 @@ pub struct Parser<I> {
     lookup_sampled_image: FastHashMap<spirv::Word, LookupSampledImage>,
     lookup_function_type: FastHashMap<spirv::Word, LookupFunctionType>,
     lookup_function: FastHashMap<spirv::Word, Handle<crate::Function>>,
+    deferred_function_calls: Vec<DeferredFunctionCall>,
 }
 
 impl<I: Iterator<Item = u32>> Parser<I> {
@@ -369,6 +376,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
             lookup_sampled_image: FastHashMap::default(),
             lookup_function_type: FastHashMap::default(),
             lookup_function: FastHashMap::default(),
+            deferred_function_calls: Vec::new(),
         }
     }
 
@@ -511,6 +519,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
         type_arena: &Arena<crate::Type>,
         const_arena: &Arena<crate::Constant>,
         global_arena: &Arena<crate::GlobalVariable>,
+        local_function_calls: &mut FastHashMap<Handle<crate::Expression>, spirv::Word>,
     ) -> Result<(), Error> {
         loop {
             use spirv::Op;
@@ -1016,6 +1025,32 @@ impl<I: Iterator<Item = u32>> Parser<I> {
                         },
                     );
                 }
+                Op::FunctionCall => {
+                    inst.expect_at_least(4)?;
+                    let result_type_id = self.next()?;
+                    let result_id = self.next()?;
+                    let func_id = self.next()?;
+
+                    let mut arguments = Vec::with_capacity(inst.wc as usize - 4);
+                    for _ in 0..arguments.capacity() {
+                        let arg_id = self.next()?;
+                        arguments.push(self.lookup_expression.lookup(arg_id)?.handle);
+                    }
+                    let expr = crate::Expression::Call {
+                        // will be replaced by `Local()` after all the functions are parsed
+                        origin: crate::FunctionOrigin::External(String::new()),
+                        arguments,
+                    };
+                    let expr_handle = fun.expressions.append(expr);
+                    local_function_calls.insert(expr_handle, func_id);
+                    self.lookup_expression.insert(
+                        result_id,
+                        LookupExpression {
+                            handle: expr_handle,
+                            type_id: result_type_id,
+                        },
+                    );
+                }
                 Op::ExtInst => {
                     inst.expect_at_least(5)?;
                     let result_type_id = self.next()?;
@@ -1043,7 +1078,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
                         arguments.push(self.lookup_expression.lookup(arg_id)?.handle);
                     }
                     let expr = crate::Expression::Call {
-                        name: name.to_string(),
+                        origin: crate::FunctionOrigin::External(name.to_string()),
                         arguments,
                     };
                     self.lookup_expression.insert(
@@ -1184,6 +1219,22 @@ impl<I: Iterator<Item = u32>> Parser<I> {
             }
         }
 
+        for dfc in self.deferred_function_calls.drain(..) {
+            let dst_handle = *self.lookup_function.lookup(dfc.dst_id)?;
+            match *module
+                .functions
+                .get_mut(dfc.source_handle)
+                .expressions
+                .get_mut(dfc.expr_handle)
+            {
+                crate::Expression::Call {
+                    ref mut origin,
+                    arguments: _,
+                } => *origin = crate::FunctionOrigin::Local(dst_handle),
+                _ => unreachable!(),
+            }
+        }
+
         if !self.future_decor.is_empty() {
             log::warn!("Unused item decorations: {:?}", self.future_decor);
             self.future_decor.clear();
@@ -1973,6 +2024,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
             }
         }
         // read body
+        let mut local_function_calls = FastHashMap::default();
         loop {
             let fun_inst = self.next_inst()?;
             log::debug!("\t\t{:?}", fun_inst.op);
@@ -1985,6 +2037,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
                         &module.types,
                         &module.constants,
                         &module.global_variables,
+                        &mut local_function_calls,
                     )?;
                 }
                 spirv::Op::FunctionEnd => {
@@ -1998,6 +2051,14 @@ impl<I: Iterator<Item = u32>> Parser<I> {
         fun.global_usage =
             crate::GlobalUse::scan(&fun.expressions, &fun.body, &module.global_variables);
         let handle = module.functions.append(fun);
+        for (expr_handle, dst_id) in local_function_calls {
+            self.deferred_function_calls.push(DeferredFunctionCall {
+                source_handle: handle,
+                expr_handle,
+                dst_id,
+            });
+        }
+
         self.lookup_function.insert(fun_id, handle);
         self.lookup_expression.clear();
         self.lookup_sampled_image.clear();

src/front/wgsl.rs

@@ -328,7 +328,6 @@ impl<'a> ExpressionContext<'a, '_, '_> {
                 self.global_vars,
                 self.local_vars,
                 &Arena::new(),
-                &FastHashMap::default(),
             )
             .map_err(Error::InvalidResolve)
     }
@@ -548,7 +547,7 @@ impl Parser {
                         arguments.push(arg);
                     }
                     crate::Expression::Call {
-                        name: name.to_owned(),
+                        origin: crate::FunctionOrigin::External(name.to_owned()),
                         arguments,
                     }
                 } else {

src/lib.rs

@@ -338,6 +338,13 @@ pub enum DerivativeAxis {
     Width,
 }
 
+/// Origin of a function to call.
+#[derive(Clone, Debug, PartialEq)]
+pub enum FunctionOrigin {
+    Local(Handle<Function>),
+    External(String),
+}
+
 /// An expression that can be evaluated to obtain a value.
 #[derive(Clone, Debug)]
 pub enum Expression {
@@ -399,9 +406,9 @@ pub enum Expression {
         //modifier,
         expr: Handle<Expression>,
     },
-    /// Call a function defined in this module.
+    /// Call another function.
     Call {
-        name: String,
+        origin: FunctionOrigin,
         arguments: Vec<Handle<Expression>>,
     },
 }

src/proc/typifier.rs

@@ -1,6 +1,6 @@
 use crate::{
     arena::{Arena, Handle},
-    FastHashMap, Type, TypeInner, VectorSize,
+    Type, TypeInner, VectorSize,
 };
 
 pub struct Typifier {
@@ -32,7 +32,6 @@ impl Typifier {
         global_vars: &Arena<crate::GlobalVariable>,
         local_vars: &Arena<crate::LocalVariable>,
         functions: &Arena<crate::Function>,
-        function_lookup: &FastHashMap<String, Handle<crate::Function>>,
     ) -> Result<Handle<crate::Type>, ResolveError> {
         if self.types.len() <= expr_handle.index() {
             for (eh, expr) in expressions.iter().skip(self.types.len()) {
@@ -93,7 +92,6 @@ impl Typifier {
                             global_vars,
                             local_vars,
                             functions,
-                            function_lookup,
                         )?;
 
                         let (kind, width) = match types[image].inner {
@@ -168,7 +166,7 @@ impl Typifier {
                     crate::Expression::CrossProduct(_, _) => unimplemented!(),
                     crate::Expression::Derivative { .. } => unimplemented!(),
                     crate::Expression::Call {
-                        ref name,
+                        origin: crate::FunctionOrigin::External(ref name),
                         ref arguments,
                     } => match name.as_str() {
                         "distance" | "length" | "dot" => {
@@ -184,14 +182,14 @@ impl Typifier {
                         "normalize" | "fclamp" | "max" | "reflect" | "pow" | "clamp" | "mix" => {
                             self.types[arguments[0].index()]
                         }
-                        other => functions[*function_lookup.get(other).ok_or(
-                            ResolveError::FunctionNotDefined {
-                                name: other.to_string(),
-                            },
-                        )?]
+                        _ => return Err(ResolveError::FunctionNotDefined { name: name.clone() }),
+                    },
+                    crate::Expression::Call {
+                        origin: crate::FunctionOrigin::Local(handle),
+                        arguments: _,
+                    } => functions[handle]
                         .return_type
                         .ok_or(ResolveError::FunctionReturnsVoid)?,
-                    },
                 };
                 log::debug!("Resolving {:?} = {:?} : {:?}", eh, expr, ty);
                 self.types.push(ty);