Merge pull request #7130 from snobee/annotate-type-signatures

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,28 @@
 use std::ffi::OsStr;
 use std::io::Write;
+use std::ops::Range;
 use std::path::{Path, PathBuf};
 
-use bumpalo::Bump;
+use bumpalo::{collections::String as BumpString, 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_problem::can::RuntimeError;
+use roc_region::all::{LineColumn, LineInfo};
+use roc_reporting::report::{RenderTarget, DEFAULT_PALETTE};
+use roc_target::Target;
+use roc_types::subs::{Subs, Variable};
 
 #[derive(Copy, Clone, Debug)]
 pub enum FormatMode {
@@ -263,10 +274,177 @@ fn fmt_all<'a>(buf: &mut Buf<'a>, ast: &'a FullAst) {
     buf.fmt_end_of_file();
 }
 
+#[derive(Debug)]
+pub enum AnnotationProblem<'a> {
+    Loading(LoadingProblem<'a>),
+    Type(TypeProblem),
+}
+
+#[derive(Debug)]
+pub struct TypeProblem {
+    pub name: String,
+    pub position: LineColumn,
+}
+
+pub fn annotate_file(arena: &Bump, file: PathBuf) -> Result<(), AnnotationProblem> {
+    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,
+    )
+    .map_err(AnnotationProblem::Loading)?;
+
+    let buf = annotate_module(arena, &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<'a>(
+    arena: &'a Bump,
+    loaded: &mut LoadedModule,
+) -> Result<BumpString<'a>, AnnotationProblem<'a>> {
+    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,
+    )
+    .map_err(AnnotationProblem::Type)?;
+    edits.sort_by_key(|(offset, _)| *offset);
+
+    let mut buffer = BumpString::new_in(arena);
+    let mut file_progress = 0;
+
+    for (position, edit) in edits {
+        buffer.push_str(&src[file_progress..position]);
+        buffer.push_str(&edit);
+
+        file_progress = position;
+    }
+    buffer.push_str(&src[file_progress..]);
+
+    Ok(buffer)
+}
+
+pub fn annotation_edits(
+    decls: &Declarations,
+    subs: &Subs,
+    abilities: &IAbilitiesStore<Resolved>,
+    src: &str,
+    module_id: ModuleId,
+    interns: &Interns,
+) -> Result<Vec<(usize, String)>, TypeProblem> {
+    let mut edits = Vec::with_capacity(decls.len());
+
+    for (index, tag) in decls.iter_bottom_up() {
+        let var = decls.variables[index];
+        let symbol = decls.symbols[index];
+        let expr = &decls.expressions[index].value;
+
+        if decls.annotations[index].is_some()
+            | matches!(
+                *expr,
+                Expr::RuntimeError(RuntimeError::ExposedButNotDefined(..)) | Expr::ImportParams(..)
+            )
+            | abilities.is_specialization_name(symbol.value)
+            | matches!(tag, DeclarationTag::MutualRecursion { .. })
+        {
+            continue;
+        }
+
+        let byte_range = match tag {
+            DeclarationTag::Destructure(i) => decls.destructs[i.index()].loc_pattern.byte_range(),
+            _ => symbol.byte_range(),
+        };
+
+        let edit = annotation_edit(src, subs, interns, module_id, var, byte_range)?;
+
+        edits.push(edit);
+    }
+    Ok(edits)
+}
+
+pub fn annotation_edit(
+    src: &str,
+    subs: &Subs,
+    interns: &Interns,
+    module_id: ModuleId,
+    var: Variable,
+    symbol_range: Range<usize>,
+) -> Result<(usize, String), TypeProblem> {
+    let symbol_str = &src[symbol_range.clone()];
+    if subs.var_contains_error(var) {
+        let line_info = LineInfo::new(src);
+        let position = line_info.convert_offset(symbol_range.start as u32);
+        return Err(TypeProblem {
+            name: symbol_str.to_owned(),
+            position,
+        });
+    }
+
+    let signature = roc_types::pretty_print::name_and_print_var(
+        var,
+        &mut subs.clone(),
+        module_id,
+        interns,
+        roc_types::pretty_print::DebugPrint::NOTHING,
+    );
+
+    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 edit = format!("{indent}{symbol_str} : {signature}\n");
+
+    Ok((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};
 
@@ -379,4 +557,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 =
+                "Hello, World!""#};
+
+        let after = HEADER.to_string()
+            + indoc! {r#"
+            main : Str
+            main =
+                "Hello, World!"
+            "#};
+
+        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: "zero", b: (1, 2)}
+
+            main = a"#};
+
+        let after = HEADER.to_string()
+            + indoc! {r#"
+            {a, b} : { a : Str, b : ( Num *, Num * )* }
+            {a, b} = {a: "zero", b: (1, 2)}
+
+            main : Str
+            main = a
+            "#};
+
+        let annotated = annotate_string(before);
+
+        assert_eq!(annotated, after);
+    }
 }

crates/cli/src/lib.rs

@@ -42,7 +42,10 @@ 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, AnnotationProblem,
+    FormatMode,
+};
 
 pub const CMD_BUILD: &str = "build";
 pub const CMD_RUN: &str = "run";
@@ -52,6 +55,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";
@@ -380,6 +384,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)))