From 52c8ca6ceeb4c40a415c9d10fa87fe8d30e018a7 Mon Sep 17 00:00:00 2001 From: Riley-Kilgore Date: Thu, 8 Aug 2024 17:15:28 -0700 Subject: [PATCH] Added fix to Fn find_node --- crates/aiken-lang/src/expr.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/crates/aiken-lang/src/expr.rs b/crates/aiken-lang/src/expr.rs index 9948002b..6e131a5b 100644 --- a/crates/aiken-lang/src/expr.rs +++ b/crates/aiken-lang/src/expr.rs @@ -377,8 +377,20 @@ impl TypedExpr { expressions.iter().find_map(|e| e.find_node(byte_index)) } - TypedExpr::Fn { body, .. } => body - .find_node(byte_index) + TypedExpr::Fn { + body, + args, + return_annotation, + .. + } => args + .iter() + .find_map(|arg| arg.find_node(byte_index)) + .or_else(|| body.find_node(byte_index)) + .or_else(|| { + return_annotation + .as_ref() + .and_then(|a| a.find_node(byte_index)) + }) .or(Some(Located::Expression(self))), TypedExpr::Tuple {