migration.rs

//! Checkpoint and restore functionality for live migration.
//!
//! This module implements serialization and deserialization of runtime state,
//! enabling process migration, fault tolerance, and debugging capabilities.
//!
//! ## Serializable State
//!
//! The checkpoint captures:
//! - Execution stacks (call frames and register state)
//! - Linear memory contents
//! - Global variable values
//! - Table entries (function references)
//!
//! ## Trigger Mechanisms
//!
//! - **wasm32-wasip1-threads**: Background thread monitors `checkpoint.trigger` file
//! - **wasm32-wasip1**: WASI-based file existence check at instruction boundaries

use crate::error::RuntimeError;
use crate::execution::func::FuncInst;
use crate::execution::global::GlobalAddr;
use crate::execution::mem::MemAddr;
use crate::execution::module::ModuleInst;
use crate::execution::table::TableAddr;
use crate::execution::value::{Ref, Val};
use crate::execution::vm::{Frame, Stacks};
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;
use std::rc::Rc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;
use std::time::Duration;

static CHECKPOINT_TRIGGERED: AtomicBool = AtomicBool::new(false);
const CHECKPOINT_TRIGGER_FILE: &str = "./checkpoint.trigger";

/// Starts background thread to monitor checkpoint trigger file.
///
/// Used on `wasm32-wasip1-threads` target for non-blocking checkpoint detection.
pub fn setup_checkpoint_monitor() {
    thread::spawn(|| loop {
        if std::path::Path::new(CHECKPOINT_TRIGGER_FILE).exists() {
            CHECKPOINT_TRIGGERED.store(true, Ordering::Relaxed);
            let _ = std::fs::remove_file(CHECKPOINT_TRIGGER_FILE);
        }
        thread::sleep(Duration::from_millis(10));
    });
}

/// Checks if checkpoint has been triggered via atomic flag.
#[inline(always)]
pub fn check_checkpoint_flag() -> bool {
    CHECKPOINT_TRIGGERED.load(Ordering::Relaxed)
}

/// Checks for checkpoint trigger via WASI file existence.
///
/// Fallback for `wasm32-wasip1` target without thread support.
pub fn check_checkpoint_trigger(frame: &Frame) -> Result<bool, RuntimeError> {
    if let Some(module) = frame.module.upgrade() {
        if let Some(ref wasi) = module.wasi_impl {
            if wasi.check_file_exists(CHECKPOINT_TRIGGER_FILE) {
                let _ = std::fs::remove_file(CHECKPOINT_TRIGGER_FILE);
                return Ok(true);
            }
        }
    }
    Ok(false)
}

/// Complete runtime state for checkpoint serialization.
///
/// Contains all information needed to restore execution:
/// - Call stack and register state
/// - Linear memory contents
/// - Global variable values
/// - Table entries
#[derive(Serialize, Deserialize, Debug)]
pub struct SerializableState {
    pub stacks: Stacks,
    pub memory_data: Vec<u8>,
    pub global_values: Vec<Val>,
    pub tables_data: Vec<Vec<Option<u32>>>,
    pub frame_func_indices: Vec<u32>,
}

