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

@@ -493,6 +493,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| {
+                cache: Cache::new(move |choices| {
-                    match Prng::from_choices(choices).sample(&self.fuzzer.program) {
+                    match Prng::from_choices(choices, iteration).sample(&self.fuzzer.program, iteration) {
                         Err(..) => Status::Invalid,
                         Ok(None) => Status::Invalid,
                         Ok(Some((_, value))) => {
@@ -447,11 +452,10 @@ 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) {
+                    let mut eval_result = self.eval(&value, plutus_version);
-                        mut eval_result => {
                     results.push(BenchmarkResult {
                         test: self.clone(),
                         cost: eval_result.cost(),
@@ -459,8 +463,7 @@ impl PropertyTest {
                         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 },
+    Seeded { 
-    Replayed { choices: Vec<u8>, uplc: PlutusData },
+        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(digest.to_vec()),
-                    Data::bytestring(vec![]),          // Random choices
+                    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,6 +576,7 @@ impl Prng {
                 ],
             ),
             choices: choices.to_vec(),
+            iteration,
         }
     }
 
@@ -571,16 +584,40 @@ impl 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());
+        let program_clone = program.clone();
-        result
+
-            .result()
+        let result = program.eval(ExBudget::max());
-            .map_err(|uplc_error| FuzzerError {
+        
+        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,
                         })
-            .map(Prng::from_result)
+                    }
+                    _ => 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

@@ -1875,6 +1875,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#"

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

@@ -357,12 +357,23 @@ fn infer_definition(
                         .map(|ann| hydrator.type_from_annotation(ann, environment))
                         .transpose()?;
 
-                    let (inferred_annotation, inferred_inner_type) = infer_fuzzer(
+                    let (inferred_annotation, inferred_inner_type) = match infer_fuzzer(
                         environment,
                         provided_inner_type.clone(),
                         &typed_via.tipo(),
                         &arg.via.location(),
-                    )?;
+                    ) {
+                        Ok(result) => Ok(result),
+                        Err(err) => match err {
+                            Error::CouldNotUnify { .. } => infer_scaled_fuzzer(
+                                environment,
+                                provided_inner_type.clone(),
+                                &typed_via.tipo(),
+                                &arg.via.location(),
+                            ),
+                            _ => Err(err),
+                        },
+                    }?;
 
                     // Ensure that the annotation, if any, matches the type inferred from the
                     // Fuzzer.
@@ -826,6 +837,71 @@ fn annotate_fuzzer(tipo: &Type, location: &Span) -> Result<Annotation, Error> {
     }
 }
 
+#[allow(clippy::result_large_err)]
+fn infer_scaled_fuzzer(
+    environment: &mut Environment<'_>,
+    expected_inner_type: Option<Rc<Type>>,
+    tipo: &Rc<Type>,
+    location: &Span,
+) -> Result<(Annotation, Rc<Type>), Error> {
+    let could_not_unify = || Error::CouldNotUnify {
+        location: *location,
+        expected: Type::scaled_fuzzer(
+            expected_inner_type
+                .clone()
+                .unwrap_or_else(|| Type::generic_var(0)),
+        ),
+        given: tipo.clone(),
+        situation: None,
+        rigid_type_names: Default::default(),
+    };
+
+    match tipo.borrow() {
+        Type::Fn { ret, args, .. } => {
+            // Check if this is a ScaledFuzzer (fn(PRNG, Int) -> Option<(PRNG, a)>)
+            if args.len() == 2 {
+                match ret.borrow() {
+                    Type::App { module, name, args: ret_args, .. }
+                        if module.is_empty() && name == "Option" && ret_args.len() == 1 => {
+                        if let Type::Tuple { elems, .. } = ret_args[0].borrow() {
+                            if elems.len() == 2 {
+                                let wrapped = &elems[1];
+
+                                // Unify with expected ScaledFuzzer type
+                                environment.unify(
+                                    tipo.clone(),
+                                    Type::scaled_fuzzer(wrapped.clone()),
+                                    *location,
+                                    false,
+                                )?;
+
+                                return Ok((annotate_fuzzer(wrapped, location)?, wrapped.clone()));
+                            }
+                        }
+                    }
+                    _ => ()
+                }
+            }
+
+            Err(could_not_unify())
+        }
+
+        Type::Var { tipo, alias } => match &*tipo.deref().borrow() {
+            TypeVar::Link { tipo } => infer_scaled_fuzzer(
+                environment,
+                expected_inner_type,
+                &Type::with_alias(tipo.clone(), alias.clone()),
+                location,
+            ),
+            _ => Err(Error::GenericLeftAtBoundary {
+                location: *location,
+            }),
+        },
+
+        Type::App { .. } | Type::Tuple { .. } | Type::Pair { .. } => Err(could_not_unify()),
+    }
+}
+
 fn put_params_in_scope<'a>(
     name: &'_ str,
     environment: &'a mut Environment,

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

@@ -9,7 +9,7 @@ Schema {
             Var {
                 tipo: RefCell {
                     value: Generic {
-                        id: 64,
+                        id: 65,
                     },
                 },
                 alias: None,

crates/aiken-project/src/lib.rs

@@ -1139,7 +1139,7 @@ where
                 Test::PropertyTest(property_test) => property_test
                     .benchmark(seed, property_max_success, plutus_version)
                     .into_iter()
-                    .map(|result| TestResult::Benchmark(result))
+                    .map(TestResult::Benchmark)
                     .collect::<Vec<_>>(),
             })
             .collect::<Vec<TestResult<(Constant, Rc<Type>), PlutusData>>>()

crates/aiken-project/src/snapshots/aiken_project__export__tests__cannot_export_generics.snap

@@ -9,7 +9,7 @@ Schema {
             Var {
                 tipo: RefCell {
                     value: Generic {
-                        id: 64,
+                        id: 65,
                     },
                 },
                 alias: None,

crates/aiken-project/src/telemetry/terminal.rs

@@ -230,7 +230,7 @@ impl EventListener for Terminal {
                     "     Complete"
                         .if_supports_color(Stderr, |s| s.bold())
                         .if_supports_color(Stderr, |s| s.green()),
-                    format!("benchmark results written to CSV")
+                    "benchmark results written to CSV"
                         .if_supports_color(Stderr, |s| s.bold())
                 );
             }