Add sourcemap support to the Aiken compiler

[Quantumplation]

Motivation

As Aiken matures, developers need better tooling for understanding what their contracts are actually doing at runtime. Three use cases drove this work:

1. Time-travel debuggers - Tools like https://github.com/SundaeSwap-finance/gastronomy can step through UPLC execution, but without source maps, users see raw UPLC with DeBruijn indices instead of their original Aiken code.
2. Code coverage - Understanding which parts of your contract are actually exercised by tests is table-stakes for most languages. This requires mapping executed UPLC nodes back to source locations.
3. Fuzzers - Coverage-guided fuzzing needs the same source mapping to know when new code paths are discovered.

All three require the same fundamental capability: given a UPLC node during execution, answer "where did this come from in the source?"

Approach

The core insight is that UPLC's Term type can carry metadata through compilation and execution. We:

1. Made Term generic over a context parameter: Term<T, C = ()> where C defaults to unit for backwards compatibility, but can hold source locations when needed.
2. Threaded source locations through code generation: As the compiler builds UPLC terms, it attaches SourceLocation (module name + byte span) to each node.
3. Made the CEK machine generic over context: The stepping interface preserves context through execution, so debuggers can query "what source location is this term from?" at each step.
4. Made optimization passes context-aware: Source locations survive inlining, lambda reduction, and other transformations.
5. Added source map generation: Post-compilation, we walk the term tree and build a mapping from node indices to source locations, which can be embedded in the blueprint or exported separately.

Alternative considered

We initially experimented with a trace-based approach: inject special Trace calls at key points during codegen, then strip them out before final output. The traces would carry source location strings that a debugger could intercept.

This had several problems, mainly that it didn't play well with optimizations and was very brittle.

The generic context approach is cleaner: the metadata is truly out-of-band and doesn't affect the compiled output when not needed.

What's included

• Term<T, C> with generic context parameter
• SourceLocation type for cross-module location tracking
• CEK machine stepping interface (get_initial_machine_state, step)
• Source map generation in aiken export --source-map
• aiken export --list to discover exportable functions/tests
• aiken export --no-optimize for debugging (larger output, but maps directly to source)
• aiken coverage command with LCOV output
• Variable name tracking in source maps (so debuggers can show x instead of i_42)

⚠️ Review note

This PR was heavily AI-assisted (Claude Code) because I was doing it over Christmas break; Without this, given my busy schedule, it likely wouldn't have happened at all, so I'm hoping the output is coherent enough to offset the influence of AI.

The high level design and approach was decided by @SupernaviX and @MicroProofs ahead of implementation, and all tests pass; Still, the volume of mechanical changes (particularly threading context through pattern matches) means a careful review is warranted. Key areas to scrutinize:

• crates/uplc/src/machine.rs - CEK machine changes
• crates/uplc/src/optimize/shrinker.rs - Optimization pass changes
• crates/aiken-project/src/blueprint/source_map.rs - Source map format (new public interface)
• crates/aiken-lang/src/gen_uplc.rs - Location threading through codegen

The test suite passes and the feature works end-to-end with Gastronomy, but fresh eyes on the implementation details would be valuable.

[Riley-Kilgore]

Really glad you had time to make this feature-set!

The only issue I can find relating to the sourcemap support itself is that the validator.title is constructed as {module_name}.{validator_name} where module_name derives from the module path. For a validator file at validators/governance/voting.ak, the title becomes validators/governance/voting.my_validator. When joined with dir, the resulting path is:

dir/validators/governance/voting.my_validator.sourcemap.json

Line 774 creates only dir with fs::create_dir_all(dir), but not dir/validators/governance/. The subsequent fs::write at line 788 will fail with "No such file or directory".

Example Test:

/// Nested module paths contain slashes which causes PathBuf::join
/// to create subdirectories instead of flat filenames.
///
/// This test SHOULD PASS but currently FAILS because the source map code
/// at lib.rs:782 uses the raw validator title (which contains slashes) as
/// a filename, causing PathBuf::join to create nested directories.
///
/// BUG: The code does:
///   let filename = format!("{}.sourcemap.json", validator.title);
///   let file_path = dir.join(&filename);
///
/// This test verifies the CORRECT behavior: source map files should be
/// written to a flat directory structure. It will FAIL until the bug is fixed.
#[test]
fn source_map_path_should_be_flat() {
    use std::path::PathBuf;

    // Simulate the current (buggy) behavior in lib.rs for external source maps
    let dir = PathBuf::from("/tmp/sourcemaps");

    // For a validator at validators/governance/voting.ak with name "my_vote",
    // the title becomes "governance/voting.my_vote.spend"
    let validator_title = "governance/voting.my_vote.spend";

    // Current buggy code: uses title directly as filename
    let filename = format!("{}.sourcemap.json", validator_title);
    let file_path = dir.join(&filename);
    let parent = file_path.parent().unwrap();

    // This assertion FAILS because the current code creates a nested path
    // Expected: /tmp/sourcemaps/governance-voting.my_vote.spend.sourcemap.json (flat)
    // Actual: /tmp/sourcemaps/governance/voting.my_vote.spend.sourcemap.json (nested)
    assert_eq!(
        parent.to_str().unwrap(),
        "/tmp/sourcemaps",
        "Source map path should be flat, but got nested directory: {}",
        parent.display()
    );
}

