Fix exporting ply files on native

crates/brush-dataset/src/lib.rs

@@ -167,6 +167,7 @@ mod async_helpers {
     pub(super) fn spawn_future<T: 'static>(
         future: impl Future<Output = T> + 'static,
     ) -> JoinHandle<T> {
+        // On wasm, just spawn locally.
         task::spawn_local(future)
     }
 }

crates/brush-desktop/src/main.rs

@@ -42,6 +42,7 @@ fn main() -> anyhow::Result<()> {
             .dyn_into::<web_sys::HtmlCanvasElement>()
             .unwrap();
 
+        // On wasm, run as a local task.
         async_std::task::spawn_local(async {
             eframe::WebRunner::new()
                 .start(

crates/brush-viewer/src/async_lib.rs

@@ -0,0 +1,88 @@
+use std::future::Future;
+use web_time::Duration;
+use wgpu::WasmNotSend;
+
+/// Wrap a future to yield back to the browser macrotasks
+/// queue each time it yields (using a setTimeout).
+///
+/// In the future this could work better with scheduler.yield.
+pub fn with_timeout_yield<F: Future<Output = T> + WasmNotSend + 'static, T: Send + 'static>(
+    future: F,
+    min_time_between_yields: Duration,
+) -> impl Future<Output = T> + WasmNotSend + 'static {
+    #[cfg(not(target_family = "wasm"))]
+    {
+        let _ = min_time_between_yields;
+        future
+    }
+
+    #[cfg(target_family = "wasm")]
+    {
+        use gloo_timers::future::TimeoutFuture;
+        use std::{
+            pin::Pin,
+            task::{Context, Poll},
+        };
+        use web_time::Instant;
+
+        struct WithTimeoutYield<F> {
+            future: Pin<Box<F>>,
+            timeout: Option<TimeoutFuture>,
+            last_yield: Option<Instant>,
+            min_duration: Duration,
+        }
+
+        impl<F: Future> Future for WithTimeoutYield<F> {
+            type Output = F::Output;
+
+            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+                let this = self.get_mut();
+
+                if let Some(timeout) = &mut this.timeout {
+                    match std::pin::pin!(timeout).poll(cx) {
+                        Poll::Ready(_) => {
+                            this.timeout = None;
+                            this.last_yield = Some(Instant::now());
+                        }
+                        Poll::Pending => return Poll::Pending,
+                    }
+                }
+
+                match this.future.as_mut().poll(cx) {
+                    Poll::Ready(output) => Poll::Ready(output),
+                    Poll::Pending => {
+                        let should_yield = this
+                            .last_yield
+                            .map(|t| t.elapsed() >= this.min_duration)
+                            .unwrap_or(true);
+
+                        if should_yield && this.timeout.is_none() {
+                            this.timeout = Some(TimeoutFuture::new(0));
+                        }
+                        Poll::Pending
+                    }
+                }
+            }
+        }
+
+        WithTimeoutYield {
+            future: Box::pin(future),
+            timeout: None,
+            last_yield: None,
+            min_duration: min_time_between_yields,
+        }
+    }
+}
+
+pub fn spawn_future<T: WasmNotSend + 'static>(
+    fut: impl Future<Output = T> + WasmNotSend + 'static,
+) {
+    #[cfg(target_family = "wasm")]
+    {
+        async_std::task::spawn_local(fut);
+    }
+    #[cfg(not(target_family = "wasm"))]
+    {
+        async_std::task::spawn(fut);
+    }
+}

crates/brush-viewer/src/lib.rs

@@ -5,9 +5,9 @@
 use egui_tiles::SimplificationOptions;
 use viewer::{ViewerContext, ViewerMessage};
 
+mod async_lib;
 mod burn_texture;
 mod orbit_controls;
-mod timeout_future;
 
 mod panels;
 mod train_loop;

crates/brush-viewer/src/panels/scene.rs

@@ -1,4 +1,3 @@
-use async_std::task;
 use brush_dataset::splat_export;
 use egui::epaint::mutex::RwLock as EguiRwLock;
 use std::sync::Arc;
@@ -243,20 +242,32 @@ runs consider using the native app."#,
 
                             ui.add_space(15.0);
 
-                            if ui.button("↑ Export").clicked() {
+                            if ui.button("⬆ Export").clicked() {
                                 let splats = splats.clone();
-                                task::spawn_local(async move {
-                                    let data = splat_export::splat_to_ply(*splats).await;
-                                    let data = match data {
-                                        Ok(data) => data,
+
+                                crate::async_lib::spawn_future(async move {
+                                    let file = rrfd::save_file("export.ply").await;
+
+                                    // Not sure where/how to show this error if any.
+                                    match file {
                                         Err(e) => {
                                             log::error!("Failed to save file: {e}");
-                                            return;
                                         }
-                                    };
-                                    // Not sure where/how to show this error if any.
-                                    if let Err(e) = rrfd::save_file("export.ply", data).await {
-                                        log::error!("Failed to save file: {e}");
+                                        Ok(file) => {
+                                            let data = splat_export::splat_to_ply(*splats).await;
+
+                                            let data = match data {
+                                                Ok(data) => data,
+                                                Err(e) => {
+                                                    log::error!("Failed to serialize file: {e}");
+                                                    return;
+                                                }
+                                            };
+
+                                            if let Err(e) = file.write(&data).await {
+                                                log::error!("Failed to write file: {e}");
+                                            }
+                                        }
                                     }
                                 });
                             }

crates/brush-viewer/src/viewer.rs

@@ -4,7 +4,6 @@ use async_fn_stream::try_fn_stream;
 use async_std::{
     channel::{Receiver, Sender, TrySendError},
     stream::{Stream, StreamExt},
-    task,
 };
 use brush_dataset::{self, splat_import, Dataset, LoadDatasetArgs, LoadInitArgs, ZipData};
 use brush_render::camera::Camera;
@@ -20,6 +19,7 @@ use glam::{Quat, Vec3};
 use web_time::Instant;
 
 use crate::{
+    async_lib,
     orbit_controls::OrbitControls,
     panels::{DatasetPanel, LoadDataPanel, PresetsPanel, ScenePanel, StatsPanel},
     train_loop::{self, TrainMessage},
@@ -97,12 +97,14 @@ fn process_loop(
     let stream = try_fn_stream(|emitter| async move {
         let _ = emitter.emit(ViewerMessage::PickFile).await;
         let picked = rrfd::pick_file().await?;
+        let name = picked.file_name();
+
+        if name.contains(".ply") {
+            let data = picked.read().await;
 
-        if picked.file_name.contains(".ply") {
             let _ = emitter
                 .emit(ViewerMessage::StartLoading { training: false })
                 .await;
-            let data = picked.data;
             let splat_stream =
                 splat_import::load_splat_from_ply::<PrimaryBackend>(data, device.clone());
             let mut splat_stream = std::pin::pin!(splat_stream);
@@ -114,13 +116,15 @@ fn process_loop(
                     })
                     .await;
             }
-        } else if picked.file_name.contains(".zip") {
+        } else if name.contains(".zip") {
+            let data = picked.read().await;
+
             let _ = emitter
                 .emit(ViewerMessage::StartLoading { training: true })
                 .await;
 
             let stream = train_loop::train_loop(
-                ZipData::from(picked.data),
+                ZipData::from(data),
                 device,
                 train_receiver,
                 load_data_args,
@@ -215,17 +219,8 @@ impl ViewerContext {
             }
         };
 
-        #[cfg(target_family = "wasm")]
-        {
-            let fut =
-                crate::timeout_future::with_timeout_yield(fut, web_time::Duration::from_millis(5));
-            task::spawn_local(fut);
-        }
-
-        #[cfg(not(target_family = "wasm"))]
-        {
-            task::spawn(fut);
-        }
+        let fut = crate::async_lib::with_timeout_yield(fut, web_time::Duration::from_millis(5));
+        async_lib::spawn_future(fut);
     }
 
     pub fn send_train_message(&self, message: TrainMessage) {

crates/rrfd/src/android.rs

@@ -1,4 +1,3 @@
-use super::PickedFile;
 use anyhow::Result;
 use async_std::channel::Sender;
 use jni::objects::{GlobalRef, JByteArray, JClass, JStaticMethodID, JString};
@@ -8,13 +7,20 @@ use lazy_static::lazy_static;
 use std::sync::Arc;
 use std::sync::RwLock;
 
+#[derive(Clone, Debug)]
+pub struct PickedFile {
+    pub data: Vec<u8>,
+    pub file_name: String,
+}
+
 lazy_static! {
     static ref VM: RwLock<Option<Arc<jni::JavaVM>>> = RwLock::new(None);
     static ref CHANNEL: RwLock<Option<Sender<Result<PickedFile>>>> = RwLock::new(None);
     static ref START_FILE_PICKER: RwLock<Option<JStaticMethodID>> = RwLock::new(None);
     static ref FILE_PICKER_CLASS: RwLock<Option<GlobalRef>> = RwLock::new(None);
 }
 
+#[allow(unused)]
 pub fn jni_initialize(vm: Arc<jni::JavaVM>) {
     let mut env = vm.get_env().expect("Cannot get reference to the JNIEnv");
     let class = env.find_class("com/splats/app/FilePicker").unwrap();
@@ -30,6 +36,7 @@ pub fn jni_initialize(vm: Arc<jni::JavaVM>) {
     *VM.write().unwrap() = Some(vm);
 }
 
+#[allow(unused)]
 pub(crate) async fn pick_file() -> Result<PickedFile> {
     let (sender, receiver) = async_std::channel::bounded(1);
     {