lemur/src/profiler.rs

use std::{
    collections::HashMap,
    path::PathBuf,
    sync::{Arc, Mutex},
    thread,
};

use anyhow::{Result, bail};
use tokio::{select, sync::mpsc};
use wholesym::{SymbolManager, SymbolMap};

use crate::emulator::{EmulatorClient, EmulatorCommand, ProfileEvent, SimEvent, SimId};

pub struct Profiler {
    sim_id: SimId,
    client: EmulatorClient,
    status: Arc<Mutex<ProfilerStatus>>,
    action: Option<mpsc::UnboundedSender<RecordingAction>>,
    killer: Option<oneshot::Sender<()>>,
}

impl Profiler {
    pub fn new(sim_id: SimId, client: EmulatorClient) -> Self {
        Self {
            sim_id,
            client,
            status: Arc::new(Mutex::new(ProfilerStatus::Disabled)),
            action: None,
            killer: None,
        }
    }

    pub fn status(&self) -> ProfilerStatus {
        self.status.lock().unwrap().clone()
    }

    pub fn enable(&mut self) {
        let sim_id = self.sim_id;
        let client = self.client.clone();
        let status = self.status.clone();
        let (action_tx, action_rx) = mpsc::unbounded_channel();
        self.action = Some(action_tx);
        let (killer_tx, killer_rx) = oneshot::channel();
        self.killer = Some(killer_tx);
        thread::spawn(move || {
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap()
                .block_on(async move {
                    select! {
                        _ = run_profile(sim_id, client, status.clone(), action_rx) => {}
                        _ = killer_rx => {
                            *status.lock().unwrap() = ProfilerStatus::Disabled;
                        }
                    }
                })
        });
    }

    pub fn disable(&mut self) {
        if let Some(killer) = self.killer.take() {
            let _ = killer.send(());
        }
    }

    pub fn start_recording(&mut self) {
        if let Some(action) = &self.action {
            let _ = action.send(RecordingAction::Start);
        }
    }

    pub fn finish_recording(&mut self) -> oneshot::Receiver<Vec<u8>> {
        let (tx, rx) = oneshot::channel();
        if let Some(action) = &self.action {
            let _ = action.send(RecordingAction::Finish(tx));
        }
        rx
    }

    pub fn cancel_recording(&mut self) {
        if let Some(action) = &self.action {
            let _ = action.send(RecordingAction::Cancel);
        }
    }
}

impl Drop for Profiler {
    fn drop(&mut self) {
        self.disable();
    }
}

async fn run_profile(
    sim_id: SimId,
    client: EmulatorClient,
    status: Arc<Mutex<ProfilerStatus>>,
    mut action_source: mpsc::UnboundedReceiver<RecordingAction>,
) {
    let (profile_sync, mut profile_source) = mpsc::unbounded_channel();
    client.send_command(EmulatorCommand::StartProfiling(sim_id, profile_sync));

    *status.lock().unwrap() = ProfilerStatus::Enabled;

    let mut session = ProfilerSession::new();
    loop {
        select! {
            maybe_event = profile_source.recv() => {
                let Some(event) = maybe_event else {
                    break; // emulator thread disconnected
                };
                if let Err(error) = handle_event(event, &mut session).await {
                    *status.lock().unwrap() = ProfilerStatus::Error(error.to_string());
                    return;
                }
            }
            maybe_action = action_source.recv() => {
                let Some(action) = maybe_action else {
                    break; // ui thread disconnected
                };
                handle_action(action, &mut session, &status);
            }
        }
    }

    *status.lock().unwrap() = ProfilerStatus::Disabled;
}

async fn handle_event(event: ProfileEvent, session: &mut ProfilerSession) -> Result<()> {
    match event {
        ProfileEvent::Start { file_path } => session.start_profiling(file_path).await,
        ProfileEvent::Update { cycles, event } => {
            session.track_elapsed_cycles(cycles)?;
            if let Some(event) = event {
                session.track_event(event)?;
            }
            Ok(())
        }
    }
}

fn handle_action(
    action: RecordingAction,
    session: &mut ProfilerSession,
    status: &Mutex<ProfilerStatus>,
) {
    match action {
        RecordingAction::Start => {
            session.start_recording();
            *status.lock().unwrap() = ProfilerStatus::Recording;
        }
        RecordingAction::Finish(rx) => {
            if let Some(bytes) = session.finish_recording() {
                let _ = rx.send(bytes);
            }
            *status.lock().unwrap() = ProfilerStatus::Enabled;
        }
        RecordingAction::Cancel => {
            session.cancel_recording();
            *status.lock().unwrap() = ProfilerStatus::Enabled;
        }
    }
}

#[derive(Clone)]
pub enum ProfilerStatus {
    Disabled,
    Enabled,
    Recording,
    Error(String),
}