Separately, is it intended behavior for coverage reporting to only run property tests a single time? It seems we should have coverage reporting run property based tests a multitude of times as normal to properly report coverage.

[Quantumplation]

@Riley-Kilgore sorry, I didn't get notified when you responded! Or, more likely, it got buried. I think both points are fair, so I'll address them soon.

How did the AI generated code stand up?

[Riley-Kilgore]

Thanks for the updates, and for working through the earlier review.

I've worked through a second pass and found some more issues which seem like they'd be worth addressing.

1. Span/context loss in optimizer cleanup paths

• Where: crates/uplc/src/optimize/shrinker.rs (afterwards, run_once_pass)
• Problem: Name <-> NamedDeBruijn cleanup strips context and restores C::default().
• Impact: source maps / coverage lose span fidelity.

Full regression tests:

#[test]
fn afterwards_preserves_term_context() {
    #[derive(Clone, Debug, PartialEq, Eq)]
    struct Ctx(u8);

    impl Default for Ctx {
        fn default() -> Self {
            Ctx(0)
        }
    }

    let program: Program<Name, Ctx> = Program {
        version: (1, 0, 0),
        term: Term::<Name, Ctx>::integer(1.into()).map_context(|_| Ctx(1)),
    };

    let optimized = program.afterwards();

    assert_eq!(
        optimized.term.context(),
        Ctx(1),
        "expected optimizer to preserve term context through afterwards()"
    );
}

#[test]
fn run_once_pass_preserves_term_context() {
    #[derive(Clone, Debug, PartialEq, Eq)]
    struct Ctx(u8);

    impl Default for Ctx {
        fn default() -> Self {
            Ctx(0)
        }
    }

    let program: Program<Name, Ctx> = Program {
        version: (1, 0, 0),
        term: Term::<Name, Ctx>::integer(1.into()).map_context(|_| Ctx(1)),
    };

    let optimized = program.run_once_pass();

    assert_eq!(
        optimized.term.context(),
        Ctx(1),
        "expected optimizer to preserve term context through run_once_pass()"
    );
}

2. Source-map index stability after parameter application

• Where: crates/aiken-project/src/blueprint/source_map.rs (visit_post_order)
• Problem: Apply traversal is argument-first, shifting indices for existing function subtrees when wrapping with Apply.
• Impact: can invalidate existing sourceMap / sourceMapFile assumptions.

Full regression test:

#[test]
fn from_term_indices_do_not_shift_when_wrapping_in_apply() {
    use aiken_lang::{ast::Span, line_numbers::LineNumbers};

    let src = "x";
    let module_name = "test";
    let data = (src.to_string(), LineNumbers::new(src));
    let mut module_sources: IndexMap<&str, &(String, LineNumbers)> = IndexMap::new();
    module_sources.insert(module_name, &data);

    let var_x: Term<Name, AstSourceLocation> = Term::Var {
        context: AstSourceLocation::new(module_name, Span { start: 0, end: 1 }),
        name: Rc::new(Name::text("x_id_0")),
    };

    let sm_before = SourceMap::from_term(&var_x, module_name, &module_sources);
    let idx_before = *sm_before
        .locations
        .keys()
        .next()
        .expect("expected a single location entry for Var(x)");

    let wrapped: Term<Name, AstSourceLocation> = Term::Apply {
        function: Rc::new(var_x.clone()),
        argument: Rc::new(Term::Constant {
            value: Rc::new(uplc::ast::Constant::Unit),
            context: AstSourceLocation::empty(),
        }),
        context: AstSourceLocation::empty(),
    };

    let sm_after = SourceMap::from_term(&wrapped, module_name, &module_sources);
    let idx_after = *sm_after
        .locations
        .keys()
        .next()
        .expect("expected Var(x) to still have a location entry after wrapping");

    assert_eq!(
        idx_after, idx_before,
        "wrapping a term in Apply should not shift indices for the original function subtree"
    );
}

3. aiken coverage can silently pass when property fuzzer crashes

• Where:
  • crates/aiken-project/src/lib.rs in collect_tests_for_coverage (Err(_) => break)
  • crates/aiken-project/src/lib.rs in coverage (empty test list => success path)