/// Serializes runtime state to a checkpoint file.
///
/// Captures memory, globals, tables, and stack state for later restoration.
pub fn checkpoint<P: AsRef<Path>>(
    module_inst: &ModuleInst,
    stacks: &Stacks,
    mem_addrs: &[MemAddr],
    global_addrs: &[GlobalAddr],
    table_addrs: &[TableAddr],
    output_path: P,
) -> Result<(), RuntimeError> {
    println!("Checkpointing state to {:?}...", output_path.as_ref());

    // 1. Gather Memory state
    let memory_data = if let Some(mem_addr) = mem_addrs.get(0) {
        mem_addr.get_data()
    } else {
        Vec::new()
    };

    // 2. Gather Global state
    let global_values = global_addrs
        .iter()
        .map(|global_addr| Ok(global_addr.get()))
        .collect::<Result<Vec<Val>, RuntimeError>>()?;

    // 3. Gather Table state (using Arc::ptr_eq)
    let tables_data = table_addrs
        .iter()
        .map(|table_addr| -> Result<Vec<Option<u32>>, RuntimeError> {
            let table_inst = table_addr.read_lock();
            let mut table_indices = Vec::with_capacity(table_inst.elem.len());
            for val in table_inst.elem.iter() {
                if let Val::Ref(Ref::FuncAddr(target_func_addr)) = val {
                    let found_index = module_inst.func_addrs.iter().position(|module_func_addr| {
                        Rc::ptr_eq(target_func_addr.get_rc(), module_func_addr.get_rc())
                    });
                    if let Some(index) = found_index {
                        table_indices.push(Some(index as u32));
                    } else {
                        eprintln!("Warning: FuncAddr in table not found in module func_addrs during checkpoint.");
                        table_indices.push(None);
                    }
                } else {
                    table_indices.push(None);
                }
            }
            Ok(table_indices)
        })
        .collect::<Result<Vec<Vec<Option<u32>>>, _>>()?;

    // 4. Compute function indices for each activation frame (using Rc::ptr_eq)
    let frame_func_indices = stacks
        .activation_frame_stack
        .iter()
        .map(|frame_stack| {
            let frame_instrs = &frame_stack.label_stack[0].processed_instrs;
            module_inst
                .func_addrs
                .iter()
                .position(|func_addr| {
                    let inst = func_addr.read_lock();
                    if let FuncInst::RuntimeFunc { code, .. } = inst {
                        Rc::ptr_eq(frame_instrs, &code.body)
                    } else {
                        false
                    }
                })
                .expect("Function not found in module func_addrs during checkpoint")
                as u32
        })
        .collect::<Vec<u32>>();

    // 5. Assemble state
    // Note: Register file is already compact because restore_offsets() truncates
    // register vectors on function return, so no checkpoint-time compaction needed.
    let state = SerializableState {
        stacks: stacks.clone(),
        memory_data,
        global_values,
        tables_data,
        frame_func_indices,
    };

    // 6. Serialize and write (with per-component size diagnostics)
    let reg_file_size = bincode::serialize(&state.stacks.reg_file)
        .map(|v| v.len())
        .unwrap_or(0);
    let frames_count = state.stacks.activation_frame_stack.len();
    let total_labels: usize = state
        .stacks
        .activation_frame_stack
        .iter()
        .map(|f| f.label_stack.len())
        .sum();
    let frames_size = bincode::serialize(&state.stacks.activation_frame_stack)
        .map(|v| v.len())
        .unwrap_or(0);
    let total_locals: usize = state
        .stacks
        .activation_frame_stack
        .iter()
        .map(|f| f.frame.locals.len())
        .sum();
    let _stacks_size = reg_file_size + frames_size;
    let memory_size = state.memory_data.len();
    let globals_size = bincode::serialize(&state.global_values)
        .map(|v| v.len())
        .unwrap_or(0);
    let tables_size = bincode::serialize(&state.tables_data)
        .map(|v| v.len())
        .unwrap_or(0);
    let indices_size = bincode::serialize(&state.frame_func_indices)
        .map(|v| v.len())
        .unwrap_or(0);
    println!("Checkpoint component sizes:");
    println!("  reg_file:           {} bytes", reg_file_size);
    println!(
        "  frames:             {} bytes ({} frames, {} labels, {} locals total)",
        frames_size, frames_count, total_labels, total_locals
    );
    println!("  memory_data:        {} bytes", memory_size);
    println!("  global_values:      {} bytes", globals_size);
    println!("  tables_data:        {} bytes", tables_size);
    println!("  frame_func_indices: {} bytes", indices_size);

    let encoded: Vec<u8> =
        bincode::serialize(&state).map_err(|e| RuntimeError::SerializationError(e.to_string()))?;

    println!("  total encoded:      {} bytes", encoded.len());

    let mut file =
        File::create(output_path).map_err(|e| RuntimeError::CheckpointSaveError(e.to_string()))?;
    file.write_all(&encoded)
        .map_err(|e| RuntimeError::CheckpointSaveError(e.to_string()))?;

    println!("Checkpoint successful.");
    Ok(())
}

