From 5cf0a4d294353cbd027798c48ed8a770601b099d Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Thu, 15 Aug 2024 17:45:26 +0200
Subject: [PATCH] Fix validator arity check

  For now, this panics, but ideally, we should return a "unknown
  purpose" error when we cannot map the name to an arity.
---
 crates/aiken-lang/src/ast.rs                  | 36 +++++++++++++++++++
 crates/aiken-lang/src/tipo/infer.rs           |  4 +--
 .../aiken-project/src/blueprint/validator.rs  | 12 +++----
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs
index feaf1c05..80ec6eb0 100644
--- a/crates/aiken-lang/src/ast.rs
+++ b/crates/aiken-lang/src/ast.rs
@@ -25,6 +25,13 @@ pub const ENV_MODULE: &str = "env";
 pub const CONFIG_MODULE: &str = "config";
 pub const DEFAULT_ENV_MODULE: &str = "default";
 
+pub const PURPOSE_SPEND: &str = "spend";
+pub const PURPOSE_MINT: &str = "mint";
+pub const PURPOSE_WITHDRAW: &str = "withdraw";
+pub const PURPOSE_PUBLISH: &str = "publish";
+pub const PURPOSE_VOTE: &str = "vote";
+pub const PURPOSE_PROPOSE: &str = "propose";
+
 pub type TypedModule = Module<TypeInfo, TypedDefinition>;
 pub type UntypedModule = Module<(), UntypedDefinition>;
 
@@ -269,6 +276,25 @@ impl TypedFunction {
                     .and_then(|a| a.find_node(byte_index))
             })
     }
+
+    pub fn validator_arity(&self) -> usize {
+        if self.name == PURPOSE_SPEND
+            || self.name == PURPOSE_PUBLISH
+            || self.name == PURPOSE_PROPOSE
+        {
+            4
+        } else if self.name == PURPOSE_MINT
+            || self.name == PURPOSE_WITHDRAW
+            || self.name == PURPOSE_VOTE
+        {
+            3
+        } else {
+            panic!(
+                "tried to get validator arity of a non-validator function {}",
+                &self.name
+            );
+        }
+    }
 }
 
 impl TypedTest {
@@ -580,6 +606,16 @@ pub struct ModuleConstant<T> {
 pub type TypedValidator = Validator<Rc<Type>, TypedArg, TypedExpr>;
 pub type UntypedValidator = Validator<(), UntypedArg, UntypedExpr>;
 
+#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
+pub enum Purpose {
+    Spend,
+    Mint,
+    Withdraw,
+    Publish,
+    Propose,
+    Vote,
+}
+
 #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct Validator<T, Arg, Expr> {
     pub doc: Option<String>,
diff --git a/crates/aiken-lang/src/tipo/infer.rs b/crates/aiken-lang/src/tipo/infer.rs
index 6fafaf86..edb0a4ff 100644
--- a/crates/aiken-lang/src/tipo/infer.rs
+++ b/crates/aiken-lang/src/tipo/infer.rs
@@ -210,10 +210,10 @@ fn infer_definition(
                         typed_fun.arguments.drain(0..params_length);
 
                         // TODO: the expected number of args comes from the script purpose
-                        if typed_fun.arguments.len() < 2 || typed_fun.arguments.len() > 3 {
+                        if typed_fun.arguments.len() != typed_fun.validator_arity() {
                             return Err(Error::IncorrectValidatorArity {
                                 count: typed_fun.arguments.len() as u32,
-                                expected: 3,
+                                expected: typed_fun.validator_arity() as u32,
                                 location: typed_fun.location,
                             });
                         }
diff --git a/crates/aiken-project/src/blueprint/validator.rs b/crates/aiken-project/src/blueprint/validator.rs
index 34ca61cc..dd202b94 100644
--- a/crates/aiken-project/src/blueprint/validator.rs
+++ b/crates/aiken-project/src/blueprint/validator.rs
@@ -846,13 +846,11 @@ mod tests {
         let mut definitions = fixture_definitions();
         definitions.insert(
             &schema,
-            Schema::Data(Data::AnyOf(vec![
-                Constructor {
-                    index: 0,
-                    fields: vec![Declaration::Referenced(Reference::new("Bool")).into()],
-                }
-                .into(),
-            ]))
+            Schema::Data(Data::AnyOf(vec![Constructor {
+                index: 0,
+                fields: vec![Declaration::Referenced(Reference::new("Bool")).into()],
+            }
+            .into()]))
             .into(),
         );