diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs index 2a32a25e..1c94123d 100644 --- a/crates/aiken-lang/src/ast.rs +++ b/crates/aiken-lang/src/ast.rs @@ -266,6 +266,16 @@ pub struct Function { pub on_test_failure: OnTestFailure, } +impl Function { + pub fn is_spend(&self) -> bool { + self.name == HANDLER_SPEND + } + + pub fn is_mint(&self) -> bool { + self.name == HANDLER_MINT + } +} + impl TypedFunction { pub fn find_node(&self, byte_index: usize) -> Option> { self.arguments @@ -279,10 +289,6 @@ impl TypedFunction { }) } - pub fn is_spend(&self) -> bool { - self.name == HANDLER_SPEND - } - pub fn has_valid_purpose_name(&self) -> bool { self.name == HANDLER_SPEND || self.name == HANDLER_PUBLISH @@ -1245,6 +1251,15 @@ impl Annotation { } } + pub fn option(inner: Annotation) -> Self { + Annotation::Constructor { + name: "Option".to_string(), + module: None, + location: inner.location(), + arguments: vec![inner], + } + } + pub fn is_logically_equal(&self, other: &Annotation) -> bool { match self { Annotation::Constructor { diff --git a/crates/aiken-lang/src/tipo/environment.rs b/crates/aiken-lang/src/tipo/environment.rs index 7f008f60..a9b4c3a6 100644 --- a/crates/aiken-lang/src/tipo/environment.rs +++ b/crates/aiken-lang/src/tipo/environment.rs @@ -1275,10 +1275,9 @@ impl<'a> Environment<'a> { location, end_position: _, }) if kind.is_validator() => { - let default_annotation = |mut arg: UntypedArg| { + let default_annotation = |mut arg: UntypedArg, ann: Annotation| { if arg.annotation.is_none() { - arg.annotation = Some(Annotation::data(arg.location)); - + arg.annotation = Some(ann); arg } else { arg @@ -1287,12 +1286,29 @@ impl<'a> Environment<'a> { let mut handler_names = vec![]; + let params_len = params.len(); + for handler in handlers { let temp_params: Vec = params .iter() .cloned() .chain(handler.arguments.clone()) - .map(default_annotation) + .enumerate() + .map(|(ix, arg)| { + let is_datum = handler.is_spend() && ix == params_len; + let is_mint_policy = handler.is_mint() && ix == params_len + 1; + let location = arg.location; + default_annotation( + arg, + if is_datum { + Annotation::option(Annotation::data(location)) + } else if is_mint_policy { + Annotation::bytearray(location) + } else { + Annotation::data(location) + }, + ) + }) .collect(); handler_names.push(handler.name.clone()); @@ -1312,7 +1328,10 @@ impl<'a> Environment<'a> { .iter() .cloned() .chain(fallback.arguments.clone()) - .map(default_annotation) + .map(|arg| { + let location = arg.location; + default_annotation(arg, Annotation::data(location)) + }) .collect(); self.register_function(