cleanup redundancy and duplication around UPLC/program generation

  The code is sufficiently complexe already, and the addition of the
  source map context makes it a bit more. That's okay. But duplicating
  entire chunks of logic with no changes when we have generics, is not!

  There's a similar cleanup to be done around the 'coverage' command,
  which is, in most part, a duplicate of the `check` command.

  Note that this changes also drop the 'finalize_minimal' and make the
  `no_optimize` flag now unused. It is not a mistake. That flag should
  not exist and shall be removed in upcoming commits.

Signed-off-by: KtorZ <matthias.benkort@gmail.com>

Author: KtorZ

crates/aiken-lang/src/ast.rs

@@ -2489,7 +2489,7 @@ impl chumsky::Span for Span {
 
 /// A source location that includes both the module name and the byte span.
 /// Used for tracking source locations across module boundaries during code generation.
-#[derive(Clone, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
 pub struct SourceLocation {
     /// The module name (e.g., "aiken/collection/list" or "validators/my_validator")
     pub module: String,
@@ -2510,11 +2510,11 @@ impl SourceLocation {
     }
 
     pub fn is_empty(&self) -> bool {
-        self.module.is_empty() && self.span.start == 0 && self.span.end == 0
+        self.module.is_empty() && self.span == Span::default()
     }
 }
 
-impl fmt::Debug for SourceLocation {
+impl fmt::Display for SourceLocation {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "{}:{:?}", self.module, self.span.range())
     }

crates/aiken-lang/src/gen_uplc.rs

@@ -56,10 +56,7 @@ use uplc::{
     builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST},
     builtins::DefaultFunction,
     machine::cost_model::ExBudget,
-    optimize::{
-        aiken_optimize_and_intern, aiken_optimize_minimal_with_context,
-        aiken_optimize_with_context, interner::CodeGenInterner, shrinker::NO_INLINE,
-    },
+    optimize::{aiken_optimize_and_intern, interner::CodeGenInterner, shrinker::NO_INLINE},
 };
 
 type Otherwise = Option<AirTree>;
@@ -138,18 +135,13 @@ impl<'a> CodeGenerator<'a> {
         }
     }
 
-    pub fn generate(&mut self, validator: &TypedValidator, module_name: &str) -> Program<Name> {
-        let (program, _term_with_spans) = self.generate_with_term(validator, module_name);
-        program
-    }
-
     /// Generate a validator program and return both the finalized program and the term with spans.
     /// The term with spans is useful for generating source maps.
