Make the GUI lazy so it can handle huge transactions

gastronomy-ui/Cargo.toml

@@ -12,14 +12,16 @@ publish = false
 tauri-build = { version = "1.5.3", features = [] }
 
 [dependencies]
-figment = { version = "0.10", features = ["env", "toml"] }
 dashmap = "6"
+figment = { version = "0.10", features = ["env", "toml"] }
 gastronomy = { path = "../gastronomy" }
+pallas-codec = "0.30"
 serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 tauri = { version = "1.7.0", features = ["dialog-open"] }
 tauri-plugin-store = { git = "https://github.com/tauri-apps/plugins-workspace", branch = "v1" }
 tokio = "1.40"
+uuid = { version = "1", features = ["v4"] }
 
 [features]
 # this feature is used for production builds or when `devPath` points to the filesystem and the built-in dev server is disabled.

gastronomy-ui/src/execution_trace.rs

@@ -0,0 +1,113 @@
+use std::{collections::BTreeMap, fmt::Display};
+
+use gastronomy::{
+    execution_trace::{parse_context, parse_env, parse_raw_frames, parse_uplc_value, RawFrame},
+    uplc::{self, LoadedProgram, Program},
+    Frame,
+};
+use pallas_codec::flat::Flat;
+use tauri::InvokeError;
+use tokio::sync::{mpsc, oneshot};
+use uuid::Uuid;
+
+pub struct ExecutionTrace {
+    pub identifier: String,
+    worker_channel: mpsc::Sender<WorkerRequest>,
+}
+
+impl ExecutionTrace {
+    pub fn from_program(program: LoadedProgram) -> Result<Self, InvokeError> {
+        let identifier = Uuid::new_v4().to_string();
+
+        // The Aiken uplc crate uses lots of Rc<T> internally, so it's not Send.
+        // The string representation of a frame of execution can get HUGE, so we need to serialize it lazily.
+        // So, send the raw bytes to another thread, and interact with it over a channel.
+        let (worker_channel, requests) = mpsc::channel(16);
+        let worker = ExecutionTraceWorker {
+            raw_program: program.program.to_flat().map_err(to_invoke_error)?,
+            source_map: program.source_map,
+            requests,
+        };
+        std::thread::Builder::new()
+            .name(identifier.clone())
+            .stack_size(4 * 1024 * 1024)
+            .spawn(|| worker.run())
+            .map_err(to_invoke_error)?;
+
+        Ok(Self {
+            identifier,
+            worker_channel,
+        })
+    }
+    pub async fn frame_count(&self) -> Result<usize, InvokeError> {
+        let (frame_count_sink, frame_count_source) = oneshot::channel();
+        let request = WorkerRequest::FrameCount(frame_count_sink);
+        self.worker_channel
+            .send(request)
+            .await
+            .map_err(to_invoke_error)?;
+        frame_count_source.await.map_err(to_invoke_error)?
+    }
+    pub async fn get_frame(&self, frame: usize) -> Result<Frame, InvokeError> {
+        let (frame_sink, frame_source) = oneshot::channel();
+        let request = WorkerRequest::GetFrame(frame, frame_sink);
+        self.worker_channel
+            .send(request)
+            .await
+            .map_err(to_invoke_error)?;
+        frame_source.await.map_err(to_invoke_error)?
+    }
+}
+
+fn to_invoke_error<T: Display>(err: T) -> InvokeError {
+    InvokeError::from(err.to_string())
+}
+
+type ResponseChannel<T> = oneshot::Sender<Result<T, InvokeError>>;
+
+enum WorkerRequest {
+    FrameCount(ResponseChannel<usize>),
+    GetFrame(usize, ResponseChannel<Frame>),
+}
+
+struct ExecutionTraceWorker {
+    raw_program: Vec<u8>,
+    source_map: BTreeMap<u64, String>,
+    requests: mpsc::Receiver<WorkerRequest>,
+}
+
+impl ExecutionTraceWorker {
+    fn run(self) {
+        let program = Program::unflat(&self.raw_program).unwrap();
+        let states = uplc::execute_program(program).unwrap();
+        let frames = parse_raw_frames(&states, &self.source_map);
+
+        let mut requests = self.requests;
+        while let Some(request) = requests.blocking_recv() {
+            match request {
+                WorkerRequest::FrameCount(res) => {
+                    let _ = res.send(Ok(frames.len()));
+                }
+                WorkerRequest::GetFrame(index, res) => {
+                    let _ = res.send(Self::get_frame(index, &frames));
+                }
+            }
+        }
+    }
+
+    fn get_frame(index: usize, frames: &[RawFrame<'_>]) -> Result<Frame, InvokeError> {
+        let Some(raw) = frames.get(index) else {
+            return Err(InvokeError::from("Invalid frame index"));
+        };
+        let frame = Frame {
+            label: raw.label.to_string(),
+            context: parse_context(raw.context),
+            env: parse_env(raw.env),
+            term: raw.term.to_string(),
+            ret_value: raw.ret_value.map(|v| parse_uplc_value(v.clone())),
+            location: raw.location.cloned(),
+            budget: raw.budget.clone(),
+        };
+        Ok(frame)
+    }
+}

gastronomy-ui/src/main.rs

@@ -5,16 +5,17 @@ use std::path::{Path, PathBuf};
 
 use api::{CreateTraceResponse, GetFrameResponse, GetTraceSummaryResponse};
 use dashmap::DashMap;
+use execution_trace::ExecutionTrace;
 use figment::providers::{Env, Serialized};
 use gastronomy::{
     chain_query::ChainQuery,
     config::{load_base_config, Config},
-    ExecutionTrace,
 };
 use tauri::{InvokeError, Manager, State, Wry};
 use tauri_plugin_store::{with_store, StoreBuilder, StoreCollection};
 
 mod api;
+mod execution_trace;
 
 struct SessionState {
     traces: DashMap<String, ExecutionTrace>,
@@ -54,11 +55,12 @@ async fn create_traces<'a>(
         ChainQuery::None
     };
 
-    let mut traces = gastronomy::trace_executions(file, &parameters, query)
+    let mut programs = gastronomy::execution_trace::load_file(file, &parameters, query)
         .await
         .map_err(InvokeError::from_anyhow)?;
     let mut identifiers = vec![];
-    for trace in traces.drain(..) {
+    for program in programs.drain(..) {
+        let trace = ExecutionTrace::from_program(program)?;
         let identifier = trace.identifier.clone();
         state.traces.insert(identifier.clone(), trace);
         identifiers.push(identifier);
@@ -67,35 +69,30 @@ async fn create_traces<'a>(
 }
 
 #[tauri::command]
-fn get_trace_summary(
+async fn get_trace_summary(
     identifier: &str,
-    state: State<SessionState>,
+    state: State<'_, SessionState>,
 ) -> Result<api::GetTraceSummaryResponse, InvokeError> {
     println!("Getting summary");
     let Some(trace) = state.traces.get(identifier) else {
         return Err(InvokeError::from("Trace not found"));
     };
-    Ok(GetTraceSummaryResponse {
-        frame_count: trace.frames.len(),
-    })
+    let frame_count = trace.frame_count().await?;
+    Ok(GetTraceSummaryResponse { frame_count })
 }
 
 #[tauri::command]
-fn get_frame(
+async fn get_frame(
     identifier: &str,
     frame: usize,
-    state: State<SessionState>,
+    state: State<'_, SessionState>,
 ) -> Result<api::GetFrameResponse, InvokeError> {
     println!("Getting frame");
     let Some(trace) = state.traces.get(identifier) else {
         return Err(InvokeError::from("Trace not found"));
     };
-    let Some(frame) = trace.frames.get(frame) else {
-        return Err(InvokeError::from("Frame not found"));
-    };
-    Ok(GetFrameResponse {
-        frame: frame.clone(),
-    })
+    let frame = trace.get_frame(frame).await?;
+    Ok(GetFrameResponse { frame })
 }
 
 const STACK_SIZE: usize = 4 * 1024 * 1024;

gastronomy/Cargo.toml

@@ -22,4 +22,3 @@ serde = "1"
 serde_json = "1"
 uplc = { git = "https://github.com/SundaeSwap-finance/aiken.git", rev = "3c2ae7c" }
 # uplc = { path = "../../aiken/crates/uplc" }
-uuid = { version = "1", features = ["v4"] }