impl ProfilerStatus {
    pub fn enabled(&self) -> bool {
        matches!(self, Self::Enabled | Self::Recording)
    }
}

enum RecordingAction {
    Start,
    Finish(oneshot::Sender<Vec<u8>>),
    Cancel,
}

struct Recording {}

impl Recording {
    fn new() -> Self {
        Self {}
    }
}

struct ProfilerSession {
    program: Option<ProgramState>,
    recording: Option<Recording>,
}

impl ProfilerSession {
    fn new() -> Self {
        Self {
            program: None,
            recording: None,
        }
    }

    async fn start_profiling(&mut self, file_path: PathBuf) -> Result<()> {
        self.program = Some(ProgramState::new(file_path).await?);
        self.recording = None;
        Ok(())
    }

    fn track_elapsed_cycles(&mut self, cycles: u32) -> Result<()> {
        if let Some(program) = &mut self.program {
            program.track_elapsed_cycles(cycles)?;
        }
        Ok(())
    }

    fn track_event(&mut self, event: SimEvent) -> Result<()> {
        if let Some(program) = &mut self.program {
            program.track_event(event)?;
        }
        Ok(())
    }

    fn start_recording(&mut self) {
        self.recording = Some(Recording::new());
    }

    fn finish_recording(&mut self) -> Option<Vec<u8>> {
        self.recording.take().map(|_| vec![])
    }

    fn cancel_recording(&mut self) {
        self.recording.take();
    }
}

struct ProgramState {
    symbol_map: SymbolMap,
    call_stacks: HashMap<u16, Vec<StackFrame>>,
    context_stack: Vec<u16>,
}

struct StackFrame {
    #[expect(dead_code)]
    address: u32,
    cycles: u64,
}

const RESET_CODE: u16 = 0xfff0;
impl ProgramState {
    async fn new(file_path: PathBuf) -> Result<Self> {
        let symbol_manager = SymbolManager::with_config(Default::default());
        let symbol_map = symbol_manager
            .load_symbol_map_for_binary_at_path(&file_path, None)
            .await?;
        let mut call_stacks = HashMap::new();
        call_stacks.insert(
            RESET_CODE,
            vec![StackFrame {
                address: 0xfffffff0,
                cycles: 0,
            }],
        );
        Ok(Self {
            symbol_map,
            call_stacks,
            context_stack: vec![RESET_CODE],
        })
    }

    fn track_elapsed_cycles(&mut self, cycles: u32) -> Result<()> {
        let Some(code) = self.context_stack.last() else {
            return Ok(()); // program is halted, CPU is idle
        };
        let Some(stack) = self.call_stacks.get_mut(code) else {
            bail!("missing stack {code:04x}");
        };
        for frame in stack {
            frame.cycles += cycles as u64;
        }
        Ok(())
    }

    fn track_event(&mut self, event: SimEvent) -> Result<()> {
        match event {
            SimEvent::Call(address) => {
                let Some(code) = self.context_stack.last() else {
                    bail!("How did we call anything when we're halted?");
                };
                let Some(stack) = self.call_stacks.get_mut(code) else {
                    bail!("missing stack {code:04x}");
                };
                let name = self
                    .symbol_map
                    .lookup_sync(wholesym::LookupAddress::Svma(address as u64));
                println!("depth {}: {:x?}", stack.len(), name);
                stack.push(StackFrame { address, cycles: 0 });
            }
            SimEvent::Return => {
                let Some(code) = self.context_stack.last() else {
                    bail!("how did we return when we're halted?");
                };
                let Some(stack) = self.call_stacks.get_mut(code) else {
                    bail!("missing stack {code:04x}");
                };
                if stack.pop().is_none() {
                    bail!("returned from {code:04x} but stack was empty");
                }
                if stack.is_empty() {
                    bail!("returned to oblivion");
                }
            }
            SimEvent::Halt => {
                let Some(RESET_CODE) = self.context_stack.pop() else {
                    bail!("halted when not in an interrupt");
                };
            }
            SimEvent::Interrupt(code, address) => {
                // if the CPU was halted before, wake it up now
                if self.context_stack.is_empty() {
                    self.context_stack.push(RESET_CODE);
                }

                self.context_stack.push(code);
                if self
                    .call_stacks
                    .insert(code, vec![StackFrame { address, cycles: 0 }])
                    .is_some()
                {
                    bail!("{code:04x} fired twice");
                }
            }
            SimEvent::Reti => {
                let Some(code) = self.context_stack.pop() else {
                    bail!("RETI when halted");
                };
                if code == RESET_CODE {
                    bail!("RETI when not in interrupt");
                }
                if self.call_stacks.remove(&code).is_none() {
                    bail!("{code:04x} popped but never called");
                }
            }
        }
        Ok(())
    }
}