Generate source maps

crates/aiken-lang/src/gen_uplc.rs

@@ -72,8 +72,9 @@ pub struct CodeGenerator<'a> {
     data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>,
     module_types: IndexMap<&'a str, &'a TypeInfo>,
     module_src: IndexMap<&'a str, &'a (String, LineNumbers)>,
-    /// immutable option
+    /// immutable options
     tracing: TraceLevel,
+    source_maps: bool,
     /// mutable index maps that are reset
     defined_functions: IndexMap<FunctionAccessKey, ()>,
     special_functions: CodeGenSpecialFuncs,
@@ -90,6 +91,7 @@ impl<'a> CodeGenerator<'a> {
         &self.data_types
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         plutus_version: PlutusVersion,
         functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
@@ -98,6 +100,7 @@ impl<'a> CodeGenerator<'a> {
         module_types: IndexMap<&'a str, &'a TypeInfo>,
         module_src: IndexMap<&'a str, &'a (String, LineNumbers)>,
         tracing: Tracing,
+        source_maps: bool,
     ) -> Self {
         CodeGenerator {
             plutus_version,
@@ -107,6 +110,7 @@ impl<'a> CodeGenerator<'a> {
             module_types,
             module_src,
             tracing: tracing.trace_level(true),
+            source_maps,
             defined_functions: IndexMap::new(),
             special_functions: CodeGenSpecialFuncs::new(),
             code_gen_functions: IndexMap::new(),
@@ -234,6 +238,38 @@ impl<'a> CodeGenerator<'a> {
         body: &TypedExpr,
         module_build_name: &str,
         context: &[TypedExpr],
+    ) -> AirTree {
+        let inner = self.do_build(body, module_build_name, context);
+        if !self.source_maps {
+            return inner;
+        }
+        let location = Some(body.location());
+        /*
+        let location = match body {
+            TypedExpr::Assignment { location, .. } => Some(location),
+            TypedExpr::BinOp { location, ..} => Some(location),
+            TypedExpr::Call { location, .. } => Some(location),
+            _ => None,
+        };
+         */
+        if let Some(location) = location.filter(|l| *l != Span::empty()) {
+            let line_column =
+                get_line_columns_by_span(module_build_name, &location, &self.module_src);
+            let msg = AirTree::string(format!(
+                "__sourcemap {}.ak:{}:{}",
+                module_build_name, line_column.line, line_column.column
+            ));
+            AirTree::trace(msg, inner.return_type(), inner)
+        } else {
+            inner
+        }
+    }
+
+    fn do_build(
+        &mut self,
+        body: &TypedExpr,
+        module_build_name: &str,
+        context: &[TypedExpr],
     ) -> AirTree {
         if !context.is_empty() {
             let TypedExpr::Assignment {

crates/aiken-lang/src/lib.rs

@@ -14,6 +14,7 @@ pub mod line_numbers;
 pub mod parser;
 pub mod plutus_version;
 pub mod pretty;
+pub mod source_map;
 pub mod test_framework;
 pub mod tipo;
 pub mod utils;

crates/aiken-lang/src/source_map.rs

@@ -0,0 +1,144 @@
+use std::{borrow::Borrow, collections::BTreeMap, rc::Rc};
+
+use serde::Serialize;
+use uplc::{
+    ast::{Constant, DeBruijn, Name, Program, Term},
+    builtins::DefaultFunction,
+};
+
+#[derive(Debug, PartialEq, Default, Clone, Serialize)]
+#[serde(transparent)]
+pub struct SourceMap {
+    locations: BTreeMap<u64, String>,
+}
+
+impl SourceMap {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn extract(program: Program<DeBruijn>) -> (Self, Program<DeBruijn>) {
+        let name_program: Program<Name> = program.try_into().unwrap();
+
+        let mut trace = vec![];
+        let mut index = 0;
+        let mut locations = BTreeMap::new();
+        let new_term = traverse(&name_program.term, &mut trace, &mut index, &mut locations);
+
+        let new_program = Program {
+            version: name_program.version,
+            term: new_term,
+        };
+
+        (Self { locations }, new_program.try_into().unwrap())
+    }
+}
+
+fn traverse(
+    term: &Term<Name>,
+    trace: &mut Vec<Rc<Name>>,
+    index: &mut u64,
+    locations: &mut BTreeMap<u64, String>,
+) -> Term<Name> {
+    // If this is our source-location trace statement, track the location
+
+    // We are trying to match [[<trace> (con string "__sourcemap blahblahblah")] <body>],
+    // where <trace> is either the "trace" builtin or an alias for it
+    let Term::Apply {
+        function: outer_func,
+        argument: outer_arg,
+    } = term
+    else {
+        return base_traverse(term, trace, index, locations);
+    };
+
+    // Sometimes aiken uses lambdas to make an alias for builtin functions like trace.
+    // Keep track of any of those aliases for later.
+    if is_trace(outer_arg, trace) {
+        if let Term::Lambda { parameter_name, .. } = outer_func.borrow() {
+            trace.push(parameter_name.clone());
+        }
+    }
+
+    let Term::Apply { function, argument } = outer_func.borrow() else {
+        return base_traverse(term, trace, index, locations);
+    };
+
+    if !is_trace(function, trace) {
+        return base_traverse(term, trace, index, locations);
+    }
+    let Term::Constant(constant) = argument.borrow() else {
+        return base_traverse(term, trace, index, locations);
+    };
+    let Constant::String(str) = constant.borrow() else {
+        return base_traverse(term, trace, index, locations);
+    };
+    let Some(source_location) = str.strip_prefix("__sourcemap ") else {
+        return base_traverse(term, trace, index, locations);
+    };
+
+    // this is definitely a sourcemap trace statement
+    locations.insert(*index, source_location.to_string());
+    base_traverse(outer_arg, trace, index, locations)
+}
+
+fn base_traverse(
+    term: &Term<Name>,
+    trace: &mut Vec<Rc<Name>>,
+    index: &mut u64,
+    locations: &mut BTreeMap<u64, String>,
+) -> Term<Name> {
+    *index += 1;
+    match term {
+        Term::Delay(then) => {
+            let then = traverse(then, trace, index, locations).into();
+            Term::Delay(then)
+        }
+        Term::Lambda {
+            parameter_name,
+            body,
+        } => {
+            let parameter_name = parameter_name.clone();
+            let body = traverse(body, trace, index, locations).into();
+            Term::Lambda {
+                parameter_name,
+                body,
+            }
+        }
+        Term::Apply { function, argument } => {
+            let function = traverse(function, trace, index, locations).into();
+            let argument = traverse(argument, trace, index, locations).into();
+            Term::Apply { function, argument }
+        }
+        Term::Force(then) => {
+            let then = traverse(then, trace, index, locations).into();
+            Term::Force(then)
+        }
+        Term::Constr { tag, fields } => {
+            let tag = *tag;
+            let fields = fields
+                .iter()
+                .map(|f| traverse(f, trace, index, locations))
+                .collect();
+            Term::Constr { tag, fields }
+        }
+        Term::Case { constr, branches } => {
+            let constr = traverse(constr, trace, index, locations).into();
+            let branches = branches
+                .iter()
+                .map(|b| traverse(b, trace, index, locations))
+                .collect();
+            Term::Case { constr, branches }
+        }
+        Term::Var(_) | Term::Constant(_) | Term::Error | Term::Builtin(_) => term.clone(),
+    }
+}
+
+fn is_trace<T: Eq>(term: &Term<T>, trace: &Vec<Rc<T>>) -> bool {
+    match term {
+        Term::Force(then) => is_trace(then, trace),
+        Term::Builtin(DefaultFunction::Trace) => true,
+        Term::Var(name) => trace.contains(name),
+        _ => false,
+    }
+}

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

@@ -10,6 +10,7 @@ use aiken_lang::{
     ast::{well_known, Annotation, TypedArg, TypedFunction, TypedValidator},
     gen_uplc::CodeGenerator,
     plutus_version::PlutusVersion,
+    source_map::SourceMap,
     tipo::{collapse_links, Type},
 };
 use miette::NamedSource;
@@ -43,6 +44,9 @@ pub struct Validator {
     #[serde(skip_serializing_if = "Definitions::is_empty")]
     #[serde(default)]
     pub definitions: Definitions<Annotated<Schema>>,
+
+    #[serde(skip)]
+    pub source_map: SourceMap,
 }
 
 impl Validator {
@@ -197,6 +201,8 @@ impl Validator {
             (datum, Some(redeemer))
         };
 
+        let (source_map, program) = SourceMap::extract(program.get(generator, def, &module.name));
+
         Ok(Validator {
             title: format!("{}.{}.{}", &module.name, &def.name, &func.name,),
             description: func.doc.clone(),
@@ -207,8 +213,9 @@ impl Validator {
                 PlutusVersion::V1 => SerializableProgram::PlutusV1Program,
                 PlutusVersion::V2 => SerializableProgram::PlutusV2Program,
                 PlutusVersion::V3 => SerializableProgram::PlutusV3Program,
-            }(program.get(generator, def, &module.name)),
+            }(program),
             definitions,
+            source_map,
         })
     }
 }

crates/aiken-project/src/lib.rs

@@ -156,7 +156,7 @@ where
         }
     }
 
-    pub fn new_generator(&'_ self, tracing: Tracing) -> CodeGenerator<'_> {
+    pub fn new_generator(&'_ self, tracing: Tracing, source_maps: bool) -> CodeGenerator<'_> {
         CodeGenerator::new(
             self.config.plutus,
             utils::indexmap::as_ref_values(&self.functions),
@@ -165,6 +165,7 @@ where
             utils::indexmap::as_str_ref_values(&self.module_types),
             utils::indexmap::as_str_ref_values(&self.module_sources),
             tracing,
+            source_maps,
         )
     }
 
@@ -296,12 +297,23 @@ where
         fs::create_dir_all(&dir)?;
 
         for validator in &blueprint.validators {
-            let path = dir.clone().join(format!("{}.uplc", validator.title));
+            let uplc_path = dir.clone().join(format!("{}.uplc", validator.title));
+            let map_path = dir.clone().join(format!("{}.map", validator.title));
 
             let program = &validator.program;
             let program: Program<Name> = program.inner().try_into().unwrap();
 
-            fs::write(&path, program.to_pretty()).map_err(|error| Error::FileIo { error, path })?;
+            fs::write(&uplc_path, program.to_pretty()).map_err(|error| Error::FileIo {
+                error,
+                path: uplc_path,
+            })?;
+
+            let source_map = &validator.source_map;
+            let source_map = serde_json::to_string_pretty(source_map)?;
+            fs::write(&map_path, &source_map).map_err(|error| Error::FileIo {
+                error,
+                path: map_path,
+            })?
         }
 
         Ok(())
@@ -366,7 +378,7 @@ where
                     m.attach_doc_and_module_comments();
                 });
 
-                let mut generator = self.new_generator(options.tracing);
+                let mut generator = self.new_generator(options.tracing, uplc_dump);
 
                 let blueprint = Blueprint::new(&self.config, &self.checked_modules, &mut generator)
                     .map_err(Error::Blueprint)?;
@@ -543,7 +555,7 @@ where
                 _ => None,
             })
             .map(|(checked_module, func)| {
-                let mut generator = self.new_generator(tracing);
+                let mut generator = self.new_generator(tracing, false);
 
                 Export::from_function(
                     func,
@@ -939,7 +951,7 @@ where
             }
         }
 
-        let mut generator = self.new_generator(tracing);
+        let mut generator = self.new_generator(tracing, false);
 
         let mut tests = Vec::new();
 

crates/aiken-project/src/test_framework.rs

@@ -93,6 +93,7 @@ mod test {
             utils::indexmap::as_str_ref_values(&module_types),
             utils::indexmap::as_str_ref_values(&module_sources),
             Tracing::All(TraceLevel::Verbose),
+            false,
         );
 
         (

crates/aiken-project/src/tests/mod.rs

@@ -72,6 +72,7 @@ impl TestProject {
             utils::indexmap::as_str_ref_values(&self.module_types),
             utils::indexmap::as_str_ref_values(&self.module_sources),
             tracing,
+            false,
         )
     }