From bc690c541039b84c9f20c2468a362a7234e9ec9e Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Fri, 17 Mar 2023 16:30:30 +0100
Subject: [PATCH] Generated wrapped schemas for multi-validators' redeemers

---
 crates/aiken-lang/src/format.rs               |  6 +---
 crates/aiken-lang/src/parser.rs               |  4 +--
 crates/aiken-project/src/blueprint/schema.rs  | 28 +++++++++++++++++++
 .../aiken-project/src/blueprint/validator.rs  | 21 ++++++--------
 4 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index 77113e5c..9881ff62 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/crates/aiken-lang/src/format.rs
@@ -521,11 +521,7 @@ impl<'comments> Formatter<'comments> {
             ))
             .nest(INDENT)
             .group()
-            .append(if other_fun.is_some() {
-                line()
-            } else {
-                nil()
-            })
+            .append(if other_fun.is_some() { line() } else { nil() })
             .append(
                 other_fun
                     .as_ref()
diff --git a/crates/aiken-lang/src/parser.rs b/crates/aiken-lang/src/parser.rs
index cd7cbd12..3a107ebc 100644
--- a/crates/aiken-lang/src/parser.rs
+++ b/crates/aiken-lang/src/parser.rs
@@ -306,8 +306,8 @@ pub fn validator_parser() -> impl Parser<Token, ast::UntypedDefinition, Error =
             ast::UntypedDefinition::Validator(ast::Validator {
                 doc: None,
                 fun: functions
-                  .next()
-                  .expect("unwrapping safe because there's 'at_least(1)' function"),
+                    .next()
+                    .expect("unwrapping safe because there's 'at_least(1)' function"),
                 other_fun: functions.next(),
                 location: Span {
                     start: span.start,
diff --git a/crates/aiken-project/src/blueprint/schema.rs b/crates/aiken-project/src/blueprint/schema.rs
index 95809ce9..6e3ef11f 100644
--- a/crates/aiken-project/src/blueprint/schema.rs
+++ b/crates/aiken-project/src/blueprint/schema.rs
@@ -2,6 +2,7 @@ use crate::blueprint::definitions::{Definitions, Reference};
 use crate::CheckedModule;
 use aiken_lang::{
     ast::{DataType, Definition, TypedDefinition},
+    builtins::wrapped_redeemer,
     tipo::{pretty, Type, TypeVar},
 };
 use owo_colors::{OwoColorize, Stream::Stdout};
@@ -12,6 +13,9 @@ use serde::{
 use std::ops::Deref;
 use std::{collections::HashMap, sync::Arc};
 
+// NOTE: Can be anything BUT 0
+pub const REDEEMER_DISCRIMINANT: usize = 1;
+
 #[derive(Debug, PartialEq, Eq, Clone, serde::Serialize, serde::Deserialize)]
 pub struct Annotated<T> {
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -71,6 +75,30 @@ impl<T> From<T> for Annotated<T> {
 }
 
 impl Annotated<Schema> {
+    pub fn as_wrapped_redeemer(
+        definitions: &mut Definitions<Annotated<Schema>>,
+        schema: Reference,
+        type_info: Arc<Type>,
+    ) -> Reference {
+        definitions
+            .register(
+                &wrapped_redeemer(type_info),
+                &HashMap::new(),
+                |_| {
+                    Ok::<_, Error>(Annotated {
+                        title: Some("Wrapped Redeemer".to_string()),
+                        description: Some("A redeemer wrapped in an extra constructor to make multi-validator detection possible on-chain.".to_string()),
+                        annotated: Schema::Data(Data::AnyOf(vec![Constructor {
+                            index: REDEEMER_DISCRIMINANT,
+                            fields: vec![schema.into()],
+                        }
+                        .into()])),
+                    })
+                },
+            )
+            .expect("cannot fail because Ok")
+    }
+
     pub fn from_type(
         modules: &HashMap<String, CheckedModule>,
         type_info: &Type,
diff --git a/crates/aiken-project/src/blueprint/validator.rs b/crates/aiken-project/src/blueprint/validator.rs
index 73958b0c..1170e5a4 100644
--- a/crates/aiken-project/src/blueprint/validator.rs
+++ b/crates/aiken-project/src/blueprint/validator.rs
@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use super::{
     definitions::{Definitions, Reference},
     error::Error,
@@ -8,7 +6,6 @@ use super::{
 use crate::module::{CheckedModule, CheckedModules};
 use aiken_lang::{
     ast::{TypedArg, TypedFunction, TypedValidator},
-    builtins::wrapped_redeemer,
     uplc::CodeGenerator,
 };
 use miette::NamedSource;
@@ -146,15 +143,15 @@ impl Validator<Reference, Annotated<Schema>> {
                         module.code.clone(),
                     ),
                 })
-                .map(|schema| match datum {
-                    Some(..) if is_multi_validator => {
-                        let wrap_redeemer_type = wrapped_redeemer(redeemer.tipo);
-
-                        definitions.register(&wrap_redeemer_type, &HashMap::new(), todo!());
-                    }
-                    _ => Argument {
-                        title: Some(redeemer.arg_name.get_label()),
-                        schema,
+                .map(|schema| Argument {
+                    title: Some(redeemer.arg_name.get_label()),
+                    schema: match datum {
+                        Some(..) if is_multi_validator => Annotated::as_wrapped_redeemer(
+                            &mut definitions,
+                            schema,
+                            redeemer.tipo.clone(),
+                        ),
+                        _ => schema,
                     },
                 })?,
             program: program.clone(),