Automatic annotation of type signatures

crates/cli/Cargo.toml

@@ -60,10 +60,12 @@ roc_module.workspace = true
 roc_mono.workspace = true
 roc_packaging.workspace = true
 roc_parse.workspace = true
+roc_problem.workspace = true
 roc_region.workspace = true
 roc_reporting.workspace = true
 roc_target.workspace = true
 roc_tracing.workspace = true
+roc_types.workspace = true
 roc_repl_cli = { workspace = true, optional = true }
 roc_wasm_interp = { workspace = true, optional = true }
 

crates/cli/src/format.rs

@@ -1,17 +1,26 @@
 use std::ffi::OsStr;
 use std::io::Write;
+use std::ops::Range;
 use std::path::{Path, PathBuf};
 
 use bumpalo::Bump;
+use roc_can::abilities::{IAbilitiesStore, Resolved};
+use roc_can::expr::{DeclarationTag, Declarations, Expr};
 use roc_error_macros::{internal_error, user_error};
 use roc_fmt::def::fmt_defs;
 use roc_fmt::header::fmt_header;
 use roc_fmt::Buf;
 use roc_fmt::MigrationFlags;
+use roc_load::{ExecutionMode, FunctionKind, LoadConfig, LoadedModule, LoadingProblem, Threading};
+use roc_module::symbol::{Interns, ModuleId};
+use roc_packaging::cache::{self, RocCacheDir};
 use roc_parse::ast::{FullAst, SpacesBefore};
 use roc_parse::header::parse_module_defs;
 use roc_parse::normalize::Normalize;
 use roc_parse::{header, parser::SyntaxError, state::State};
+use roc_reporting::report::{RenderTarget, DEFAULT_PALETTE};
+use roc_target::Target;
+use roc_types::subs::{Content, Subs, Variable};
 
 #[derive(Copy, Clone, Debug)]
 pub enum FormatMode {
@@ -263,10 +272,162 @@ fn fmt_all<'a>(buf: &mut Buf<'a>, ast: &'a FullAst) {
     buf.fmt_end_of_file();
 }
 
+pub fn annotate_file(arena: &Bump, file: PathBuf) -> Result<(), LoadingProblem> {
+    let load_config = LoadConfig {
+        target: Target::default(),
+        function_kind: FunctionKind::from_env(),
+        render: RenderTarget::ColorTerminal,
+        palette: DEFAULT_PALETTE,
+        threading: Threading::AllAvailable,
+        exec_mode: ExecutionMode::Check,
+    };
+
+    let mut loaded = roc_load::load_and_typecheck(
+        arena,
+        file.clone(),
+        None,
+        RocCacheDir::Persistent(cache::roc_cache_dir().as_path()),
+        load_config,
+    )?;
+
+    let buf = annotate_module(&mut loaded);
+
+    std::fs::write(&file, buf.as_str())
+        .unwrap_or_else(|e| internal_error!("failed to write annotated file to {file:?}: {e}"));
+
+    Ok(())
+}
+
+fn annotate_module(loaded: &mut LoadedModule) -> String {
+    let (decls, subs, abilities) =
+        if let Some(decls) = loaded.declarations_by_id.get(&loaded.module_id) {
+            let subs = loaded.solved.inner_mut();
+            let abilities = &loaded.abilities_store;
+
+            (decls, subs, abilities)
+        } else if let Some(checked) = loaded.typechecked.get_mut(&loaded.module_id) {
+            let decls = &checked.decls;
+            let subs = checked.solved_subs.inner_mut();
+            let abilities = &checked.abilities_store;
+
+            (decls, subs, abilities)
+        } else {
+            internal_error!("Could not find file's module");
+        };
+
+    let src = &loaded
+        .sources
+        .get(&loaded.module_id)
+        .unwrap_or_else(|| internal_error!("Could not find the file's source"))
+        .1;
+
+    let mut edits = annotation_edits(
+        decls,
+        subs,
+        abilities,
+        src,
+        loaded.module_id,
+        &loaded.interns,
+    );
+    edits.sort_by_key(|(offset, _)| *offset);
+
+    let mut buffer = String::new();
+    let mut file_progress = 0;
+
+    for (position, edit) in edits {
+        debug_assert!(
+            file_progress <= position,
+            "Module definitions are out of order"
+        );
+
+        buffer.push_str(&src[file_progress..position]);
+        buffer.push_str(&edit);
+
+        file_progress = position;
+    }
+    buffer.push_str(&src[file_progress..]);
+
+    buffer
+}
+
+pub fn annotation_edits(
+    decls: &Declarations,
+    subs: &Subs,
+    abilities: &IAbilitiesStore<Resolved>,
+    src: &str,
+    module_id: ModuleId,
+    interns: &Interns,
+) -> Vec<(usize, String)> {
+    decls
+        .iter_bottom_up()
+        .flat_map(|(index, tag)| {
+            let var = decls.variables[index];
+            let symbol = decls.symbols[index];
+            let expr = &decls.expressions[index].value;
+
+            use roc_problem::can::RuntimeError::ExposedButNotDefined;
+            if decls.annotations[index].is_some()
+                | matches!(
+                    *expr,
+                    Expr::RuntimeError(ExposedButNotDefined(..)) | Expr::ImportParams(..)
+                )
+                | abilities.is_specialization_name(symbol.value)
+                | matches!(subs.get_content_without_compacting(var), Content::Error)
+            {
+                return None;
+            }
+
+            let byte_range = match tag {
+                DeclarationTag::Destructure(i) => {
+                    let region = decls.destructs[i.index()].loc_pattern.region;
+                    region.start().byte_offset()..region.end().byte_offset()
+                }
+                _ => symbol.byte_range(),
+            };
+
+            let edit = annotation_edit(src, subs, interns, module_id, var, byte_range);
+
+            Some(edit)
+        })
+        .collect()
+}
+
+pub fn annotation_edit(
+    src: &str,
+    subs: &Subs,
+    interns: &Interns,
+    module_id: ModuleId,
+    var: Variable,
+    symbol_range: Range<usize>,
+) -> (usize, String) {
+    let signature = roc_types::pretty_print::name_and_print_var(
+        var,
+        &mut subs.clone(),
+        module_id,
+        interns,
+        roc_types::pretty_print::DebugPrint::NOTHING,
+    )
+    // Generated names for errors start with `#`
+    .replace('#', "");
+
+    let line_start = src[..symbol_range.start]
+        .rfind('\n')
+        .map_or(symbol_range.start, |pos| pos + 1);
+    let indent = src[line_start..]
+        .split_once(|c: char| !c.is_ascii_whitespace())
+        .map_or("", |pair| pair.0);
+
+    let symbol_str = &src[symbol_range.clone()];
+    let edit = format!("{indent}{symbol_str} : {signature}\n");
+
+    (line_start, edit)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::fs::File;
+    use indoc::indoc;
+    use std::fs::{read_to_string, File};
     use std::io::Write;
     use tempfile::{tempdir, TempDir};
 
@@ -367,4 +528,75 @@ main =
 
         cleanup_temp_dir(dir);
     }