-    pub fn generate_with_term(
+    pub fn generate(
         &mut self,
         validator: &TypedValidator,
         module_name: &str,
-    ) -> (Program<Name>, Term<Name, SourceLocation>) {
+    ) -> Program<Name, SourceLocation> {
         let context_name = "__context__".to_string();
         let context_name_interned = introduce_name(&mut self.interner, &context_name);
         validator.params.iter().for_each(|arg| {
@@ -190,20 +182,15 @@ impl<'a> CodeGenerator<'a> {
         });
 
         // Finalize with spans preserved for source map generation
-        let program_with_spans = self.finalize_with_spans(term);
-
-        // Strip spans for the program (compiled code doesn't need them)
-        let program = program_with_spans.clone().map_context(|_| ());
-
-        (program, program_with_spans.term)
+        self.finalize(term)
     }
 
     pub fn generate_raw(
         &mut self,
         body: &TypedExpr,
         args: &[TypedArg],
         module_name: &str,
-    ) -> Program<Name> {
+    ) -> Program<Name, SourceLocation> {
         args.iter().for_each(|arg| {
             arg.get_variable_name()
                 .iter()
@@ -233,50 +220,7 @@ impl<'a> CodeGenerator<'a> {
                 .for_each(|arg_name| self.interner.pop_text(arg_name.to_string()))
         });
 
-        self.finalize(term.map_context(|_| ()))
-    }
-
-    /// Returns the raw Term with Span context preserved.
-    /// This applies used functions (critical for recursive functions to work)
-    /// and is useful for source map generation.
-    pub fn generate_raw_with_spans(
-        &mut self,
-        body: &TypedExpr,
-        args: &[TypedArg],
-        module_name: &str,
-    ) -> Term<Name, SourceLocation> {
-        args.iter().for_each(|arg| {
-            arg.get_variable_name()
-                .iter()
-                .for_each(|arg_name| self.interner.intern(arg_name.to_string()))
-        });
-
-        let mut air_tree = self.build(body, module_name, &[]);
-
-        air_tree = AirTree::no_op(air_tree, SourceLocation::empty());
-
-        let full_tree = self.hoist_functions_to_validator(air_tree);
-
-        let full_vec = full_tree.to_vec();
-
-        let mut term = self.uplc_code_gen(full_vec);
-
-        term = if args.is_empty() {
-            term
-        } else {
-            cast_validator_args(term, args, &self.interner, &self.data_types)
-        };
-
-        args.iter().for_each(|arg| {
-            arg.get_variable_name()
-                .iter()
-                .for_each(|arg_name| self.interner.pop_text(arg_name.to_string()))
-        });
-
-        // Apply used functions (critical for recursive functions to work)
-        term = self.special_functions.apply_used_functions(term);
-
-        term
+        self.finalize(term)
     }
 
     fn new_program<T, C>(&self, term: Term<T, C>) -> Program<T, C> {
@@ -288,7 +232,12 @@ impl<'a> CodeGenerator<'a> {
         Program { version, term }
     }
 
-    fn finalize(&mut self, mut term: Term<Name>) -> Program<Name> {
+    /// Finalize and optimize a program while preserving source location context.
+    /// Uses the subset of optimizations that are generic over context type.
+    fn finalize<C: Default + Clone + PartialEq>(
+        &mut self,
+        mut term: Term<Name, C>,
+    ) -> Program<Name, C> {
         term = self.special_functions.apply_used_functions(term);
 
         let program = aiken_optimize_and_intern(self.new_program(term));
@@ -305,41 +254,6 @@ impl<'a> CodeGenerator<'a> {
         program
     }
 
-    /// Finalize and optimize a program while preserving source location context.
-    /// Uses the subset of optimizations that are generic over context type.
-    pub fn finalize_with_spans(
-        &mut self,
-        mut term: Term<Name, SourceLocation>,
-    ) -> Program<Name, SourceLocation> {
-        // Apply used functions (critical for recursive functions to work)
-        term = self.special_functions.apply_used_functions(term);
-
-        let program = aiken_optimize_with_context(self.new_program(term));
-
-        // Reset is important for reusing the generator instance
-        self.reset(true);
-
-        program
-    }
-
-    /// Finalize with minimal optimization, preserving source location context.
-    /// Skips performance optimizations like inlining and lambda reduction.
-    /// Produces larger but more readable code that maps directly to source.
-    pub fn finalize_minimal_with_spans(
-        &mut self,
-        mut term: Term<Name, SourceLocation>,
-    ) -> Program<Name, SourceLocation> {
-        // Apply used functions (critical for recursive functions to work)
-        term = self.special_functions.apply_used_functions(term);
-
-        let program = aiken_optimize_minimal_with_context(self.new_program(term));
-
-        // Reset is important for reusing the generator instance
-        self.reset(true);
-
-        program
-    }
-
     // TODO: pass mono types to build so monomorphization
     // happens as we build the AIR Tree rather than after
     fn build(
@@ -4259,7 +4173,7 @@ impl<'a> CodeGenerator<'a> {
 
                     let mut program = self.new_program(
                         self.special_functions
-                            .apply_used_functions(term.map_context(|_| ())),
+                            .apply_used_functions(term.strip_context()),
                     );
 
                     let mut interner = CodeGenInterner::new();
@@ -4275,7 +4189,7 @@ impl<'a> CodeGenerator<'a> {
                         .unwrap_or_else(|e| panic!("Failed to evaluate constant: {e:#?}"))
                         .try_into()
                         .unwrap();
-                    Some(result.map_context(|_| location.clone()))
+                    Some(result.map_context(&|_| location.clone()))
                 }
                 ValueConstructorVariant::ModuleFn {
                     name: func_name,
@@ -4438,7 +4352,7 @@ impl<'a> CodeGenerator<'a> {
                                 term = term.lambda(format!("arg_{index}"))
                             }
                         }
-                        Some(term.map_context(|_| location.clone()))
+                        Some(term.map_context(&|_| location.clone()))
                     }
                 }
             },
@@ -5048,7 +4962,7 @@ impl<'a> CodeGenerator<'a> {
                 };
 
                 if extract_constant(term.pierce_no_inlines_ref()).is_some() {
-                    let mut program = self.new_program(term.clone().map_context(|_| ()));
+                    let mut program = self.new_program(term.clone().strip_context());
 
                     let mut interner = CodeGenInterner::new();
 
@@ -5063,7 +4977,7 @@ impl<'a> CodeGenerator<'a> {
                         .expect("Evaluated on unwrapping a data constant and got an error");
 
                     let result: Term<Name> = evaluated_term.try_into().unwrap();
-                    term = result.map_context(|_| location.clone());
+                    term = result.map_context(&|_| location.clone());
                 }
 
                 Some(term)
@@ -5074,7 +4988,7 @@ impl<'a> CodeGenerator<'a> {
                 if extract_constant(term.pierce_no_inlines_ref()).is_some() {
                     term = builder::convert_type_to_data(term, &tipo, &self.data_types);
 
-                    let mut program = self.new_program(term.clone().map_context(|_| ()));
+                    let mut program = self.new_program(term.clone().strip_context());
 
                     let mut interner = CodeGenInterner::new();
 
@@ -5089,7 +5003,7 @@ impl<'a> CodeGenerator<'a> {
                         .expect("Evaluated on wrapping a constant into data and got an error");
 
                     let result: Term<Name> = evaluated_term.try_into().unwrap();
-                    term = result.map_context(|_| location.clone());
+                    term = result.map_context(&|_| location.clone());
                 } else {
                     term = builder::convert_type_to_data(term, &tipo, &self.data_types);
                 }
@@ -5294,7 +5208,7 @@ impl<'a> CodeGenerator<'a> {
                     let maybe_const = extract_constant(item.pierce_no_inlines_ref());
                     maybe_const.is_some()
                 }) {
-                    let mut program = self.new_program(term.clone().map_context(|_| ()));
+                    let mut program = self.new_program(term.clone().strip_context());
 
                     let mut interner = CodeGenInterner::new();
 
@@ -5309,7 +5223,7 @@ impl<'a> CodeGenerator<'a> {
                         .expect("Evaluated a constant record with args and got an error");
 
                     let result: Term<Name> = evaluated_term.try_into().unwrap();
-                    term = result.map_context(|_| location.clone());
+                    term = result.map_context(&|_| location.clone());
                 }
 
                 Some(term)

crates/aiken-lang/src/gen_uplc/builder.rs

@@ -150,7 +150,7 @@ impl CodeGenSpecialFuncs {
     ) -> Term<Name, C> {
         for func_name in self.used_funcs.iter() {
             // Convert the stored function (Term<Name>) to Term<Name, C>
-            let func = self.get_function(func_name).map_context(|_| C::default());
+            let func = self.get_function(func_name).map_context(&|_| C::default());
             term = term.lambda(func_name).apply(func);
         }
         term

crates/aiken-lang/src/test_framework.rs

@@ -68,18 +68,21 @@ impl Test {
         module_name: String,
         input_path: PathBuf,
     ) -> Test {
-        let program = generator.generate_raw(&test.body, &[], &module_name);
+        let program = generator
+            .generate_raw(&test.body, &[], &module_name)
+            .strip_context();
 
         let assertion = match test.body.try_into() {
             Err(..) => None,
             Ok(Assertion { bin_op, head, tail }) => {
                 let as_constant = |generator: &mut CodeGenerator<'_>, side| {
-                    Program::<NamedDeBruijn>::try_from(generator.generate_raw(
+                    Program::<NamedDeBruijn, _>::try_from(generator.generate_raw(
                         &side,
                         &[],
                         &module_name,
                     ))
                     .expect("failed to convert assertion operaand to NamedDeBruijn")
+                    .strip_context()
                     .eval(ExBudget::max())
                     .unwrap_constant()
                     .map(|cst| (cst, side.tipo()))
@@ -146,19 +149,25 @@ impl Test {
 
             let stripped_type_info = convert_opaque_type(&type_info, generator.data_types(), true);
 
-            let program = generator.clone().generate_raw(
-                &test.body,
-                &[TypedArg {
-                    tipo: stripped_type_info.clone(),
-                    ..parameter.clone().into()
-                }],
-                &module_name,
-            );
+            let program = generator
+                .clone()
+                .generate_raw(
+                    &test.body,
+                    &[TypedArg {
+                        tipo: stripped_type_info.clone(),
+                        ..parameter.clone().into()
+                    }],
+                    &module_name,
+                )
+                .strip_context();
 
             // NOTE: We need not to pass any parameter to the fuzzer/sampler here because the fuzzer
             // argument is a Data constructor which needs not any conversion. So we can just safely
             // apply onto it later.
-            let generator_program = generator.clone().generate_raw(&via, &[], &module_name);
+            let generator_program = generator
+                .clone()
+                .generate_raw(&via, &[], &module_name)
+                .strip_context();
 
             match kind {
                 RunnableKind::Bench => Test::Benchmark(Benchmark {

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

@@ -1,10 +1,12 @@
-use aiken_lang::{ast::SourceLocation, ast::TypedValidator, gen_uplc::CodeGenerator};
-use uplc::ast::{DeBruijn, Name, Program, Term};
+use aiken_lang::{
+    ast::{SourceLocation, TypedValidator},
+    gen_uplc::CodeGenerator,
+};
+use uplc::ast::{Name, Program};
 
 #[derive(Default)]
 pub struct MemoProgram {
-    program: Option<Program<DeBruijn>>,
-    term_with_spans: Option<Term<Name, SourceLocation>>,
+    program: Option<Program<Name, SourceLocation>>,
 }
 
 impl MemoProgram {
@@ -13,24 +15,14 @@ impl MemoProgram {
         generator: &mut CodeGenerator,
         def: &TypedValidator,
         module_name: &str,
-    ) -> Program<DeBruijn> {
+    ) -> Program<Name, SourceLocation> {
         match self.program.take() {
             None => {
-                let (program, term_with_spans) = generator.generate_with_term(def, module_name);
-                let new_program = program.to_debruijn().unwrap();
-
+                let new_program = generator.generate(def, module_name);
                 self.program.replace(new_program.clone());
-                self.term_with_spans.replace(term_with_spans);
-
                 new_program
             }
             Some(program) => program,
         }
     }
-
-    /// Get the term with source locations, if available.
-    /// This is only available after `get()` has been called.
-    pub fn get_term_with_spans(&self) -> Option<&Term<Name, SourceLocation>> {
-        self.term_with_spans.as_ref()
-    }
 }