From 4287fa3f4af105d12f9c887284ca038712459c4c Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Tue, 30 Jul 2024 18:01:54 -0400
Subject: [PATCH] feat: new formatting for validators v3

---
 crates/aiken-lang/src/format.rs     | 91 ++++++++++++++++-------------
 crates/aiken-lang/src/tipo/error.rs |  2 +-
 crates/aiken-lang/src/tipo/infer.rs |  2 +-
 3 files changed, 53 insertions(+), 42 deletions(-)

diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index 544f181e..1d176021 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/crates/aiken-lang/src/format.rs
@@ -239,8 +239,9 @@ impl<'comments> Formatter<'comments> {
                 handlers,
                 fallback,
                 params,
+                name,
                 ..
-            }) => self.definition_validator(params, handlers, fallback, *end_position),
+            }) => self.definition_validator(name, params, handlers, fallback, *end_position),
 
             Definition::Test(Function {
                 name,
@@ -581,57 +582,67 @@ impl<'comments> Formatter<'comments> {
 
     fn definition_validator<'a>(
         &mut self,
+        name: &'a str,
         params: &'a [UntypedArg],
         handlers: &'a [UntypedFunction],
         fallback: &'a UntypedFunction,
         end_position: usize,
     ) -> Document<'a> {
-        // validator(params)
-        let v_head = "validator".to_doc().append(if !params.is_empty() {
-            wrap_args(params.iter().map(|e| (self.fn_arg(e), false)))
-        } else {
-            nil()
-        });
+        // validator name(params)
+        let v_head = "validator"
+            .to_doc()
+            .append(" ")
+            .append(name)
+            .append(if !params.is_empty() {
+                wrap_args(params.iter().map(|e| (self.fn_arg(e), false)))
+            } else {
+                nil()
+            });
 
-        let fun_comments = self.pop_comments(fun.location.start);
-        let fun_doc_comments = self.doc_comments(fun.location.start);
-        let first_fn = self
+        let mut handler_docs = vec![];
+
+        for handler in handlers.iter() {
+            let fun_comments = self.pop_comments(handler.location.start);
+            let fun_doc_comments = self.doc_comments(handler.location.start);
+
+            let first_fn = self
+                .definition_fn(
+                    &handler.public,
+                    &handler.name,
+                    &handler.arguments,
+                    &handler.return_annotation,
+                    &handler.body,
+                    handler.end_position,
+                )
+                .group();
+
+            let first_fn = commented(fun_doc_comments.append(first_fn).group(), fun_comments);
+
+            handler_docs.push(first_fn);
+        }
+
+        let fallback_comments = self.pop_comments(fallback.location.start);
+        let fallback_doc_comments = self.doc_comments(fallback.location.start);
+
+        let fallback_fn = self
             .definition_fn(
-                &fun.public,
-                &fun.name,
-                &fun.arguments,
-                &fun.return_annotation,
-                &fun.body,
-                fun.end_position,
+                &fallback.public,
+                &fallback.name,
+                &fallback.arguments,
+                &fallback.return_annotation,
+                &fallback.body,
+                fallback.end_position,
             )
             .group();
-        let first_fn = commented(fun_doc_comments.append(first_fn).group(), fun_comments);
 
-        let other_fn = match other_fun {
-            None => nil(),
-            Some(other) => {
-                let other_comments = self.pop_comments(other.location.start);
-                let other_doc_comments = self.doc_comments(other.location.start);
+        let fallback_fn = commented(
+            fallback_doc_comments.append(fallback_fn).group(),
+            fallback_comments,
+        );
 
-                let other_fn = self
-                    .definition_fn(
-                        &other.public,
-                        &other.name,
-                        &other.arguments,
-                        &other.return_annotation,
-                        &other.body,
-                        other.end_position,
-                    )
-                    .group();
+        handler_docs.push(fallback_fn);
 
-                commented(other_doc_comments.append(other_fn).group(), other_comments)
-            }
-        };
-
-        let v_body = line()
-            .append(first_fn)
-            .append(if other_fun.is_some() { lines(2) } else { nil() })
-            .append(other_fn);
+        let v_body = line().append(join(handler_docs, lines(2)));
 
         let v_body = match printed_comments(self.pop_comments(end_position), false) {
             Some(comments) => v_body.append(lines(2)).append(comments).nest(INDENT),
diff --git a/crates/aiken-lang/src/tipo/error.rs b/crates/aiken-lang/src/tipo/error.rs
index fada415c..32a19911 100644
--- a/crates/aiken-lang/src/tipo/error.rs
+++ b/crates/aiken-lang/src/tipo/error.rs
@@ -1100,7 +1100,7 @@ impl ExtraData for Error {
             | Error::LastExpressionIsAssignment { .. }
             | Error::LogicalOpChainMissingExpr { .. }
             | Error::MissingVarInAlternativePattern { .. }
-            | Error::MultiValidatorEqualArgs { .. }
+            | Error::NonLocalClauseGuardVariable { .. }
             | Error::NotIndexable { .. }
             | Error::NotExhaustivePatternMatch { .. }
             | Error::NotFn { .. }
diff --git a/crates/aiken-lang/src/tipo/infer.rs b/crates/aiken-lang/src/tipo/infer.rs
index 1851374d..a2a65eba 100644
--- a/crates/aiken-lang/src/tipo/infer.rs
+++ b/crates/aiken-lang/src/tipo/infer.rs
@@ -9,7 +9,7 @@ use crate::{
     ast::{
         Annotation, ArgName, ArgVia, DataType, Definition, Function, ModuleConstant, ModuleKind,
         RecordConstructor, RecordConstructorArg, Tracing, TypeAlias, TypedArg, TypedDefinition,
-        TypedFunction, TypedModule, UntypedArg, UntypedDefinition, UntypedModule, Use, Validator,
+        TypedModule, UntypedArg, UntypedDefinition, UntypedModule, Use, Validator,
     },
     builtins::{self, fuzzer, generic_var},
     tipo::{expr::infer_function, Span, Type, TypeVar},