/// Restores runtime state from a checkpoint file.
///
/// Reads serialized state and restores memory, globals, tables, and stacks.
pub fn restore<P: AsRef<Path>>(
    module_inst: Rc<ModuleInst>,
    input_path: P,
) -> Result<Stacks, RuntimeError> {
    println!("Restoring state from {:?}...", input_path.as_ref());

    // 1. Read from file
    let mut file =
        File::open(input_path).map_err(|e| RuntimeError::CheckpointLoadError(e.to_string()))?;
    let mut encoded = Vec::new();
    file.read_to_end(&mut encoded)
        .map_err(|e| RuntimeError::CheckpointLoadError(e.to_string()))?;

    // 2. Deserialize the state using bincode
    let mut state: SerializableState = bincode::deserialize(&encoded[..])
        .map_err(|e| RuntimeError::DeserializationError(e.to_string()))?;

    // 3. Restore memory state into module_inst
    // Assuming only one memory instance for now
    if let Some(mem_addr) = module_inst.mem_addrs.get(0) {
        mem_addr.set_data(state.memory_data);
        println!("Memory state restored into module instance.");
    } else if !state.memory_data.is_empty() {
        eprintln!("Warning: Checkpoint contains memory data, but module has no memory instance.");
    }

    // 4. Restore global state into module_inst
    if module_inst.global_addrs.len() == state.global_values.len() {
        for (global_addr, value) in module_inst.global_addrs.iter().zip(state.global_values) {
            global_addr.set(value)?;
        }
        println!("Global state restored into module instance.");
    } else {
        eprintln!(
            "Warning: Mismatch in global variable count between module ({}) and checkpoint ({}). Globals not restored.",
            module_inst.global_addrs.len(),
            state.global_values.len()
        );
    }

    // 5. Restore table state into module_inst
    if module_inst.table_addrs.len() == state.tables_data.len() {
        for (table_idx, table_indices) in state.tables_data.iter().enumerate() {
            let table_addr = &module_inst.table_addrs[table_idx];
            let restored_elements = table_indices
                .iter()
                .map(|maybe_index| {
                    maybe_index
                        .map(|index| {
                            module_inst
                                .func_addrs
                                .get(index as usize)
                                .cloned()
                                .ok_or_else(|| {
                                    RuntimeError::RestoreError(format!(
                                        "Invalid function index {} found in table data",
                                        index
                                    ))
                                })
                        })
                        .transpose()
                })
                .collect::<Result<Vec<Option<_>>, _>>()?;
            table_addr.set_elements(restored_elements)?;
        }
        println!("Table state restored into module instance.");
    } else {
        eprintln!(
            "Warning: Mismatch in table count between module ({}) and checkpoint ({}). Tables not restored.",
            module_inst.table_addrs.len(),
            state.tables_data.len()
        );
    }

    // 6. Reconstruct skipped fields in Stacks (Frame::module, primary_mem, processed_instrs)
    let primary_mem = module_inst.mem_addrs.first().cloned();
    for (frame_stack, &func_idx) in state
        .stacks
        .activation_frame_stack
        .iter_mut()
        .zip(state.frame_func_indices.iter())
    {
        frame_stack.frame.module = Rc::downgrade(&module_inst);
        frame_stack.primary_mem = primary_mem.clone();

        // Reconstruct processed_instrs from module function body
        let func_addr = &module_inst.func_addrs[func_idx as usize];
        let func_inst = func_addr.read_lock();
        if let FuncInst::RuntimeFunc { code, .. } = func_inst {
            let body = code.body.clone();
            for label_stack in frame_stack.label_stack.iter_mut() {
                label_stack.processed_instrs = body.clone();
            }
        }
    }
    println!("Frame module references and processed instructions restored.");

    println!("Restore successful (state applied to module). Returning Stacks.");
    Ok(state.stacks)
}