+
+    const HEADER: &str = indoc! {r#"
+        interface Test
+            exposes []
+            imports []
+
+    "#};
+
+    fn annotate_string(before: String) -> String {
+        let dir = tempdir().unwrap();
+        let file_path = setup_test_file(dir.path(), "before.roc", &before);
+
+        let arena = Bump::new();
+        let result = annotate_file(&arena, file_path.clone());
+        result.unwrap();
+
+        let annotated = read_to_string(file_path).unwrap();
+
+        cleanup_temp_dir(dir);
+        annotated
+    }
+
+    #[test]
+    fn test_annotate_simple() {
+        let before = HEADER.to_string()
+            + indoc! {r#"
+            main =
+                Task.ok {}"#};
+
+        let after = HEADER.to_string()
+            + indoc! {r#"
+            main : Task {} *
+            main =
+                Task.ok {}
+            "#};
+
+        let annotated = annotate_string(before);
+
+        assert_eq!(annotated, after);
+    }
+
+    #[test]
+    fn test_annotate_empty() {
+        let before = HEADER.to_string();
+        let after = HEADER.to_string() + "\n";
+        let annotated = annotate_string(before);
+
+        assert_eq!(annotated, after);
+    }
+
+    #[test]
+    fn test_annotate_destructure() {
+        let before = HEADER.to_string()
+            + indoc! {r#"
+            {a, b} = {a: Task.ok {}, b: (1, 2)}
+
+            main = a"#};
+
+        let after = HEADER.to_string()
+            + indoc! {r#"
+            {a, b} : { a : Task {} *, b : ( Num *, Num * )* }
+            {a, b} = {a: Task.ok {}, b: (1, 2)}
+
+            main : Task {} *
+            main = a
+            "#};
+
+        let annotated = annotate_string(before);
+
+        assert_eq!(annotated, after);
+    }
 }

crates/cli/src/lib.rs

@@ -42,7 +42,9 @@ use strum::IntoEnumIterator;
 use tempfile::TempDir;
 
 mod format;
-pub use format::{format_files, format_src, FormatMode};
+pub use format::{
+    annotate_file, annotation_edit, annotation_edits, format_files, format_src, FormatMode,
+};
 
 pub const CMD_BUILD: &str = "build";
 pub const CMD_RUN: &str = "run";
@@ -52,6 +54,7 @@ pub const CMD_DOCS: &str = "docs";
 pub const CMD_CHECK: &str = "check";
 pub const CMD_VERSION: &str = "version";
 pub const CMD_FORMAT: &str = "format";
+pub const CMD_FORMAT_ANNOTATE: &str = "annotate";
 pub const CMD_TEST: &str = "test";
 pub const CMD_GLUE: &str = "glue";
 pub const CMD_PREPROCESS_HOST: &str = "preprocess-host";
@@ -366,6 +369,16 @@ pub fn build_app() -> Command {
                     .required(false),
             )
             .after_help("If DIRECTORY_OR_FILES is omitted, the .roc files in the current working\ndirectory are formatted.")
+            .subcommand(Command::new(CMD_FORMAT_ANNOTATE)
+                .about("Annotate all top level definitions from a .roc file")
+                .arg(
+                    Arg::new(ROC_FILE)
+                        .help("The .roc file ot annotate")
+                        .value_parser(value_parser!(PathBuf))
+                        .required(false)
+                        .default_value(DEFAULT_ROC_FILENAME),
+                )
+            )
         )
         .subcommand(Command::new(CMD_VERSION)
             .about(concatcp!("Print the Roc compiler’s version, which is currently ", VERSION)))

crates/cli/src/main.rs

@@ -3,15 +3,15 @@ use bumpalo::Bump;
 use roc_build::link::LinkType;
 use roc_build::program::{check_file, CodeGenBackend};
 use roc_cli::{
-    build_app, format_files, format_src, test, BuildConfig, FormatMode, CMD_BUILD, CMD_CHECK,
-    CMD_DEV, CMD_DOCS, CMD_FORMAT, CMD_GLUE, CMD_PREPROCESS_HOST, CMD_REPL, CMD_RUN, CMD_TEST,
-    CMD_VERSION, DIRECTORY_OR_FILES, FLAG_CHECK, FLAG_DEV, FLAG_LIB, FLAG_MAIN, FLAG_MIGRATE,
-    FLAG_NO_COLOR, FLAG_NO_HEADER, FLAG_NO_LINK, FLAG_OUTPUT, FLAG_PP_DYLIB, FLAG_PP_HOST,
-    FLAG_PP_PLATFORM, FLAG_STDIN, FLAG_STDOUT, FLAG_TARGET, FLAG_TIME, GLUE_DIR, GLUE_SPEC,
-    ROC_FILE, VERSION,
+    annotate_file, build_app, format_files, format_src, test, BuildConfig, FormatMode, CMD_BUILD,
+    CMD_CHECK, CMD_DEV, CMD_DOCS, CMD_FORMAT, CMD_FORMAT_ANNOTATE, CMD_GLUE, CMD_PREPROCESS_HOST,
+    CMD_REPL, CMD_RUN, CMD_TEST, CMD_VERSION, DIRECTORY_OR_FILES, FLAG_CHECK, FLAG_DEV, FLAG_LIB,
+    FLAG_MAIN, FLAG_MIGRATE, FLAG_NO_COLOR, FLAG_NO_HEADER, FLAG_NO_LINK, FLAG_OUTPUT,
+    FLAG_PP_DYLIB, FLAG_PP_HOST, FLAG_PP_PLATFORM, FLAG_STDIN, FLAG_STDOUT, FLAG_TARGET, FLAG_TIME,
+    GLUE_DIR, GLUE_SPEC, ROC_FILE, VERSION,
 };
 use roc_docs::generate_docs_html;
-use roc_error_macros::user_error;
+use roc_error_macros::{internal_error, user_error};
 use roc_fmt::MigrationFlags;
 use roc_gen_dev::AssemblyBackendMode;
 use roc_gen_llvm::llvm::build::LlvmBackendMode;
@@ -311,6 +311,29 @@ fn main() -> io::Result<()> {
 
             Ok(0)
         }
+        Some((CMD_FORMAT, fmatches)) if Some(CMD_FORMAT_ANNOTATE) == fmatches.subcommand_name() => {
+            let matches = fmatches
+                .subcommand_matches(CMD_FORMAT_ANNOTATE)
+                .unwrap_or_else(|| internal_error!("No annotate subcommand present"));
+
+            let arena = Bump::new();
+            let roc_file_path = matches
+                .get_one::<PathBuf>(ROC_FILE)
+                .unwrap_or_else(|| internal_error!("No default for ROC_FILE"));
+
+            let annotate_exit_code = match annotate_file(&arena, roc_file_path.to_owned()) {
+                Ok(()) => 0,
+                Err(LoadingProblem::FormattedReport(report, ..)) => {
+                    eprintln!("{report}");
+                    1
+                }
+                Err(other) => {
+                    internal_error!("build_file failed with error:\n{other:?}");
+                }
+            };
+
+            Ok(annotate_exit_code)
+        }
         Some((CMD_FORMAT, matches)) => {
             let from_stdin = matches.get_flag(FLAG_STDIN);
             let to_stdout = matches.get_flag(FLAG_STDOUT);