ruka_codegen_rust/

program_cfg.rs

use super::*;

#[derive(Debug, Clone)]
struct SourceMap {
    file: String,
    function_lines: BTreeMap<String, usize>,
}

impl SourceMap {
    #[cfg(test)]
    fn unknown() -> Self {
        Self {
            file: "<unknown>".to_owned(),
            function_lines: BTreeMap::new(),
        }
    }

    fn from_function_lines(file: String, function_lines: &BTreeMap<String, usize>) -> Self {
        Self {
            file,
            function_lines: function_lines.clone(),
        }
    }

    fn function_line(&self, name: &str) -> usize {
        self.function_lines.get(name).copied().unwrap_or(0)
    }
}

/// Error produced while generating or writing Rust source.
#[derive(Debug, Error)]
pub enum CodegenError {
    /// Token rendering failed.
    #[error("failed to format generated rust source: {0}")]
    Format(#[from] genco::fmt::Error),
    /// Generated source could not be parsed for pretty formatting.
    #[error("failed to parse generated rust source for pretty formatting: {0}")]
    PrettyFormat(#[from] syn::Error),
    /// Writing generated source to disk failed.
    #[error("failed to write generated rust source: {0}")]
    Io(#[from] std::io::Error),
}

/// Emit Rust source for `program` and write it to `path`.
pub fn emit_to_path(
    program: &MirProgram,
    path: &Path,
    source_file: &Path,
    function_lines: &BTreeMap<String, usize>,
) -> Result<(), CodegenError> {
    let source = emit_program_with_function_lines(program, source_file, function_lines)?;
    std::fs::write(path, source)?;
    Ok(())
}

/// Emit Rust source for a MIR program without source mapping metadata.
#[cfg(test)]
pub fn emit_program(program: &MirProgram) -> Result<String, CodegenError> {
    emit_program_with_map(program, &SourceMap::unknown())
}

/// Emit Rust source for a MIR program with source location metadata.
pub fn emit_program_with_function_lines(
    program: &MirProgram,
    source_file: &Path,
    function_lines: &BTreeMap<String, usize>,
) -> Result<String, CodegenError> {
    let source_map =
        SourceMap::from_function_lines(source_file.display().to_string(), function_lines);
    emit_program_with_map(program, &source_map)
}

fn emit_program_with_map(
    program: &MirProgram,
    source_map: &SourceMap,
) -> Result<String, CodegenError> {
    for (_, function) in program.functions.iter() {
        function.assert_valid();
    }

    let program_names = ProgramNames::from_program(program);

    let mut tokens: rust::Tokens = quote! {
        #[derive(Debug, Clone)]
        /// Error returned by the generated Rust runtime shim.
        pub enum RuntimeError {
            MissingMain,
            InvalidMainArity { actual: usize },
        }
    };

    for decl in &program.structs {
        let decl_tokens = emit_struct_decl_tokens(decl);
        quote_in!(tokens => $decl_tokens);
    }

    for decl in &program.enums {
        let decl_tokens = emit_enum_decl_tokens(decl);
        quote_in!(tokens => $decl_tokens);
    }

    for (func_id, function) in program.functions.iter() {
        let function_names = FunctionNames::from_function(function);
        let function_tokens = emit_function_tokens(
            func_id,
            function,
            &program_names,
            &function_names,
            &source_map.file,
            source_map.function_line(&function.name),
        );
        quote_in!(tokens => $function_tokens);
    }

    if let Some(main_id) = program.function_names.get("main") {
        let main_arity = program.functions[*main_id].arity;
        let main_ident = program_names.function_ident(*main_id);
        quote_in!(tokens =>
            /// Execute the generated `main` function when present and arity-zero.
            pub fn run_main() -> Result<(), RuntimeError> {
                if $main_arity != 0 {
                    Err(RuntimeError::InvalidMainArity { actual: $main_arity })
                } else {
                    let _ = $main_ident();
                    ruka_runtime::ptr::assert_no_leaks();
                    Ok(())
                }
            }
        );
    } else {
        quote_in!(tokens =>
            /// Report that no callable `main` function was generated.
            pub fn run_main() -> Result<(), RuntimeError> {
                Err(RuntimeError::MissingMain)
            }
        );
    }

    let source = tokens.to_file_string()?;
    let syntax = parse_file(&source)?;
    Ok(prettyplease::unparse(&syntax))
}

fn emit_struct_decl_tokens(decl: &MirStructDecl) -> rust::Tokens {
    let mut tokens = rust::Tokens::new();
    let name = mangle_struct_ident(&decl.name);
    let mut params = rust::Tokens::new();
    for (index, param) in decl.type_params.iter().enumerate() {
        if index > 0 {
            quote_in!(params => ,);
        }
        let param = param.clone();
        quote_in!(params => $param);
    }

    let mut fields = rust::Tokens::new();
    for field in &decl.fields {
        let field_name = field.name.clone();
        let field_ty = emit_type_expr_tokens(&field.ty);
        quote_in!(fields => $field_name: $field_ty,);
    }

    if decl.type_params.is_empty() {
        quote_in!(tokens =>
            #[allow(non_camel_case_types)]
            #[derive(Debug, Clone)]
            struct $name {
                $fields
            }
        );
    } else {
        quote_in!(tokens =>
            #[allow(non_camel_case_types)]
            #[derive(Debug, Clone)]
            struct $name<$params> {
                $fields
            }
        );
    }

    tokens
}

fn emit_enum_decl_tokens(decl: &MirEnumDecl) -> rust::Tokens {
    let mut tokens = rust::Tokens::new();
    let name = mangle_struct_ident(&decl.name);

    let mut params = rust::Tokens::new();
    for (index, param) in decl.type_params.iter().enumerate() {
        if index > 0 {
            quote_in!(params => ,);
        }
        let param = param.clone();
        quote_in!(params => $param);
    }

    let mut variants = rust::Tokens::new();
    for variant in &decl.variants {
        let variant_name = variant.name.clone();
        if variant.payload.is_empty() {
            quote_in!(variants => $variant_name,);
        } else {
            let mut payload_tokens = rust::Tokens::new();
            for (index, payload) in variant.payload.iter().enumerate() {
                if index > 0 {
                    quote_in!(payload_tokens => ,);
                }
                let payload = emit_type_expr_tokens(payload);
                quote_in!(payload_tokens => $payload);
            }
            quote_in!(variants => $variant_name($payload_tokens),);
        }
    }

    if decl.type_params.is_empty() {
        quote_in!(tokens =>
            #[allow(non_camel_case_types)]
            #[allow(dead_code)]
            #[derive(Debug, Clone)]
            enum $name {
                $variants
            }
        );
    } else {
        quote_in!(tokens =>
            #[allow(non_camel_case_types)]
            #[allow(dead_code)]
            #[derive(Debug, Clone)]
            enum $name<$params> {
                $variants
            }
        );
    }

    tokens
}

fn emit_function_tokens(
    func_id: MirFuncId,
    function: &MirFunction,
    program_names: &ProgramNames,
    function_names: &FunctionNames,
    source_file: &str,
    source_line: usize,
) -> rust::Tokens {
    let mut tokens = rust::Tokens::new();
    let source_doc = format!("source: {}:{}", source_file, source_line);
    let func_ident = program_names.function_ident(func_id);
    let return_ty = emit_ty_tokens(&function.return_ty);

    let structured_body = structurize_cfg_body(function);

    let mut mut_locals = HashSet::new();
    collect_mutable_locals(&structured_body, &mut mut_locals);

    let mut read_locals = HashSet::new();
    collect_read_locals(&structured_body, &mut read_locals);

    let (ref_ro_locals, ref_mut_locals) = collect_ref_locals(function);
    let slice_locals = collect_slice_locals(function);

    let mut param_inits = rust::Tokens::new();
    for binding in function.param_bindings() {
        if binding.requires_materialization() {
            assert!(
                binding.expects_view(),
                "only view params may require Rust materialization"
            );
            assert!(
                !binding.local.is_place(),
                "materialized Rust params should lower to value locals"
            );
        }
        if !read_locals.contains(&binding.local_id) {
            continue;
        }
        let local_name = function_names.local_ident(binding.local_id);
        let arg_name = incoming_param_ident(function_names, binding.local_id, binding.index, true);
        if binding.expects_view() || binding.expects_mut_borrow() {
            if binding.materializes_view_from_owned() {
                quote_in!(param_inits => let $local_name = (*$arg_name).clone(););
            } else {
                quote_in!(param_inits => let $local_name = $arg_name;);
            }
        } else {
            let mut_kw = if mut_locals.contains(&binding.local_id) {
                quote!(mut)
            } else {
                quote!()
            };
            quote_in!(param_inits => let $mut_kw $local_name = $arg_name;);
        }
    }

    let body_tokens = emit_stmt_list_tokens(
        function,
        &structured_body,
        source_file,
        source_line,
        &mut_locals,
        &read_locals,
        &ref_ro_locals,
        &ref_mut_locals,
        &slice_locals,
        function_names,
        program_names,
    );

    let params = emit_function_params(function, function_names, &read_locals);

    quote_in!(tokens =>
        #[doc = $(quoted(source_doc.as_str()))]
        #[allow(non_snake_case)]
        #[allow(unused_assignments)]
        fn $func_ident($params) -> $return_ty {
            $param_inits
            $body_tokens
        }
    );

    tokens
}

#[derive(Debug, Clone)]
enum CfgEnd {
    Jump {
        target: MirBlockId,
        args: Vec<MirLocalId>,
    },
    Return,
}

fn structurize_cfg_body(function: &MirFunction) -> Vec<MirStmt> {
    let mut active_loops = Vec::new();
    let (stmts, end) = structurize_from_block(function, function.entry, &mut active_loops);
    if !matches!(end, CfgEnd::Return) {
        panic!("cfg structurizer expected function to end in return");
    }
    stmts
}

fn structurize_from_block(
    function: &MirFunction,
    block_id: MirBlockId,
    active_loops: &mut Vec<MirBlockId>,
) -> (Vec<MirStmt>, CfgEnd) {
    let block = &function.blocks[block_id];
    let mut out = Vec::new();
    for instr in &block.instrs {
        out.push(MirStmt::Instr(instr.clone()));
    }

    match &block.terminator {
        MirTerminator::Return { value } => {
            out.push(MirStmt::Return { value: *value });
            (out, CfgEnd::Return)
        }
        MirTerminator::Jump { target, args } => {
            if active_loops.contains(target) {
                return (
                    out,
                    CfgEnd::Jump {
                        target: *target,
                        args: args.clone(),
                    },
                );
            }

            if let Some((while_stmt, exit)) =
                structurize_while_from_jump(function, *target, args, active_loops)
            {
                out.push(while_stmt);
                let (rest, end) = structurize_from_block(function, exit, active_loops);
                out.extend(rest);
                (out, end)
            } else {
                (
                    out,
                    CfgEnd::Jump {
                        target: *target,
                        args: args.clone(),
                    },
                )
            }
        }
        MirTerminator::Branch {
            cond,
            then_target,
            else_target,
            ..
        } => {
            let (then_body, then_end) =
                structurize_from_block(function, *then_target, active_loops);
            let (else_body, else_end) =
                structurize_from_block(function, *else_target, active_loops);

            out.push(MirStmt::If {
                cond: *cond,
                then_body,
                else_body,
            });

            match (then_end, else_end) {
                (CfgEnd::Return, CfgEnd::Return) => (out, CfgEnd::Return),
                (
                    CfgEnd::Jump {
                        target: then_merge,
                        args: then_args,
                    },
                    CfgEnd::Jump {
                        target: else_merge,
                        args: else_args,
                    },
                ) if then_merge == else_merge && then_args.is_empty() && else_args.is_empty() => {
                    let (rest, end) = structurize_from_block(function, then_merge, active_loops);
                    out.extend(rest);
                    (out, end)
                }
                (CfgEnd::Return, CfgEnd::Jump { target, args }) if args.is_empty() => {
                    let (rest, end) = structurize_from_block(function, target, active_loops);
                    out.extend(rest);
                    (out, end)
                }
                (CfgEnd::Jump { target, args }, CfgEnd::Return) if args.is_empty() => {
                    let (rest, end) = structurize_from_block(function, target, active_loops);
                    out.extend(rest);
                    (out, end)
                }
                _ => panic!("unsupported cfg branch shape for structurization"),
            }
        }
    }
}

fn structurize_while_from_jump(
    function: &MirFunction,
    header: MirBlockId,
    init_args: &[MirLocalId],
    active_loops: &mut Vec<MirBlockId>,
) -> Option<(MirStmt, MirBlockId)> {
    let header_block = &function.blocks[header];
    let (cond, body_start, exit) = match &header_block.terminator {
        MirTerminator::Branch {
            cond,
            then_target,
            then_args,
            else_target,
            else_args,
        } if then_args.is_empty() && else_args.is_empty() => (*cond, *then_target, *else_target),
        _ => return None,
    };

    active_loops.push(header);
    let (body, body_end) = structurize_from_block(function, body_start, active_loops);
    let popped = active_loops.pop();
    assert_eq!(popped, Some(header));
    let step_args = match body_end {
        CfgEnd::Jump { target, args } if target == header => args,
        CfgEnd::Return => Vec::new(),
        _ => return None,
    };

    let cond_body = header_block
        .instrs
        .iter()
        .cloned()
        .map(MirStmt::Instr)
        .collect::<Vec<_>>();

    Some((
        MirStmt::While {
            loop_params: header_block.params.clone(),
            init_args: init_args.to_vec(),
            cond_body,
            cond,
            body,
            step_args,
        },
        exit,
    ))
}

pub(crate) fn emit_stmt_list_tokens(
    function: &MirFunction,
    stmts: &[MirStmt],
    source_file: &str,
    source_line: usize,
    mut_locals: &HashSet<MirLocalId>,
    read_locals: &HashSet<MirLocalId>,
    ref_ro_locals: &HashSet<MirLocalId>,
    ref_mut_locals: &HashSet<MirLocalId>,
    slice_locals: &HashSet<MirLocalId>,
    function_names: &FunctionNames,
    program_names: &ProgramNames,
) -> rust::Tokens {
    let mut tokens = rust::Tokens::new();
    let mut index = 0;
    while index < stmts.len() {
        if let Some((fused_tokens, consumed)) = try_emit_pointer_if_binding_tokens(
            function,
            stmts,
            index,
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        ) {
            quote_in!(tokens => $fused_tokens);
            index += consumed;
            continue;
        }

        let stmt_tokens = emit_stmt_tokens(
            function,
            &stmts[index],
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        );
        quote_in!(tokens => $stmt_tokens);
        index += 1;
    }
    tokens
}

fn try_emit_pointer_if_binding_tokens(
    function: &MirFunction,
    stmts: &[MirStmt],
    index: usize,
    source_file: &str,
    source_line: usize,
    mut_locals: &HashSet<MirLocalId>,
    read_locals: &HashSet<MirLocalId>,
    ref_ro_locals: &HashSet<MirLocalId>,
    ref_mut_locals: &HashSet<MirLocalId>,
    slice_locals: &HashSet<MirLocalId>,
    function_names: &FunctionNames,
    program_names: &ProgramNames,
) -> Option<(rust::Tokens, usize)> {
    let MirStmt::Instr(MirInstr::PointerIsSome { pointer, dst }) = stmts.get(index)? else {
        return None;
    };
    let MirStmt::If {
        cond,
        then_body,
        else_body,
    } = stmts.get(index + 1)?
    else {
        return None;
    };
    if cond != dst {
        return None;
    }
    if then_body.len() < 2 {
        return None;
    }

    let mut tokens = rust::Tokens::new();
    let pointer_name = function_names.local_ident(*pointer);
    let then_rest;

    if let (
        MirStmt::Instr(MirInstr::PointerBorrowRo {
            pointer: borrow_pointer,
            dst: borrowed_ref,
        }),
        MirStmt::Instr(MirInstr::DerefCopy {
            src: deref_src,
            dst: binding_local,
        }),
    ) = (&then_body[0], &then_body[1])
    {
        if borrow_pointer != pointer || deref_src != borrowed_ref {
            return None;
        }
        let borrowed_ref_name = function_names.local_ident(*borrowed_ref);
        let binding_name = function_names.local_ident(*binding_local);
        let mut_kw = if mut_locals.contains(binding_local) {
            quote!(mut)
        } else {
            quote!()
        };
        let binding_init = if read_locals.contains(binding_local) {
            quote!(let $mut_kw $binding_name = (*$borrowed_ref_name).clone();)
        } else {
            quote!(let _ = (*$borrowed_ref_name).clone();)
        };

        then_rest = &then_body[2..];
        let then_tokens = emit_stmt_list_tokens(
            function,
            then_rest,
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        );
        let else_tokens = emit_stmt_list_tokens(
            function,
            else_body,
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        );
        quote_in!(tokens =>
            if let Some($borrowed_ref_name) = $pointer_name.as_ref() {
                let $borrowed_ref_name = $borrowed_ref_name.borrow();
                $binding_init
                $then_tokens
            } else {
                $else_tokens
            }
        );
        return Some((tokens, 2));
    }

    if let (
        MirStmt::Instr(MirInstr::PointerBorrowRo {
            pointer: borrow_pointer,
            dst: borrowed_ref,
        }),
        MirStmt::Instr(MirInstr::Move {
            src: move_src,
            dst: binding_local,
        }),
    ) = (&then_body[0], &then_body[1])
    {
        if borrow_pointer != pointer || move_src != borrowed_ref {
            return None;
        }
        let borrowed_ref_name = function_names.local_ident(*borrowed_ref);
        let binding_name = if read_locals.contains(binding_local) {
            function_names.local_ident(*binding_local).to_owned()
        } else {
            "_".to_owned()
        };
        then_rest = &then_body[2..];
        let then_tokens = emit_stmt_list_tokens(
            function,
            then_rest,
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        );
        let else_tokens = emit_stmt_list_tokens(
            function,
            else_body,
            source_file,
            source_line,
            mut_locals,
            read_locals,
            ref_ro_locals,
            ref_mut_locals,
            slice_locals,
            function_names,
            program_names,
        );
        quote_in!(tokens =>
            if let Some($borrowed_ref_name) = $pointer_name.as_ref() {
                let $binding_name = $borrowed_ref_name.borrow();
                $then_tokens
            } else {
                $else_tokens
            }
        );
        return Some((tokens, 2));
    }

    if then_body.len() >= 2 {
        if let (
            MirStmt::Instr(MirInstr::PointerBorrowMut {
                pointer: borrow_pointer,
                dst: borrowed_ref,
            }),
            MirStmt::Instr(MirInstr::Move {
                src: move_src,
                dst: binding_local,
            }),
        ) = (&then_body[0], &then_body[1])
        {
            if borrow_pointer != pointer || move_src != borrowed_ref {
                return None;
            }
            let binding_name = if read_locals.contains(binding_local) {
                function_names.local_ident(*binding_local).to_owned()
            } else {
                "_".to_owned()
            };
            let binding_ptr_name = synthetic_temp_ident(&format!("{}_ptr", binding_name));
            let binding_ptr_pattern = binding_ptr_name.clone();
            then_rest = &then_body[2..];
            let then_tokens = emit_stmt_list_tokens(
                function,
                then_rest,
                source_file,
                source_line,
                mut_locals,
                read_locals,
                ref_ro_locals,
                ref_mut_locals,
                slice_locals,
                function_names,
                program_names,
            );
            let else_tokens = emit_stmt_list_tokens(
                function,
                else_body,
                source_file,
                source_line,
                mut_locals,
                read_locals,
                ref_ro_locals,
                ref_mut_locals,
                slice_locals,
                function_names,
                program_names,
            );
            quote_in!(tokens =>
                if let Some($binding_ptr_pattern) = $pointer_name.as_mut() {
                    let $binding_name = $binding_ptr_name.borrow_mut();
                    $then_tokens
                } else {
                    $else_tokens
                }
            );
            return Some((tokens, 2));
        }
    }

    None
}