Fix SSA parser to handle function as value

noir-lang · Jan 17, 2025 · 7f2cdef · 7f2cdef
1 parent da1fe6a
commit 7f2cdef
Showing 1 changed file with 29 additions and 13 deletions.
diff --git a/compiler/noirc_evaluator/src/ssa/parser/into_ssa.rs b/compiler/noirc_evaluator/src/ssa/parser/into_ssa.rs
@@ -24,7 +24,7 @@ impl ParsedSsa {
 struct Translator {
     builder: FunctionBuilder,
 
-    /// Maps function names to their IDs
+    /// Maps internal function names (e.g. "f1") to their IDs
     functions: HashMap<String, FunctionId>,
 
     /// Maps block names to their IDs
@@ -135,14 +135,14 @@ impl Translator {
 
         match block.terminator {
             ParsedTerminator::Jmp { destination, arguments } => {
-                let block_id = self.lookup_block(destination)?;
+                let block_id = self.lookup_block(&destination)?;
                 let arguments = self.translate_values(arguments)?;
                 self.builder.terminate_with_jmp(block_id, arguments);
             }
             ParsedTerminator::Jmpif { condition, then_block, else_block } => {
                 let condition = self.translate_value(condition)?;
-                let then_destination = self.lookup_block(then_block)?;
-                let else_destination = self.lookup_block(else_block)?;
+                let then_destination = self.lookup_block(&then_block)?;
+                let else_destination = self.lookup_block(&else_block)?;
                 self.builder.terminate_with_jmpif(condition, then_destination, else_destination);
             }
             ParsedTerminator::Return(values) => {
@@ -187,8 +187,17 @@ impl Translator {
                 let function_id = if let Some(id) = self.builder.import_intrinsic(&function.name) {
                     id
                 } else {
-                    let function_id = self.lookup_function(function)?;
-                    self.builder.import_function(function_id)
+                    match self.lookup_function(&function) {
+                        Ok(f) => self.builder.import_function(f),
+                        Err(e) => {
+                            if let Ok(v) = self.lookup_variable(&function) {
+                                // e.g. `v2 = call v0(v1) -> u32`, a lambda passed as a parameter
+                                v
+                            } else {
+                                return Err(e);
+                            }
+                        }
+                    }
                 };
 
                 let arguments = self.translate_values(arguments)?;
@@ -293,7 +302,14 @@ impl Translator {
             ParsedValue::NumericConstant { constant, typ } => {
                 Ok(self.builder.numeric_constant(constant, typ.unwrap_numeric()))
             }
-            ParsedValue::Variable(identifier) => self.lookup_variable(identifier),
+            ParsedValue::Variable(identifier) => self.lookup_variable(&identifier).or_else(|e| {
+                if let Ok(f) = self.lookup_function(&identifier) {
+                    // e.g. `v3 = call f1(f2, v0) -> u32`
+                    Ok(self.builder.import_function(f))
+                } else {
+                    Err(e)
+                }
+            }),
         }
     }
 
@@ -314,27 +330,27 @@ impl Translator {
         Ok(())
     }
 
-    fn lookup_variable(&mut self, identifier: Identifier) -> Result<ValueId, SsaError> {
+    fn lookup_variable(&mut self, identifier: &Identifier) -> Result<ValueId, SsaError> {
         if let Some(value_id) = self.variables[&self.current_function_id()].get(&identifier.name) {
             Ok(*value_id)
         } else {
-            Err(SsaError::UnknownVariable(identifier))
+            Err(SsaError::UnknownVariable(identifier.clone()))
         }
     }
 
-    fn lookup_block(&mut self, identifier: Identifier) -> Result<BasicBlockId, SsaError> {
+    fn lookup_block(&mut self, identifier: &Identifier) -> Result<BasicBlockId, SsaError> {
         if let Some(block_id) = self.blocks[&self.current_function_id()].get(&identifier.name) {
             Ok(*block_id)
         } else {
-            Err(SsaError::UnknownBlock(identifier))
+            Err(SsaError::UnknownBlock(identifier.clone()))
         }
     }
 
-    fn lookup_function(&mut self, identifier: Identifier) -> Result<FunctionId, SsaError> {
+    fn lookup_function(&mut self, identifier: &Identifier) -> Result<FunctionId, SsaError> {
         if let Some(function_id) = self.functions.get(&identifier.name) {
             Ok(*function_id)
         } else {
-            Err(SsaError::UnknownFunction(identifier))
+            Err(SsaError::UnknownFunction(identifier.clone()))
         }
     }