Added ScaledFuzzer capabilities

crates/aiken-lang/src/ast/well_known.rs

@@ -9,6 +9,7 @@ pub const BOOL_CONSTRUCTORS: &[&str] = &["False", "True"];
 pub const BYTE_ARRAY: &str = "ByteArray";
 pub const DATA: &str = "Data";
 pub const FUZZER: &str = "Fuzzer";
+pub const SCALED_FUZZER: &str = "ScaledFuzzer";
 pub const G1_ELEMENT: &str = "G1Element";
 pub const G2_ELEMENT: &str = "G2Element";
 pub const INT: &str = "Int";
@@ -220,6 +221,53 @@ impl Type {
         })
     }
 
+    pub fn scaled_fuzzer(a: Rc<Type>) -> Rc<Type> {
+        let prng_annotation = Annotation::Constructor {
+            location: Span::empty(),
+            module: None,
+            name: PRNG.to_string(),
+            arguments: vec![],
+        };
+
+        Rc::new(Type::Fn {
+            args: vec![Type::prng(), Type::int()],
+            ret: Type::option(Type::tuple(vec![Type::prng(), a])),
+            alias: Some(
+                TypeAliasAnnotation {
+                    alias: SCALED_FUZZER.to_string(),
+                    parameters: vec!["a".to_string()],
+                    annotation: Annotation::Fn {
+                        location: Span::empty(),
+                        arguments: vec![
+                            prng_annotation.clone(),
+                            Annotation::Constructor {
+                                location: Span::empty(),
+                                module: None,
+                                name: INT.to_string(),
+                                arguments: vec![],
+                            },
+                        ],
+                        ret: Annotation::Constructor {
+                            location: Span::empty(),
+                            module: None,
+                            name: OPTION.to_string(),
+                            arguments: vec![Annotation::Tuple {
+                                location: Span::empty(),
+                                elems: vec![
+                                    prng_annotation,
+                                    Annotation::Var {
+                                        location: Span::empty(),
+                                        name: "a".to_string(),
+                                    },
+                                ],
+                            }],
+                        }.into(),
+                    },
+                }.into(),
+            ),
+        })
+    }
+
     pub fn map(k: Rc<Type>, v: Rc<Type>) -> Rc<Type> {
         Rc::new(Type::App {
             public: true,

crates/aiken-lang/src/builtins.rs

@@ -490,6 +490,22 @@ pub fn prelude(id_gen: &IdGenerator) -> TypeInfo {
         },
     );
 
+    // ScaledFuzzer
+    //
+    // pub type ScaledFuzzer<a> =
+    //   fn(PRNG, Int) -> Option<(PRNG, a)>
+    let scaled_fuzzer_value = Type::generic_var(id_gen.next());
+    prelude.types.insert(
+        well_known::SCALED_FUZZER.to_string(), 
+        TypeConstructor {
+            location: Span::empty(),
+            parameters: vec![scaled_fuzzer_value.clone()],
+            tipo: Type::scaled_fuzzer(scaled_fuzzer_value),
+            module: "".to_string(),
+            public: true,
+        },
+    );
+
     prelude
 }
 

crates/aiken-lang/src/test_framework.rs

@@ -331,10 +331,14 @@ impl PropertyTest {
     ) -> Result<Option<Counterexample<'a>>, FuzzerError> {
         let mut prng = initial_prng;
         let mut counterexample = None;
+        let mut iteration = 0;
 
         while *remaining > 0 && counterexample.is_none() {
-            (prng, counterexample) = self.run_once(prng, labels, plutus_version)?;
+            let (next_prng, cex) = self.run_once(prng, labels, plutus_version, iteration)?;
+            prng = next_prng;
+            counterexample = cex;
             *remaining -= 1;
+            iteration += 1;
         }
 
         Ok(counterexample)
@@ -345,11 +349,12 @@ impl PropertyTest {
         prng: Prng,
         labels: &mut BTreeMap<String, usize>,
         plutus_version: &'a PlutusVersion,
+        iteration: usize,
     ) -> Result<(Prng, Option<Counterexample<'a>>), FuzzerError> {
         use OnTestFailure::*;
 
         let (next_prng, value) = prng
-            .sample(&self.fuzzer.program)?
+            .sample(&self.fuzzer.program, iteration)?
             .expect("A seeded PRNG returned 'None' which indicates a fuzzer is ill-formed and implemented wrongly; please contact library's authors.");
 
         let mut result = self.eval(&value, plutus_version);
@@ -379,8 +384,8 @@ impl PropertyTest {
             let mut counterexample = Counterexample {
                 value,
                 choices: next_prng.choices(),
-                cache: Cache::new(|choices| {
-                    match Prng::from_choices(choices).sample(&self.fuzzer.program) {
+                cache: Cache::new(move |choices| {
+                    match Prng::from_choices(choices, iteration).sample(&self.fuzzer.program, iteration) {
                         Err(..) => Status::Invalid,
                         Ok(None) => Status::Invalid,
                         Ok(Some((_, value))) => {
@@ -447,20 +452,18 @@ impl PropertyTest {
         let mut prng = Prng::from_seed(seed);
 
         while remaining > 0 {
-            match prng.sample(&self.fuzzer.program) {
+            match prng.sample(&self.fuzzer.program, n - remaining) {
                 Ok(Some((new_prng, value))) => {
                     prng = new_prng;
-                    match self.eval(&value, plutus_version) {
-                        mut eval_result => {
-                            results.push(BenchmarkResult {
-                                test: self.clone(),
-                                cost: eval_result.cost(),
-                                success: true,
-                                traces: eval_result.logs().to_vec(),
-                            });
-                        }
-                    }
+                    let mut eval_result = self.eval(&value, plutus_version);
+                    results.push(BenchmarkResult {
+                        test: self.clone(),
+                        cost: eval_result.cost(),
+                        success: true,
+                        traces: eval_result.logs().to_vec(),
+                    });
                 }
+
                 Ok(None) => {
                     break;
                 }
@@ -501,8 +504,16 @@ impl PropertyTest {
 ///
 #[derive(Debug)]
 pub enum Prng {
-    Seeded { choices: Vec<u8>, uplc: PlutusData },
-    Replayed { choices: Vec<u8>, uplc: PlutusData },
+    Seeded { 
+        choices: Vec<u8>, 
+        uplc: PlutusData,
+        iteration: usize,
+    },
+    Replayed { 
+        choices: Vec<u8>, 
+        uplc: PlutusData,
+        iteration: usize,
+    },
 }
 
 impl Prng {
@@ -546,15 +557,16 @@ impl Prng {
             uplc: Data::constr(
                 Prng::SEEDED,
                 vec![
-                    Data::bytestring(digest.to_vec()), // Prng's seed
-                    Data::bytestring(vec![]),          // Random choices
+                    Data::bytestring(digest.to_vec()),
+                    Data::bytestring(vec![]),
                 ],
             ),
+            iteration: 0,
         }
     }
 
     /// Construct a Pseudo-random number generator from a pre-defined list of choices.
-    pub fn from_choices(choices: &[u8]) -> Prng {
+    pub fn from_choices(choices: &[u8], iteration: usize) -> Prng {
         Prng::Replayed {
             uplc: Data::constr(
                 Prng::REPLAYED,
@@ -564,23 +576,48 @@ impl Prng {
                 ],
             ),
             choices: choices.to_vec(),
+            iteration,
         }
     }
 
     /// Generate a pseudo-random value from a fuzzer using the given PRNG.
     pub fn sample(
         &self,
         fuzzer: &Program<Name>,
+        iteration: usize,
     ) -> Result<Option<(Prng, PlutusData)>, FuzzerError> {
+        // First try evaluating as a regular fuzzer
         let program = Program::<NamedDeBruijn>::try_from(fuzzer.apply_data(self.uplc())).unwrap();
-        let mut result = program.eval(ExBudget::max());
-        result
-            .result()
-            .map_err(|uplc_error| FuzzerError {
-                traces: result.logs(),
-                uplc_error,
-            })
-            .map(Prng::from_result)
+        let program_clone = program.clone();
+
+        let result = program.eval(ExBudget::max());
+
+        match result.result() {
+            Ok(term) if matches!(term, Term::Constant(_)) => {
+                // If we got a valid constant result, process it
+                Ok(Prng::from_result(term, iteration))
+            }
+            _ => {
+                // Use the cloned program for the second attempt
+                let program_with_iteration = Program::<NamedDeBruijn>::try_from(
+                    program_clone.apply_data(Data::integer(num_bigint::BigInt::from(iteration as i64)))
+                ).unwrap();
+
+                let mut result = program_with_iteration.eval(ExBudget::max());
+                match result.result() {
+                    Ok(term) if matches!(term, Term::Constant(_)) => {
+                        Ok(Prng::from_result(term, iteration))
+                    }
+                    Err(uplc_error) => {
+                        Err(FuzzerError {
+                            traces: result.logs(),
+                            uplc_error,
+                        })
+                    }
+                    _ => unreachable!("Fuzzer returned a malformed result? {result:#?}")
+                }
+            }
+        }
     }
 
     /// Obtain a Prng back from a fuzzer execution. As a reminder, fuzzers have the following
@@ -593,9 +630,9 @@ impl Prng {
     /// made during shrinking aren't breaking underlying invariants (if only, because we run out of
     /// values to replay). In such case, the replayed sequence is simply invalid and the fuzzer
     /// aborted altogether with 'None'.
-    pub fn from_result(result: Term<NamedDeBruijn>) -> Option<(Self, PlutusData)> {
+    pub fn from_result(result: Term<NamedDeBruijn>, iteration: usize) -> Option<(Self, PlutusData)> {
         /// Interpret the given 'PlutusData' as one of two Prng constructors.
-        fn as_prng(cst: &PlutusData) -> Prng {
+        fn as_prng(cst: &PlutusData, iteration: usize) -> Prng {
             if let PlutusData::Constr(Constr { tag, fields, .. }) = cst {
                 if *tag == 121 + Prng::SEEDED {
                     if let [PlutusData::BoundedBytes(bytes), PlutusData::BoundedBytes(choices)] =
@@ -612,6 +649,7 @@ impl Prng {
                                     PlutusData::BoundedBytes(vec![].into()),
                                 ],
                             ),
+                            iteration,
                         };
                     }
                 }
@@ -622,6 +660,7 @@ impl Prng {
                         return Prng::Replayed {
                             choices: choices.to_vec(),
                             uplc: cst.clone(),
+                            iteration,
                         };
                     }
                 }
@@ -635,7 +674,7 @@ impl Prng {
                 if *tag == 121 + Prng::SOME {
                     if let [PlutusData::Array(elems)] = &fields[..] {
                         if let [new_seed, value] = &elems[..] {
-                            return Some((as_prng(new_seed), value.clone()));
+                            return Some((as_prng(new_seed, iteration), value.clone()));
                         }
                     }
                 }

crates/aiken-lang/src/tests/check.rs

@@ -1872,6 +1872,45 @@ fn fuzzer_err_unify_3() {
     ))
 }
 
+#[test]
+fn scaled_fuzzer_ok_basic() {
+    let source_code = r#"
+        fn int() -> ScaledFuzzer<Int> { todo }
+        test prop(n via int()) { True }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok());
+}
+
+#[test]
+fn scaled_fuzzer_ok_explicit() {
+    let source_code = r#"
+        fn int(prng: PRNG, complexity: Int) -> Option<(PRNG, Int)> { todo }
+        test prop(n via int) { True }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok());
+}
+
+#[test]
+fn scaled_fuzzer_err_unify() {
+    let source_code = r#"
+        fn int() -> ScaledFuzzer<Int> { todo }
+        test prop(n: Bool via int()) { True }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((
+            _,
+            Error::CouldNotUnify {
+                situation: Some(UnifyErrorSituation::FuzzerAnnotationMismatch),
+                ..
+            }
+        ))
+    ));
+}
+
 #[test]
 fn utf8_hex_literal_warning() {
     let source_code = r#"