• Problem: fuzzer sampling errors can be swallowed, producing success with empty coverage output.
• Impact: aiken check can fail while aiken coverage appears successful.

Full regression test:

#[cfg(test)]
mod coverage_regressions {
    use super::*;
    use crate::watch::with_project;
    use aiken_lang::ast::{TraceLevel, Tracing};
    use std::{
        fs,
        path::{Path, PathBuf},
        time::{SystemTime, UNIX_EPOCH},
    };

    fn mk_temp_project(name: &str, source: &str) -> PathBuf {
        let nonce = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("clock went backwards")
            .as_nanos();

        let root = std::env::temp_dir().join(format!("aiken_cov_{name}_{nonce}"));
        fs::create_dir_all(root.join("lib")).expect("failed to create lib dir");

        fs::write(
            root.join("aiken.toml"),
            r#"name = "demo/coverage-regression"
version = "0.0.0"
license = "Apache-2.0"
description = "tmp"

[repository]
user = "demo"
project = "coverage-regression"
platform = "github"
"#,
        )
        .expect("failed to write aiken.toml");

        fs::write(root.join("lib/tests.ak"), source).expect("failed to write tests.ak");

        root
    }

    #[test]
    fn coverage_fails_when_property_fuzzer_errors() {
        let root = mk_temp_project(
            "fuzzer_error",
            r#"
fn bad_fuzzer() -> Fuzzer<Int> {
  todo
}

test fuzzer_boom(n via bad_fuzzer()) {
  True
}
"#,
        );

        let output = root.join("coverage.lcov");
        let result = with_project(Some(Path::new(&root)), false, true, true, |p| {
            p.coverage(
                None,
                false,
                1,
                5,
                Tracing::All(TraceLevel::Silent),
                None,
                output.clone(),
            )
        });

        assert!(
            result.is_err(),
            "coverage should fail when a property fuzzer crashes; current behavior can return success with an empty report"
        );

        let _ = fs::remove_dir_all(root);
    }
}

4. Coverage fail once semantics evaluated per iteration instead of per property

• Where: crates/aiken-project/src/lib.rs (coverage result loop)
• Problem: coverage expands property tests into name[0], name[1], ... and checks each independently.
• Impact: test ... fail once can incorrectly fail coverage even when one iteration fails (which should satisfy fail once).

Full regression test:

#[test]
fn coverage_fail_once_should_aggregate_over_property_iterations() {
    use crate::coverage::CoverageData;
    use aiken_lang::ast::OnTestFailure;
    use uplc::machine::cost_model::ExBudget;

    let results = vec![
        (
            "mod.prop[0]".to_string(),
            OnTestFailure::SucceedImmediately,
            crate::coverage::CoverageResult {
                success: true,
                errored: false,
                returned_false: false,
                machine_error: false,
                coverage: CoverageData::new(),
                remaining_budget: ExBudget::max(),
                logs: vec![],
            },
        ),
        (
            "mod.prop[1]".to_string(),
            OnTestFailure::SucceedImmediately,
            crate::coverage::CoverageResult {
                success: false,
                errored: false,
                returned_false: true,
                machine_error: false,
                coverage: CoverageData::new(),
                remaining_budget: ExBudget::max(),
                logs: vec![],
            },
        ),
    ];

    let mut failed_tests = Vec::new();
    for (name, on_test_failure, cov_result) in results {
        let test_passed = match on_test_failure {
            OnTestFailure::SucceedEventually | OnTestFailure::SucceedImmediately => {
                cov_result.errored || cov_result.returned_false
            }
            OnTestFailure::FailImmediately => cov_result.success,
        };
        if !test_passed {
            failed_tests.push(name);
        }
    }

    assert!(
        failed_tests.is_empty(),
        "for `fail once`, coverage should aggregate iterations by base property and pass when any iteration fails"
    );
}

5. External sourceMapFile entries are not resolvable from blueprint location

• Where: crates/aiken-project/src/lib.rs (external source-map write path and source_map_file assignment)
• Problem: blueprint stores filename only (e.g. foo.foo.spend.sourcemap.json) and drops configured --source-map-dir prefix.
• Impact: consumers resolving sourceMapFile relative to plutus.json cannot find the files.

Full regression test:

#[cfg(test)]
mod source_map_reference_regressions {
    use crate::{options::SourceMapMode, watch::with_project};
    use aiken_lang::ast::{TraceLevel, Tracing};
    use std::{
        fs,
        path::{Path, PathBuf},
        time::{SystemTime, UNIX_EPOCH},
    };

