From 5b4fedd08416813bd2c7e0181db041357c0dc36b Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Fri, 1 Mar 2024 16:47:17 +0100
Subject: [PATCH] Add PRNG to the Prelude.

---
 crates/aiken-lang/src/ast.rs                  | 28 ++++++
 crates/aiken-lang/src/builtins.rs             | 87 +++++++++++++++++++
 crates/aiken-project/src/lib.rs               |  2 +
 .../acceptance_tests/095/lib/aiken/fuzz.ak    |  9 +-
 4 files changed, 119 insertions(+), 7 deletions(-)

diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs
index b1b82824..c5ffd510 100644
--- a/crates/aiken-lang/src/ast.rs
+++ b/crates/aiken-lang/src/ast.rs
@@ -258,6 +258,34 @@ pub struct FunctionAccessKey {
 pub type TypedDataType = DataType<Rc<Type>>;
 
 impl TypedDataType {
+    pub fn prng() -> Self {
+        DataType {
+            constructors: vec![
+                RecordConstructor {
+                    location: Span::empty(),
+                    name: "Seeded".to_string(),
+                    arguments: vec![],
+                    doc: None,
+                    sugar: false,
+                },
+                RecordConstructor {
+                    location: Span::empty(),
+                    name: "Replayed".to_string(),
+                    arguments: vec![],
+                    doc: None,
+                    sugar: false,
+                },
+            ],
+            doc: None,
+            location: Span::empty(),
+            name: "PRNG".to_string(),
+            opaque: false,
+            parameters: vec![],
+            public: true,
+            typed_parameters: vec![],
+        }
+    }
+
     pub fn ordering() -> Self {
         DataType {
             constructors: vec![
diff --git a/crates/aiken-lang/src/builtins.rs b/crates/aiken-lang/src/builtins.rs
index 1142b0fe..90c6b947 100644
--- a/crates/aiken-lang/src/builtins.rs
+++ b/crates/aiken-lang/src/builtins.rs
@@ -28,6 +28,8 @@ pub const STRING: &str = "String";
 pub const OPTION: &str = "Option";
 pub const ORDERING: &str = "Ordering";
 pub const REDEEMER_WRAPPER: &str = "RedeemerWrapper";
+pub const PRNG: &str = "PRNG";
+pub const FUZZER: &str = "FUZZER";
 
 /// Build a prelude that can be injected
 /// into a compiler pipeline
@@ -414,6 +416,72 @@ pub fn prelude(id_gen: &IdGenerator) -> TypeInfo {
         ),
     );
 
+    // PRNG
+    //
+    // pub type PRNG {
+    //   Seeded { seed: Int, choices: List<Int> }
+    //   Replayed { choices: List<Int> }
+    // }
+
+    prelude.types.insert(
+        PRNG.to_string(),
+        TypeConstructor {
+            location: Span::empty(),
+            parameters: vec![],
+            tipo: prng(),
+            module: "".to_string(),
+            public: true,
+        },
+    );
+
+    prelude.types_constructors.insert(
+        PRNG.to_string(),
+        vec!["Seeded".to_string(), "Replayed".to_string()],
+    );
+
+    let mut seeded_fields = HashMap::new();
+    seeded_fields.insert("seed".to_string(), (0, Span::empty()));
+    seeded_fields.insert("choices".to_string(), (1, Span::empty()));
+    prelude.values.insert(
+        "Seeded".to_string(),
+        ValueConstructor::public(
+            function(vec![int(), list(int())], prng()),
+            ValueConstructorVariant::Record {
+                module: "".into(),
+                name: "Seeded".to_string(),
+                field_map: Some(FieldMap {
+                    arity: 2,
+                    fields: seeded_fields,
+                    is_function: false,
+                }),
+                arity: 2,
+                location: Span::empty(),
+                constructors_count: 2,
+            },
+        ),
+    );
+
+    let mut replayed_fields = HashMap::new();
+    replayed_fields.insert("choices".to_string(), (0, Span::empty()));
+    prelude.values.insert(
+        "Replayed".to_string(),
+        ValueConstructor::public(
+            function(vec![list(int())], prng()),
+            ValueConstructorVariant::Record {
+                module: "".into(),
+                name: "Replayed".to_string(),
+                field_map: Some(FieldMap {
+                    arity: 1,
+                    fields: replayed_fields,
+                    is_function: false,
+                }),
+                arity: 1,
+                location: Span::empty(),
+                constructors_count: 2,
+            },
+        ),
+    );
+
     prelude
 }
 
@@ -1145,6 +1213,16 @@ pub fn prelude_data_types(id_gen: &IdGenerator) -> IndexMap<DataTypeKey, TypedDa
         option_data_type,
     );
 
+    // PRNG
+    let prng_data_type = TypedDataType::prng();
+    data_types.insert(
+        DataTypeKey {
+            module_name: "".to_string(),
+            defined_type: "PRNG".to_string(),
+        },
+        prng_data_type,
+    );
+
     data_types
 }
 
@@ -1215,6 +1293,15 @@ pub fn bool() -> Rc<Type> {
     })
 }
 
+pub fn prng() -> Rc<Type> {
+    Rc::new(Type::App {
+        args: vec![],
+        public: true,
+        name: PRNG.to_string(),
+        module: "".to_string(),
+    })
+}
+
 pub fn list(t: Rc<Type>) -> Rc<Type> {
     Rc::new(Type::App {
         public: true,
diff --git a/crates/aiken-project/src/lib.rs b/crates/aiken-project/src/lib.rs
index 8849a091..43a06893 100644
--- a/crates/aiken-project/src/lib.rs
+++ b/crates/aiken-project/src/lib.rs
@@ -847,6 +847,8 @@ where
 
                 let type_info = convert_opaque_type(&parameter.tipo, generator.data_types());
 
+                // TODO: Possibly refactor 'generate_raw' to accept arguments and do this wrapping
+                // itself.
                 let body = TypedExpr::Fn {
                     location: Span::empty(),
                     tipo: Rc::new(Type::Fn {
diff --git a/examples/acceptance_tests/095/lib/aiken/fuzz.ak b/examples/acceptance_tests/095/lib/aiken/fuzz.ak
index be0f647a..4c783742 100644
--- a/examples/acceptance_tests/095/lib/aiken/fuzz.ak
+++ b/examples/acceptance_tests/095/lib/aiken/fuzz.ak
@@ -2,11 +2,6 @@ use aiken/builtin
 
 const max_int: Int = 255
 
-pub type PRNG {
-  Seeded { seed: Int, choices: List<Int> }
-  Replayed { choices: List<Int> }
-}
-
 pub type Fuzzer<a> =
   fn(PRNG) -> Option<(PRNG, a)>
 
@@ -112,11 +107,11 @@ pub fn map4(
 
 // Builders
 
-fn any_bool() -> Fuzzer<Bool> {
+pub fn any_bool() -> Fuzzer<Bool> {
   any_int() |> map(fn(n) { n <= 127 })
 }
 
-fn any_list(fuzz_a: Fuzzer<a>) -> Fuzzer<List<a>> {
+pub fn any_list(fuzz_a: Fuzzer<a>) -> Fuzzer<List<a>> {
   any_bool()
     |> and_then(
          fn(continue) {