Allow annotating Data for blueprint

This commit allows Data to be optionally annotated with a
  phantom-type. This doesn't change anything in codegen but we can now
  leverage this information to generate better blueprint schemas.

crates/aiken-lang/src/ast.rs

@@ -684,7 +684,7 @@ impl UnqualifiedImport {
 }
 
 // TypeAst
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Annotation {
     Constructor {
         location: Span,
@@ -1411,7 +1411,7 @@ impl Display for TraceLevel {
     }
 }
 
-#[derive(Copy, Clone, PartialEq, Eq)]
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct Span {
     pub start: usize,
     pub end: usize,

crates/aiken-lang/src/builtins.rs

@@ -38,6 +38,7 @@ pub fn prelude(id_gen: &IdGenerator) -> TypeInfo {
         types_constructors: HashMap::new(),
         values: HashMap::new(),
         accessors: HashMap::new(),
+        annotations: HashMap::new(),
     };
 
     // Int
@@ -423,6 +424,7 @@ pub fn plutus(id_gen: &IdGenerator) -> TypeInfo {
         types_constructors: HashMap::new(),
         values: HashMap::new(),
         accessors: HashMap::new(),
+        annotations: HashMap::new(),
     };
 
     for builtin in DefaultFunction::iter() {

crates/aiken-lang/src/tipo.rs

@@ -1,6 +1,6 @@
 use self::{environment::Environment, pretty::Printer};
 use crate::{
-    ast::{Constant, DefinitionLocation, ModuleKind, Span},
+    ast::{Annotation, Constant, DefinitionLocation, ModuleKind, Span},
     builtins::{G1_ELEMENT, G2_ELEMENT, MILLER_LOOP_RESULT},
     tipo::fields::FieldMap,
 };
@@ -755,6 +755,7 @@ pub struct TypeInfo {
     pub types_constructors: HashMap<String, Vec<String>>,
     pub values: HashMap<String, ValueConstructor>,
     pub accessors: HashMap<String, AccessorsMap>,
+    pub annotations: HashMap<Annotation, Rc<Type>>,
 }
 
 #[derive(Debug, Clone)]

crates/aiken-lang/src/tipo/environment.rs

@@ -72,6 +72,9 @@ pub struct Environment<'a> {
 
     pub unused_modules: HashMap<String, Span>,
 
+    /// A mapping from known annotations to their resolved type.
+    pub annotations: HashMap<Annotation, Rc<Type>>,
+
     /// Warnings
     pub warnings: &'a mut Vec<Warning>,
 }
@@ -88,6 +91,12 @@ impl<'a> Environment<'a> {
         self.scope = data.local_values;
     }
 
+    pub fn annotate(&mut self, return_type: Rc<Type>, annotation: &Annotation) -> Rc<Type> {
+        self.annotations
+            .insert(annotation.clone(), return_type.clone());
+        return_type
+    }
+
     /// Converts entities with a usage count of 0 to warnings
     pub fn convert_unused_to_warnings(&mut self) {
         let unused = self
@@ -657,6 +666,7 @@ impl<'a> Environment<'a> {
             imported_types: HashSet::new(),
             current_module,
             current_kind,
+            annotations: HashMap::new(),
             warnings,
             entity_usages: vec![HashMap::new()],
         }

crates/aiken-lang/src/tipo/hydrator.rs

@@ -123,7 +123,7 @@ impl Hydrator {
         environment: &mut Environment,
         unbounds: &mut Vec<&'a Span>,
     ) -> Result<Rc<Type>, Error> {
-        match annotation {
+        let return_type = match annotation {
             Annotation::Constructor {
                 location,
                 module,
@@ -153,8 +153,16 @@ impl Hydrator {
                     environment.increment_usage(name);
                 }
 
-                // Ensure that the correct number of arguments have been given to the constructor
-                if args.len() != parameters.len() {
+                // Ensure that the correct number of arguments have been given to the constructor.
+                //
+                // NOTE:
+                // We do consider a special case for 'Data', where we allow them to optionally
+                // carry a phantom type. That type has no effect whatsoever on the semantic (since
+                // anything can be cast to `Data` anyway) but it does provide some nice context for
+                // blueprint schema generation.
+                if args.len() != parameters.len() && !return_type.is_data()
+                    || args.len() > 1 && return_type.is_data()
+                {
                     return Err(Error::IncorrectTypeArity {
                         location: *location,
                         name: name.to_string(),
@@ -240,6 +248,8 @@ impl Hydrator {
 
                 Ok(tuple(typed_elems))
             }
-        }
+        }?;
+
+        Ok(environment.annotate(return_type, annotation))
     }
 }

crates/aiken-lang/src/tipo/infer.rs

@@ -126,6 +126,7 @@ impl UntypedModule {
             module_types_constructors: types_constructors,
             module_values: values,
             accessors,
+            annotations,
             ..
         } = environment;
 
@@ -141,6 +142,7 @@ impl UntypedModule {
                 types_constructors,
                 values,
                 accessors,
+                annotations,
                 kind,
                 package: package.to_string(),
             },

crates/aiken-project/src/blueprint/snapshots/aiken_project__blueprint__validator__tests__annotated_data.snap

@@ -0,0 +1,46 @@
+---
+source: crates/aiken-project/src/blueprint/validator.rs
+description: "Code:\n\npub type Foo {\n    foo: Int\n}\n\nvalidator {\n    fn annotated_data(datum: Data<Foo>, redeemer: Data, ctx: Void) {\n        True\n    }\n}\n"
+---
+{
+  "title": "test_module.annotated_data",
+  "datum": {
+    "title": "datum",
+    "schema": {
+      "$ref": "#/definitions/test_module~1Foo"
+    }
+  },
+  "redeemer": {
+    "title": "redeemer",
+    "schema": {
+      "$ref": "#/definitions/Data"
+    }
+  },
+  "compiledCode": "5833010000323222253330044a22930a99802a491856616c696461746f722072657475726e65642066616c736500136565734ae701",
+  "hash": "52a21f2b4f282074cb6c5aefef20d18c25f3657ca348c73875810c37",
+  "definitions": {
+    "Data": {
+      "title": "Data",
+      "description": "Any Plutus data."
+    },
+    "Int": {
+      "dataType": "integer"
+    },
+    "test_module/Foo": {
+      "title": "Foo",
+      "anyOf": [
+        {
+          "title": "Foo",
+          "dataType": "constructor",
+          "index": 0,
+          "fields": [
+            {
+              "title": "foo",
+              "$ref": "#/definitions/Int"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}

crates/aiken-project/src/blueprint/validator.rs

@@ -5,14 +5,15 @@ use super::{
     schema::{Annotated, Schema},
 };
 use crate::module::{CheckedModule, CheckedModules};
-use std::rc::Rc;
-
 use aiken_lang::{
-    ast::{TypedArg, TypedFunction, TypedValidator},
+    ast::{Annotation, TypedArg, TypedFunction, TypedValidator},
     gen_uplc::CodeGenerator,
+    tipo::Type,
 };
 use miette::NamedSource;
 use serde;
+use std::borrow::Borrow;
+use std::rc::Rc;
 use uplc::{
     ast::{Constant, DeBruijn, Program, Term},
     PlutusData,
@@ -96,7 +97,11 @@ impl Validator {
         let parameters = params
             .iter()
             .map(|param| {
-                Annotated::from_type(modules.into(), &param.tipo, &mut definitions)
+                Annotated::from_type(
+                    modules.into(),
+                    tipo_or_annotation(module, param),
+                    &mut definitions,
+                )
                 .map(|schema| Parameter {
                     title: Some(param.arg_name.get_label()),
                     schema,
@@ -118,23 +123,30 @@ impl Validator {
             parameters,
             datum: datum
                 .map(|datum| {
-                    Annotated::from_type(modules.into(), &datum.tipo, &mut definitions).map_err(
-                        |error| Error::Schema {
+                    Annotated::from_type(
+                        modules.into(),
+                        tipo_or_annotation(module, datum),
+                        &mut definitions,
+                    )
+                    .map_err(|error| Error::Schema {
                         error,
                         location: datum.location,
                         source_code: NamedSource::new(
                             module.input_path.display().to_string(),
                             module.code.clone(),
                         ),
-                        },
-                    )
+                    })
                 })
                 .transpose()?
                 .map(|schema| Parameter {
                     title: datum.map(|datum| datum.arg_name.get_label()),
                     schema,
                 }),
-            redeemer: Annotated::from_type(modules.into(), &redeemer.tipo, &mut definitions)
+            redeemer: Annotated::from_type(
+                modules.into(),
+                tipo_or_annotation(module, redeemer),
+                &mut definitions,
+            )
             .map_err(|error| Error::Schema {
                 error,
                 location: redeemer.location,
@@ -160,6 +172,29 @@ impl Validator {
     }
 }
 
+fn tipo_or_annotation<'a>(module: &'a CheckedModule, arg: &'a TypedArg) -> &'a Type {
+    match *arg.tipo.borrow() {
+        Type::App {
+            module: ref module_name,
+            name: ref type_name,
+            ..
+        } if module_name.is_empty() && &type_name[..] == "Data" => match arg.annotation {
+            Some(Annotation::Constructor { ref arguments, .. }) if !arguments.is_empty() => module
+                .ast
+                .type_info
+                .annotations
+                .get(
+                    arguments
+                        .first()
+                        .expect("guard ensures at least one element"),
+                )
+                .unwrap_or(&arg.tipo),
+            _ => &arg.tipo,
+        },
+        _ => &arg.tipo,
+    }
+}
+
 impl Validator {
     pub fn apply(
         self,
@@ -543,6 +578,23 @@ mod tests {
         );
     }
 
+    #[test]
+    fn annotated_data() {
+        assert_validator!(
+            r#"
+            pub type Foo {
+                foo: Int
+            }
+
+            validator {
+                fn annotated_data(datum: Data<Foo>, redeemer: Data, ctx: Void) {
+                    True
+                }
+            }
+            "#
+        );
+    }
+
     #[test]
     fn validate_arguments_integer() {
         let definitions = fixture_definitions();