    fn mk_temp_project(name: &str, source: &str) -> PathBuf {
        let nonce = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("clock went backwards")
            .as_nanos();

        let root = std::env::temp_dir().join(format!("aiken_sm_{name}_{nonce}"));
        fs::create_dir_all(root.join("validators")).expect("failed to create validators dir");

        fs::write(
            root.join("aiken.toml"),
            r#"name = "demo/source-map-regression"
version = "0.0.0"
license = "Apache-2.0"
description = "tmp"

[repository]
user = "demo"
project = "source-map-regression"
platform = "github"
"#,
        )
        .expect("failed to write aiken.toml");

        fs::write(root.join("validators/foo.ak"), source).expect("failed to write validator");
        root
    }

    #[test]
    fn external_source_map_file_reference_should_include_configured_directory() {
        let root = mk_temp_project(
            "source_map_file_reference",
            r#"
validator foo {
  spend(_d, _r, _o, _c) {
    True
  }
}
"#,
        );

        let map_dir_name = "maps_reference_case";
        let cwd_map_dir = std::env::current_dir()
            .expect("cwd")
            .join(map_dir_name);
        let project_map_dir = root.join(map_dir_name);
        let _ = fs::remove_dir_all(&cwd_map_dir);
        let _ = fs::remove_dir_all(&project_map_dir);

        let result = with_project(Some(Path::new(&root)), false, true, true, |p| {
            p.build(
                false,
                Tracing::All(TraceLevel::Silent),
                p.blueprint_path(None),
                None,
                SourceMapMode::External(PathBuf::from(map_dir_name)),
            )
        });
        assert!(result.is_ok(), "control check: build should succeed");

        let blueprint_path = root.join("plutus.json");
        let blueprint_src = fs::read_to_string(&blueprint_path).expect("failed to read blueprint");
        let blueprint_json: serde_json::Value =
            serde_json::from_str(&blueprint_src).expect("failed to parse blueprint json");

        let validators = blueprint_json
            .get("validators")
            .and_then(|v| v.as_array())
            .expect("expected validators array");

        for validator in validators {
            let source_map_file = validator
                .get("sourceMapFile")
                .and_then(|v| v.as_str())
                .expect("expected sourceMapFile entry");

            assert!(
                source_map_file.starts_with(&format!("{map_dir_name}/")),
                "sourceMapFile should include configured source map directory; got: {source_map_file}"
            );

            let resolved = blueprint_path
                .parent()
                .expect("blueprint should have parent dir")
                .join(source_map_file);

            assert!(
                resolved.exists(),
                "sourceMapFile should resolve relative to blueprint location; missing: {}",
                resolved.display()
            );
        }

        let _ = fs::remove_dir_all(&cwd_map_dir);
        let _ = fs::remove_dir_all(root);
    }
}

[KtorZ]

I've looked into it for a good part of the afternoon and I am afraid this will require more work. Not so much regarding the outcome itself (I didn't really get to the point of even trying it out), but the code:

• Quite a lot of logic has been duplicated across commands and internal functions. I've done a first round of cleanup for the internals; leaning more towards generics. The commands are still TODO. In particular, from the quick glance I had on the coverage command; it seems that it would better be done via adding flags to check and would likely allow to avoid the duplication altogether.

• I am a little skeptical about context now introduces a sheer amount of clones. In particular, removing context from a program doesn't look like a cheap operation given the size of the ASTs. Maybe it's not a big deal, but maybe it is. Especially around the test framework. I'd like to at least see high-level execution benchmarks (using e.g. hyperfine) to analyze the impact on the compiler.

• I have made the no_optimize flag redundant; because I don't quite agree with the idea of producing a non-optimized build. Unless you can clearly argue for it, I would prefer not to expose users to this.

[Quantumplation]

I'll take another pass at this, no worries :)

I agree with your first two suggestions, and partially agree with your third; I can definitely see not exposing this to users, but what about making it a secret flag that doesn't show up on man pages etc? it's quite useful when mucking around in the guts, such as I did here, or when adding new optimizations, as you can compare the semantics of the "pure" implementation of the UPLC (i.e. the "most obviously correct" version generated by the compiler just from the raw semantics of the code) with the "optimized" version;

[Quantumplation]

@KtorZ I removed the duplication, added a benchmark script (which showed this branch was only 3-7% slower), switched from strings to RC<&str> (got a few percentage points of speedup so we're within noise of main), and found a 24-35% speedup elsewhere in the compiler to make the whole thing a net win :)

I'll resolve merge conflicts after lunch.

Please don't spend a ton of time reviewing/refactoring, as I fear that will sour you on the idea overall 😅 Feel free to just skim and let me know what you're generally unhappy with (if anything), and I can iterate on that (potentially with @SupernaviX as a second set of eyes) before you do another deep dive review.