clusterflux-public/crates/disasmer-core/src/bundle.rs

use serde::{Deserialize, Serialize};
use syn::parse::Parser;
use syn::{Expr, Item, Lit, Meta, Token};

use crate::{Digest, EnvironmentResource, SourceTransferPolicy, TaskDefinitionId};

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SelectedInput {
    pub path: String,
    pub digest: Digest,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleIdentityInputs {
    pub wasm_code: Digest,
    pub task_abi: Digest,
    pub entrypoints: Vec<String>,
    pub default_entrypoint: String,
    pub environments: Vec<EnvironmentResource>,
    pub source_provider_manifest: Digest,
    pub source_transfer_policy: SourceTransferPolicy,
    pub selected_inputs: Vec<SelectedInput>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleMetadata {
    pub identity: Digest,
    pub wasm_code: Digest,
    pub task_metadata: BundleTaskMetadata,
    pub source_metadata: BundleSourceMetadata,
    pub debug_metadata: BundleDebugMetadata,
    pub large_input_policy: BundleLargeInputPolicy,
    pub restart_compatibility: BundleRestartCompatibility,
    pub environments: Vec<EnvironmentResource>,
    pub selected_inputs: Vec<SelectedInput>,
    pub embeds_full_container_images: bool,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleTaskMetadata {
    pub task_abi: Digest,
    pub entrypoints: Vec<String>,
    pub default_entrypoint: String,
    pub boundary: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleSourceMetadata {
    pub source_provider_manifest: Digest,
    pub transfer_policy: SourceTransferPolicy,
    pub selected_inputs: Vec<SelectedInput>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleDebugMetadata {
    pub available: bool,
    pub source_level_breakpoints: bool,
    pub dap_virtual_process: bool,
    pub variables_pane_supported: bool,
    pub probes: Vec<BundleDebugProbe>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleDebugProbe {
    pub id: String,
    pub source_path: String,
    pub line_start: u32,
    pub line_end: u32,
    pub function: String,
    pub task: TaskDefinitionId,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleLargeInputPolicy {
    pub selected_inputs_are_content_digests: bool,
    pub selected_input_bytes_included: bool,
    pub full_repository_bytes_included: bool,
    pub silent_task_argument_serialization: bool,
    pub supported_handle_types: Vec<String>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleRestartCompatibility {
    pub source_edits_can_restart_from_clean_task_boundary: bool,
    pub requires_clean_checkpoint_boundary: bool,
    pub compares_task_abi: Digest,
    pub compares_environment_digests: bool,
    pub compares_serialized_args: bool,
    pub discards_unflushed_task_local_changes: bool,
    pub incompatible_changes_require_whole_process_restart: bool,
}

impl BundleIdentityInputs {
    pub fn identity(&self) -> Digest {
        let mut parts = vec![
            b"bundle:v1".to_vec(),
            self.wasm_code.as_str().as_bytes().to_vec(),
            self.task_abi.as_str().as_bytes().to_vec(),
            self.source_provider_manifest.as_str().as_bytes().to_vec(),
            self.default_entrypoint.as_bytes().to_vec(),
            format!("{:?}", self.source_transfer_policy).into_bytes(),
        ];

        let mut entrypoints = self.entrypoints.clone();
        entrypoints.sort();
        for entrypoint in entrypoints {
            parts.push(entrypoint.into_bytes());
        }

        let mut environments = self.environments.clone();
        environments.sort_by(|left, right| left.name.cmp(&right.name));
        for environment in environments {
            parts.push(environment.name.as_bytes().to_vec());
            parts.push(format!("{:?}", environment.kind).into_bytes());
            parts.push(environment.digest.as_str().as_bytes().to_vec());
        }

        let mut inputs = self.selected_inputs.clone();
        inputs.sort_by(|left, right| left.path.cmp(&right.path));
        for input in inputs {
            parts.push(input.path.into_bytes());
            parts.push(input.digest.as_str().as_bytes().to_vec());
        }

        Digest::from_parts(parts)
    }

    pub fn inspectable_metadata(&self) -> BundleMetadata {
        let mut entrypoints = self.entrypoints.clone();
        entrypoints.sort();

        BundleMetadata {
            identity: self.identity(),
            wasm_code: self.wasm_code.clone(),
            task_metadata: BundleTaskMetadata {
                task_abi: self.task_abi.clone(),
                entrypoints,
                default_entrypoint: self.default_entrypoint.clone(),
                boundary: "disasmer_task_exports".to_owned(),
            },
            source_metadata: BundleSourceMetadata {
                source_provider_manifest: self.source_provider_manifest.clone(),
                transfer_policy: self.source_transfer_policy.clone(),
                selected_inputs: self.selected_inputs.clone(),
            },
            debug_metadata: BundleDebugMetadata {
                available: true,
                source_level_breakpoints: true,
                dap_virtual_process: true,
                variables_pane_supported: true,
                probes: Vec::new(),
            },
            large_input_policy: BundleLargeInputPolicy {
                selected_inputs_are_content_digests: true,
                selected_input_bytes_included: false,
                full_repository_bytes_included: false,
                silent_task_argument_serialization: false,
                supported_handle_types: vec![
                    "SourceSnapshot".to_owned(),
                    "Blob".to_owned(),
                    "Artifact".to_owned(),
                    "VFS".to_owned(),
                ],
            },
            restart_compatibility: BundleRestartCompatibility {
                source_edits_can_restart_from_clean_task_boundary: true,
                requires_clean_checkpoint_boundary: true,
                compares_task_abi: self.task_abi.clone(),
                compares_environment_digests: true,
                compares_serialized_args: true,
                discards_unflushed_task_local_changes: true,
                incompatible_changes_require_whole_process_restart: true,
            },
            environments: self.environments.clone(),
            selected_inputs: self.selected_inputs.clone(),
            embeds_full_container_images: false,
        }
    }
}

pub fn discover_source_debug_probes(
    source_path: impl Into<String>,
    source: &str,
) -> Vec<BundleDebugProbe> {
    let source_path = source_path.into();
    let lines = source.lines().collect::<Vec<_>>();
    let function_starts = lines
        .iter()
        .enumerate()
        .filter_map(|(index, line)| {
            parse_rust_function_name(line).map(|function| (index, function))
        })
        .collect::<Vec<_>>();
    let Ok(file) = syn::parse_file(source) else {
        return Vec::new();
    };

    file.items
        .iter()
        .filter_map(|item| {
            let Item::Fn(function) = item else {
                return None;
            };
            let function_name = function.sig.ident.to_string();
            let task = disasmer_probe_task(function)?;
            let (function_index, (line_index, _)) = function_starts
                .iter()
                .enumerate()
                .find(|(_, (_, candidate))| candidate == &function_name)?;
            let line_start = (*line_index + 1) as u32;
            let next_start = function_starts
                .get(function_index + 1)
                .map(|(next_index, _)| *next_index)
                .unwrap_or(lines.len());
            let line_end = next_start.max(*line_index + 1) as u32;
            Some(debug_probe(
                &source_path,
                line_start,
                line_end,
                &function_name,
                task,
            ))
        })
        .collect()
}

fn debug_probe(
    source_path: &str,
    line_start: u32,
    line_end: u32,
    function: &str,
    task: TaskDefinitionId,
) -> BundleDebugProbe {
    let id = Digest::from_parts([
        b"bundle-debug-probe:v1".as_slice(),
        source_path.as_bytes(),
        function.as_bytes(),
        task.as_str().as_bytes(),
        line_start.to_string().as_bytes(),
        line_end.to_string().as_bytes(),
    ])
    .as_str()
    .to_owned();
    BundleDebugProbe {
        id,
        source_path: source_path.to_owned(),
        line_start,
        line_end,
        function: function.to_owned(),
        task,
    }
}

fn parse_rust_function_name(line: &str) -> Option<String> {
    let start = line.find("fn ")? + 3;
    let rest = &line[start..];
    let name = rest
        .chars()
        .take_while(|ch| ch.is_ascii_alphanumeric() || *ch == '_')
        .collect::<String>();
    (!name.is_empty()).then_some(name)
}

fn disasmer_probe_task(function: &syn::ItemFn) -> Option<TaskDefinitionId> {
    for attribute in &function.attrs {
        let mut segments = attribute.path().segments.iter();
        let Some(namespace) = segments.next() else {
            continue;
        };
        let Some(kind) = segments.next() else {
            continue;
        };
        if namespace.ident != "disasmer" || segments.next().is_some() {
            continue;
        }
        let function_name = function.sig.ident.to_string();
        let default_name = match kind.ident.to_string().as_str() {
            "main" => function_name
                .strip_suffix("_main")
                .unwrap_or(&function_name)
                .to_owned(),
            "task" => function_name,
            _ => continue,
        };
        let declared_name = match &attribute.meta {
            Meta::List(list) => {
                let parser = syn::punctuated::Punctuated::<Meta, Token![,]>::parse_terminated;
                parser
                    .parse2(list.tokens.clone())
                    .ok()
                    .and_then(|items| {
                        items.into_iter().find_map(|item| {
                            let Meta::NameValue(name_value) = item else {
                                return None;
                            };
                            if !name_value.path.is_ident("name") {
                                return None;
                            }
                            let Expr::Lit(value) = name_value.value else {
                                return None;
                            };
                            let Lit::Str(value) = value.lit else {
                                return None;
                            };
                            Some(value.value())
                        })
                    })
                    .unwrap_or(default_name)
            }
            _ => default_name,
        };
        return Some(TaskDefinitionId::new(declared_name));
    }
    None
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use crate::{EnvironmentKind, EnvironmentRequirements};

    use super::*;

    fn env(digest: &str) -> EnvironmentResource {
        EnvironmentResource {
            name: "linux".to_owned(),
            kind: EnvironmentKind::Containerfile,
            recipe_path: PathBuf::from("envs/linux/Containerfile"),
            context_path: PathBuf::from("envs/linux"),
            digest: Digest::sha256(digest),
            requirements: EnvironmentRequirements::linux_container(),
        }
    }

    #[test]
    fn bundle_identity_changes_when_environment_recipe_changes() {
        let base = BundleIdentityInputs {
            wasm_code: Digest::sha256("wasm"),
            task_abi: Digest::sha256("abi"),
            entrypoints: vec!["build".to_owned()],
            default_entrypoint: "build".to_owned(),
            environments: vec![env("recipe-a")],
            source_provider_manifest: Digest::sha256("source"),
            source_transfer_policy: SourceTransferPolicy::local_first_snapshot_chunks(),
            selected_inputs: vec![],
        };
        let mut changed = base.clone();
        changed.environments = vec![env("recipe-b")];

        assert_ne!(base.identity(), changed.identity());
    }

    #[test]
    fn bundle_metadata_is_inspectable_and_does_not_vendor_images_by_default() {
        let inputs = BundleIdentityInputs {
            wasm_code: Digest::sha256("wasm"),
            task_abi: Digest::sha256("abi"),
            entrypoints: vec!["build".to_owned(), "test".to_owned()],
            default_entrypoint: "build".to_owned(),
            environments: vec![env("recipe")],
            source_provider_manifest: Digest::sha256("source"),
            source_transfer_policy: SourceTransferPolicy::local_first_snapshot_chunks(),
            selected_inputs: vec![SelectedInput {
                path: "inputs/config.json".to_owned(),
                digest: Digest::sha256("config"),
            }],
        };

        let metadata = inputs.inspectable_metadata();

        assert!(metadata.wasm_code.as_str().starts_with("sha256:"));
        assert_eq!(metadata.task_metadata.default_entrypoint, "build");
        assert!(metadata
            .task_metadata
            .entrypoints
            .contains(&"test".to_owned()));
        assert!(metadata.debug_metadata.dap_virtual_process);
        assert!(
            metadata
                .source_metadata
                .transfer_policy
                .local_source_bytes_remain_node_local
        );
        assert!(
            metadata
                .large_input_policy
                .selected_inputs_are_content_digests
        );
        assert!(!metadata.large_input_policy.selected_input_bytes_included);
        assert!(!metadata.large_input_policy.full_repository_bytes_included);
        assert!(
            !metadata
                .large_input_policy
                .silent_task_argument_serialization
        );
        assert!(metadata
            .large_input_policy
            .supported_handle_types
            .contains(&"Artifact".to_owned()));
        assert!(
            metadata
                .restart_compatibility
                .source_edits_can_restart_from_clean_task_boundary
        );
        assert!(
            metadata
                .restart_compatibility
                .requires_clean_checkpoint_boundary
        );
        assert_eq!(
            metadata.restart_compatibility.compares_task_abi,
            inputs.task_abi
        );
        assert!(
            metadata
                .restart_compatibility
                .incompatible_changes_require_whole_process_restart
        );
        assert_eq!(metadata.environments.len(), 1);
        assert!(!metadata.embeds_full_container_images);
    }

    #[test]
    fn source_debug_probe_metadata_maps_function_ranges_to_tasks() {
        let probes = discover_source_debug_probes(
            "src/build.rs",
            r#"#[disasmer::main]
fn build_main() {
    let linux = compile_linux();
}

#[disasmer::task]
fn compile_linux() {
    println!("linux");
}

fn helper_without_runtime_probe() {}

#[disasmer::task(name = "release")]
fn package_release() {
    println!("package");
}
"#,
        );

        assert_eq!(probes.len(), 3);
        assert_eq!(probes[0].source_path, "src/build.rs");
        assert_eq!(probes[0].function, "build_main");
        assert_eq!(probes[0].task, TaskDefinitionId::from("build"));
        assert_eq!(probes[0].line_start, 2);
        assert_eq!(probes[0].line_end, 6);
        assert_eq!(probes[1].function, "compile_linux");
        assert_eq!(probes[1].task, TaskDefinitionId::from("compile_linux"));
        assert_eq!(probes[2].function, "package_release");
        assert_eq!(probes[2].task, TaskDefinitionId::from("release"));
        assert!(probes.iter().all(|probe| probe.id.starts_with("sha256:")));
    }
}