lemur/src/profiler.rs

use std::{
    collections::HashMap,
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    thread,
};

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,
    running: Arc<AtomicBool>,
    killer: Option<oneshot::Sender<()>>,
}

impl Profiler {
    pub fn new(sim_id: SimId, client: EmulatorClient) -> Self {
        Self {
            sim_id,
            client,
            running: Arc::new(AtomicBool::new(false)),
            killer: None,
        }
    }

    pub fn started(&self) -> bool {
        self.running.load(Ordering::Relaxed)
    }

    pub fn start(&mut self) {
        let sim_id = self.sim_id;
        let client = self.client.clone();
        let running = self.running.clone();
        let (tx, rx) = oneshot::channel();
        self.killer = Some(tx);
        thread::spawn(move || {
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap()
                .block_on(async move {
                    select! {
                        _ = run_profile(sim_id, client, running.clone()) => {}
                        _ = rx => {
                            running.store(false, Ordering::Relaxed);
                        }
                    }
                })
        });
    }

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

async fn run_profile(sim_id: SimId, client: EmulatorClient, running: Arc<AtomicBool>) {
    let (profile_sync, mut profile_source) = mpsc::unbounded_channel();
    client.send_command(EmulatorCommand::StartProfiling(sim_id, profile_sync));

    running.store(true, Ordering::Relaxed);

    let mut session = None;
    while let Some(event) = profile_source.recv().await {
        match event {
            ProfileEvent::Start { file_path } => {
                session = Some(ProfileSession::new(file_path).await);
            }
            ProfileEvent::Update { cycles, event } => {
                if let Some(session) = &mut session {
                    session.track_elapsed_cycles(cycles);
                    if let Some(event) = event {
                        session.track_event(event);
                    }
                }
            }
        }
    }
    running.store(false, Ordering::Release);
}

struct ProfileSession {
    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 ProfileSession {
    async fn new(file_path: PathBuf) -> 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
            .expect("cannae load symbols");
        let mut call_stacks = HashMap::new();
        call_stacks.insert(
            RESET_CODE,
            vec![StackFrame {
                address: 0xfffffff0,
                cycles: 0,
            }],
        );
        Self {
            symbol_map,
            call_stacks,
            context_stack: vec![RESET_CODE],
        }
    }

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

    fn track_event(&mut self, event: SimEvent) {
        match event {
            SimEvent::Call(address) => {
                let Some(code) = self.context_stack.last() else {
                    panic!("How did we call anything when we're halted?");
                };
                let Some(stack) = self.call_stacks.get_mut(code) else {
                    panic!("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 {
                    panic!("how did we return when we're halted?");
                };
                let Some(stack) = self.call_stacks.get_mut(code) else {
                    panic!("missing stack {code:04x}");
                };
                if stack.pop().is_none() {
                    panic!("returned from {code:04x} but stack was empty");
                }
                if stack.is_empty() {
                    panic!("returned to oblivion");
                }
            }
            SimEvent::Halt => {
                let Some(RESET_CODE) = self.context_stack.pop() else {
                    panic!("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()
                {
                    panic!("{code:04x} fired twice");
                }
            }
            SimEvent::Reti => {
                let Some(code) = self.context_stack.pop() else {
                    panic!("RETI when halted");
                };
                if code == RESET_CODE {
                    panic!("RETI when not in interrupt");
                }
                if self.call_stacks.remove(&code).is_none() {
                    panic!("{code:04x} popped but never called");
                }
            }
        }
    }
}