ranim/cmd/render/

mod.rs

// MARK: Render api
use std::collections::VecDeque;

use crate::cmd::render::file_writer::OutputFormatExt;
use crate::{Output, Scene, SceneConstructor};
use bevy_app::SubApp;
use file_writer::{FileWriter, FileWriterBuilder};
use indicatif::{ProgressState, ProgressStyle};
use ranim_core::color::{self, LinearSrgb};
use ranim_core::{SceneEvaluator, TimeMark, frame::EvaluatedFrame};
use ranim_render::resource::RenderTextures;
use ranim_render::{RenderAppExt, RenderSetup, build_render_app, utils::WgpuContext};
use std::path::{Path, PathBuf};
use std::time::Duration;
use std::time::Instant;
use tracing::{Span, info, instrument, trace};
use tracing_indicatif::span_ext::IndicatifSpanExt;

pub(crate) mod audio;
pub(crate) mod file_writer;

/// Render the output basename template by replacing placeholders.
///
/// Supported placeholders:
/// - `{name}`: the scene/output name
/// - `{width}`: output width in pixels
/// - `{height}`: output height in pixels
/// - `{fps}`: output frame rate
fn render_output_basename(template: &str, name: &str, width: u32, height: u32, fps: u32) -> String {
    template
        .replace("{name}", name)
        .replace("{width}", &width.to_string())
        .replace("{height}", &height.to_string())
        .replace("{fps}", &fps.to_string())
}

/// Render a scene with all its outputs
pub fn render_scene(scene: &Scene, setup: &RenderSetup, buffer_count: usize) {
    for (i, output) in scene.outputs.iter().enumerate() {
        info!(
            "Rendering output {}/{} ({})",
            i + 1,
            scene.outputs.len(),
            output.format
        );
        render_scene_output(scene, setup, output, buffer_count);
    }
}

/// Render a scene output
pub fn render_scene_output(
    scene: &Scene,
    setup: &RenderSetup,
    output: &Output,
    buffer_count: usize,
) {
    render_scene_output_with_progress(scene, setup, output, buffer_count, None);
}

/// A single rendering job: one [`Output`] plus whether capture marks should
/// be processed after the main pass.
///
/// The CLI currently only creates full-scene jobs. Time-range and still-frame
/// rendering will extend this type together with the debug rendering commands.
pub struct RenderJob {
    /// Output settings used for this job.
    pub output: Output,
    /// Render and save every [`TimeMark::Capture`] after the main pass.
    ///
    /// This is enabled by `ranim output` (declared outputs) and disabled by
    /// the ad-hoc `ranim render` command.
    pub capture_marks: bool,
}

impl RenderJob {
    /// Creates a job that renders the whole scene.
    pub fn full(output: Output, capture_marks: bool) -> Self {
        Self {
            output,
            capture_marks,
        }
    }
}

/// Render one output with optional progress callback.
///
/// The callback receives `(current_frame, total_frames)` each frame.
pub fn render_scene_output_with_progress(
    scene: &Scene,
    setup: &RenderSetup,
    output: &Output,
    buffer_count: usize,
    on_progress: Option<Box<dyn Fn(u64, u64) + Send>>,
) {
    render_scene_job(
        scene,
        setup,
        RenderJob::full(output.clone(), true),
        buffer_count,
        on_progress,
    );
}

/// Render a single ad-hoc output for a scene.
///
/// This is the low-level rendering entry point used by `ranim render`: it
/// renders the whole scene once with the given [`Output`], but unlike
/// [`render_scene_output_with_progress`] it does not process capture marks.
pub fn render_scene_once(scene: &Scene, setup: &RenderSetup, output: &Output, buffer_count: usize) {
    render_scene_job(
        scene,
        setup,
        RenderJob::full(output.clone(), false),
        buffer_count,
        None,
    );
}

/// Build the scene and run a [`RenderJob`].
pub fn render_scene_job(
    scene: &Scene,
    setup: &RenderSetup,
    job: RenderJob,
    buffer_count: usize,
    on_progress: Option<Box<dyn Fn(u64, u64) + Send>>,
) {
    use std::time::Instant;

    let RenderJob {
        output,
        capture_marks,
    } = job;

    info!(
        "Output: {}x{} {}fps {} dir={:?} save_frames={}",
        output.width, output.height, output.fps, output.format, output.dir, output.save_frames
    );

    let t = Instant::now();
    let timeline = scene.build_scene();
    trace!("Build timeline cost: {:?}", t.elapsed());

    // Default logic grid 120 Hz (per the time model design); render fps only
    // decides which logic states are sampled.
    let mut evaluator = timeline.into_evaluator(DEFAULT_LOGIC_FPS);

    let mut app = RanimRenderApp::new(scene, setup, &output, buffer_count);
    app.render_scene_with_progress(&mut evaluator, on_progress);

    if capture_marks && !evaluator.time_marks().is_empty() {
        app.render_capture_marks(&mut evaluator);
    }

    // The audio plane was baked in one offline pass at seal: mux the baked
    // buffer into the finished file.
    if evaluator.has_audio() {
        match audio::audio_codec(output.format) {
            Some(_) => {
                let pcm = evaluator.audio();
                // Finalize the encoder before reading the file back for muxing.
                app.finish_video();
                let video_path = app.video_path();
                if !video_path.exists() {
                    tracing::warn!("no video file was written, skipping audio muxing");
                } else if let Err(err) = audio::mux_audio_into_video(video_path, pcm, output.format)
                {
                    tracing::error!("failed to mux audio into {}: {err:#}", video_path.display());
                }
            }
            None => tracing::warn!(
                "the scene has audio, but the {:?} format cannot store it; rendering silently",
                output.format
            ),
        }
    }
}

/// Default logic grid resolution (Hz), per the time model design.
const DEFAULT_LOGIC_FPS: f64 = 120.0;

/// Handle to the background render thread, used to submit frame data from the main thread.
///
/// Dropping this handle closes the submission channel and terminates the worker thread loop.
pub struct RenderThreadHandle {
    submit_frame_tx: async_channel::Sender<EvaluatedFrame>,
    back_rx: async_channel::Receiver<EvaluatedFrame>,
    worker_rx: async_channel::Receiver<RenderWorker>,
}

impl RenderThreadHandle {
    /// Reuses an [`EvaluatedFrame`], applies `f` to it, and submits it to the background render thread.
    pub fn sync_and_submit(&self, f: impl FnOnce(&mut EvaluatedFrame)) {
        let mut store = self.get_store();
        f(&mut store);
        self.submit_frame_tx.send_blocking(store).unwrap();
    }

    /// Retrieves a previously rendered [`EvaluatedFrame`] from the background thread for reuse.
    pub fn get_store(&self) -> EvaluatedFrame {
        self.back_rx.recv_blocking().unwrap()
    }

    /// Closes the submission channel, waits for the background thread to finish, and returns the [`RenderWorker`].
    pub fn retrive(&self) -> RenderWorker {
        self.submit_frame_tx.close(); // This terminates the worker thread loop
        self.worker_rx.recv_blocking().unwrap()
    }
}

/// Background rendering worker that owns the wgpu context, render sub-app, output textures, and video/image writers.
pub struct RenderWorker {
    ctx: WgpuContext,
    render_app: SubApp,
    render_textures: Vec<RenderTextures>,
    clear_color: wgpu::Color,
    // video writer
    video_writer: Option<FileWriter>,
    video_writer_builder: Option<FileWriterBuilder>,
    video_path: PathBuf,
    save_frames: bool,
    output_dir: PathBuf,
    scene_name: String,
    width: u32,
    height: u32,
    fps: u32,
}

impl RenderWorker {
    /// Creates a new [`RenderWorker`].
    ///
    /// This checks for or downloads `ffmpeg`, initializes the wgpu context and the
    /// render sub-app (configured by `setup`), and configures the video writer from [`Output`].
    ///
    /// # Panics
    ///
    /// Panics if `buffer_count < 1`.
    pub fn new(scene: &Scene, setup: &RenderSetup, output: &Output, buffer_count: usize) -> Self {
        assert!(buffer_count >= 1, "buffer_count must be at least 1");
        let scene_name = scene.name.clone();
        info!("Checking ffmpeg...");
        let t = Instant::now();
        if let Ok(ffmpeg_path) = which::which("ffmpeg") {
            info!("ffmpeg found at {ffmpeg_path:?}");
        } else {
            use std::path::Path;

            info!(
                "ffmpeg not found from path env, searching in {:?}...",
                Path::new("./").canonicalize().unwrap()
            );
            if Path::new("./ffmpeg").exists() {
                info!("ffmpeg found at current working directory")
            } else {
                info!("ffmpeg not found at current working directory, downloading...");
                download_ffmpeg("./").expect("failed to download ffmpeg");
            }
        }
        trace!("Check ffmmpeg cost: {:?}", t.elapsed());

        let t = Instant::now();
        info!("Creating wgpu context...");
        let ctx = pollster::block_on(WgpuContext::new());
        trace!("Create wgpu context cost: {:?}", t.elapsed());

        let mut output_dir = PathBuf::from(&output.dir);
        if !output_dir.is_absolute() {
            output_dir = std::env::current_dir().unwrap().join(output_dir);
        }
        let mut render_app = build_render_app(&ctx, output.width, output.height, 8);
        render_app.add_plugins(setup.clone());
        render_app.finish();
        render_app.cleanup();
        let render_textures: Vec<RenderTextures> = (0..buffer_count)
            .map(|_| RenderTextures::new(&ctx, output.width, output.height))
            .collect();
        let clear_color = color::try_color(&scene.config.clear_color)
            .unwrap_or(color::color("#333333ff"))
            .convert::<LinearSrgb>();
        let [r, g, b, a] = clear_color.components.map(|x| x as f64);
        let clear_color = wgpu::Color { r, g, b, a };
        let (_, _, ext) = output.format.encoding_params();
        let video_path = output_dir.join({
            let template = output
                .name_template
                .as_deref()
                .unwrap_or("{name}_{width}x{height}_{fps}");
            let base_name = render_output_basename(
                template,
                output.name.as_deref().unwrap_or(&scene_name),
                output.width,
                output.height,
                output.fps,
            );
            format!("{base_name}.{ext}")
        });
        Self {
            ctx,
            render_app,
            render_textures,
            clear_color,
            video_writer: None,
            video_writer_builder: Some(
                FileWriterBuilder::default()
                    .with_fps(output.fps)
                    .with_size(output.width, output.height)
                    .with_file_path(video_path.clone())
                    .with_output_format(output.format),
            ),
            video_path,
            save_frames: output.save_frames,
            output_dir,
            scene_name,
            width: output.width,
            height: output.height,
            fps: output.fps,
        }
    }

    /// Returns the directory path used to save individual frame images.
    pub fn save_frame_dir(&self) -> PathBuf {
        self.output_dir.join(format!(
            "{}_{}x{}_{}-frames",
            self.scene_name, self.width, self.height, self.fps
        ))
    }

    /// Returns the path of the video file being written.
    pub fn video_path(&self) -> &Path {
        &self.video_path
    }

    /// Finalize the video file: flush the encoder and wait for it to exit.
    ///
    /// After this call the file on disk is complete (valid container index)
    /// and no further frames can be written.
    pub fn finish_video(&mut self) {
        self.video_writer = None;
    }

    /// Moves this [`RenderWorker`] into a newly spawned background thread and returns a [`RenderThreadHandle`].
    ///
    /// The background thread loops over incoming [`EvaluatedFrame`]s, rendering and outputting frames using multi-buffered async readback.
    pub fn yeet(self) -> RenderThreadHandle {
        let (submit_frame_tx, submit_frame_rx) = async_channel::bounded(1);
        let (back_tx, back_rx) = async_channel::bounded(2);
        let (worker_tx, worker_rx) = async_channel::bounded(1);

        // Evaluate one frame while the worker renders the other, reusing both.
        for _ in 0..2 {
            back_tx.send_blocking(EvaluatedFrame::default()).unwrap();
        }
        std::thread::spawn(move || {
            let mut worker = self;
            let n = worker.render_textures.len();
            let mut frame_count = 0u64;
            let mut cur = 0usize;
            let mut pending: VecDeque<(usize, u64)> = VecDeque::new();

            while let Ok(mut store) = submit_frame_rx.recv_blocking() {
                // Drain oldest pending readback if all targets are occupied
                if pending.len() >= n {
                    let (prev, prev_fc) = pending.pop_front().unwrap();
                    worker.render_textures[prev].finish_readback(&worker.ctx);
                    worker.output_frame_from(prev, prev_fc);
                }

                // Render current frame and start async readback
                worker.render_app.render_frame(
                    &mut worker.render_textures[cur],
                    worker.clear_color,
                    &mut store,
                );
                worker.render_textures[cur].start_readback(&worker.ctx);

                pending.push_back((cur, frame_count));
                frame_count += 1;
                cur = (cur + 1) % n;

                // Return store early so main thread can eval next frame
                // while GPU processes the readback
                back_tx.send_blocking(store).unwrap();

                // Now try to drain any completed readbacks while we wait
                // for the next frame from the main thread
                while let Some(&(prev, _)) = pending.front() {
                    // Non-blocking: check if the oldest readback is ready
                    if !worker.render_textures[prev].try_finish_readback(&worker.ctx) {
                        break;
                    }
                    let (prev, prev_fc) = pending.pop_front().unwrap();
                    worker.output_frame_from(prev, prev_fc);
                }
            }

            // Flush all remaining pending frames
            while let Some((prev, prev_fc)) = pending.pop_front() {
                worker.render_textures[prev].finish_readback(&worker.ctx);
                worker.output_frame_from(prev, prev_fc);
            }

            worker_tx.send_blocking(worker).unwrap();
        });
        RenderThreadHandle {
            submit_frame_tx,
            back_rx,
            worker_rx,
        }
    }

    /// Renders a single frame synchronously on this worker (typically used for captures).
    ///
    /// The frame's data and capacity remain available for reuse after rendering.
    pub fn render_store(&mut self, store: &mut EvaluatedFrame) {
        self.render_app
            .render_frame(&mut self.render_textures[0], self.clear_color, store);
    }

    /// Write and save (if [`Self::save_frames`] is true)
    pub fn output_frame_from(&mut self, target_idx: usize, frame_number: u64) {
        self.write_frame_from(target_idx);
        if self.save_frames {
            self.save_frame_from(target_idx, frame_number);
        }
    }

    /// Write frame data from the given target to the video file.
    pub fn write_frame_from(&mut self, target_idx: usize) {
        let data = self.render_textures[target_idx]
            .render_texture
            .texture_data();
        if let Some(video_writer) = self.video_writer.as_mut() {
            video_writer.write_frame(data);
        } else if let Some(builder) = self.video_writer_builder.as_ref() {
            self.video_writer
                .get_or_insert(builder.clone().build())
                .write_frame(data);
        }
    }

    /// Save frame from the given target as a PNG image.
    pub fn save_frame_from(&mut self, target_idx: usize, frame_number: u64) {
        let path = self.save_frame_dir().join(format!("{frame_number:04}.png"));
        let dir = path.parent().unwrap();
        if !dir.exists() || !dir.is_dir() {
            std::fs::create_dir_all(dir).unwrap();
        }
        // Data is already in cpu buffer after finish_readback, this won't trigger GPU work
        let buffer = self.render_textures[target_idx].get_rendered_texture_img_buffer(&self.ctx);
        buffer.save(path).unwrap();
    }

    /// Capture frame to image file (sync path, uses target 0).
    pub fn capture_frame(&mut self, path: impl AsRef<Path>) {
        let path = path.as_ref();
        let path = if !path.is_absolute() {
            self.output_dir
                .join(format!(
                    "{}_{}x{}_{}",
                    self.scene_name, self.width, self.height, self.fps
                ))
                .join(path)
        } else {
            path.to_path_buf()
        };
        let dir = path.parent().unwrap();
        if !dir.exists() || !dir.is_dir() {
            std::fs::create_dir_all(dir).unwrap();
        }
        let buffer = self.render_textures[0].get_rendered_texture_img_buffer(&self.ctx);
        buffer.save(path).unwrap();
    }
}

/// MARK: RanimRenderApp
/// Application wrapper that drives frame-by-frame scene rendering, video output, and capture screenshots.
pub struct RanimRenderApp {
    render_worker: Option<RenderWorker>,
    fps: u32,
    store: EvaluatedFrame,
}

impl RanimRenderApp {
    /// Creates the application and initializes the internal [`RenderWorker`].
    pub fn new(scene: &Scene, setup: &RenderSetup, output: &Output, buffer_count: usize) -> Self {
        let render_worker = RenderWorker::new(scene, setup, output, buffer_count);
        Self {
            render_worker: Some(render_worker),
            fps: output.fps,
            store: EvaluatedFrame::default(),
        }
    }

    /// Returns the path of the video file written by this app's worker.
    pub fn video_path(&self) -> &Path {
        self.render_worker
            .as_ref()
            .expect("render worker must exist outside rendering")
            .video_path()
    }

    /// Finalizes the video file (flushes the encoder and waits for it).
    pub fn finish_video(&mut self) {
        if let Some(worker) = self.render_worker.as_mut() {
            worker.finish_video();
        }
    }

    /// Renders the entire scene and reports progress via a tracing progress bar and an optional `on_progress` callback.
    ///
    /// The frame count is `evaluator.total_secs() * fps` rounded up, with an extra frame appended when necessary to sample the final state.
    #[instrument(skip_all)]
    pub fn render_scene_with_progress(
        &mut self,
        evaluator: &mut SceneEvaluator,
        on_progress: Option<Box<dyn Fn(u64, u64) + Send>>,
    ) {
        let start = Instant::now();
        let worker_thread = self.render_worker.take().unwrap().yeet();

        let total_secs = evaluator.total_secs();
        let fps = self.fps as f64;
        let raw_frames = total_secs * fps;
        // Add an extra frame to sample the final state exactly,
        // unless total_secs * fps is already an integer (last frame lands on total_secs).
        let n = raw_frames.ceil() as u64;
        let num_frames = if (raw_frames - raw_frames.round()).abs() < 1e-9 {
            n
        } else {
            n + 1
        };
        let style =             ProgressStyle::with_template(
                "[{elapsed_precise}] [{wide_bar:.cyan/blue}] frame {human_pos}/{human_len} (eta {eta}) {msg}",
            )
            .unwrap()
            .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| {
                write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
            })
            .progress_chars("#>-");

        let span = Span::current();
        span.pb_set_style(&style);
        span.pb_set_length(num_frames);

        (0..num_frames)
            .map(|f| (f as f64 / fps).min(total_secs))
            .enumerate()
            .for_each(|(i, sec)| {
                worker_thread.sync_and_submit(|store| {
                    evaluator.sample_at(sec, store);
                });

                span.pb_inc(1);
                if let Some(cb) = &on_progress {
                    cb(i as u64 + 1, num_frames);
                }
                span.pb_set_message(
                    format!(
                        "rendering {:.1?}/{:.1?}",
                        Duration::from_secs_f64(sec),
                        Duration::from_secs_f64(total_secs)
                    )
                    .as_str(),
                );
            });
        self.render_worker.replace(worker_thread.retrive());

        info!(
            "rendered {} frames({:?}) in {:?}",
            num_frames,
            Duration::from_secs_f64(evaluator.total_secs()),
            start.elapsed(),
        );
        trace!("render timeline cost: {:?}", start.elapsed());
    }

    /// Renders and saves screenshots for every [`TimeMark::Capture`] mark on the timeline.
    #[instrument(skip_all)]
    pub fn render_capture_marks(&mut self, evaluator: &mut SceneEvaluator) {
        let start = Instant::now();
        let timemarks = evaluator
            .time_marks()
            .iter()
            .filter(|mark| matches!(mark.1, TimeMark::Capture(_)))
            .map(|(sec, mark)| (*sec, mark.clone()))
            .collect::<Vec<_>>();

        let style =             ProgressStyle::with_template(
                "[{elapsed_precise}] [{wide_bar:.cyan/blue}] frame {human_pos}/{human_len} (eta {eta}) {msg}",
            )
            .unwrap()
            .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| {
                write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
            })
            .progress_chars("#>-");

        let span = Span::current();
        span.pb_set_style(&style);
        span.pb_set_length(timemarks.len() as u64);
        let _enter = span.enter();

        let mut captured = 0usize;
        for (sec, TimeMark::Capture(filename)) in timemarks {
            // The render has advanced to the end; captures seek back and replay
            // (deterministic contract, direction handled inside sample_at).
            evaluator.sample_at(sec, &mut self.store);
            let worker = self.render_worker.as_mut().unwrap();
            worker.render_store(&mut self.store);
            worker.capture_frame(filename);
            span.pb_inc(1);
            captured += 1;
        }
        info!("saved {} capture frames from time marks", captured);

        trace!("save capture frames cost: {:?}", start.elapsed());
    }
}

// MARK: Download ffmpeg
const FFMPEG_RELEASE_URL: &str = "https://github.com/eugeneware/ffmpeg-static/releases/latest";

/// Returns the directory of the current executable.
#[allow(unused)]
pub(crate) fn exe_dir() -> PathBuf {
    std::env::current_exe()
        .unwrap()
        .parent()
        .unwrap()
        .to_path_buf()
}

/// Download latest release of ffmpeg from <https://github.com/eugeneware/ffmpeg-static/releases/latest> to <target_dir>/ffmpeg
pub fn download_ffmpeg(target_dir: impl AsRef<Path>) -> Result<PathBuf, anyhow::Error> {
    use anyhow::Context;
    use itertools::Itertools;
    use std::io::Read;
    use tracing::info;

    let target_dir = target_dir.as_ref();

    let res = reqwest::blocking::get(FFMPEG_RELEASE_URL).context("failed to get release url")?;
    let url = res.url().to_string();
    let url = url.split("tag").collect_array::<2>().unwrap();
    let url = format!("{}/download/{}", url[0], url[1]);
    info!("ffmpeg release url: {url:?}");

    #[cfg(all(target_os = "windows", target_arch = "x86_64"))]
    let url = format!("{url}/ffmpeg-win32-x64.gz");
    #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
    let url = format!("{url}/ffmpeg-linux-x64.gz");
    #[cfg(all(target_os = "linux", target_arch = "aarch64"))]
    let url = format!("{url}/ffmpeg-linux-arm64.gz");
    #[cfg(all(target_os = "macos", target_arch = "x86_64"))]
    let url = format!("{url}/ffmpeg-darwin-x64.gz");
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    let url = format!("{url}/ffmpeg-darwin-arm64.gz");

    info!("downloading ffmpeg from {url:?}...");

    let res = reqwest::blocking::get(&url).context("get err")?;
    let mut decoder = flate2::bufread::GzDecoder::new(std::io::BufReader::new(
        std::io::Cursor::new(res.bytes().unwrap()),
    ));
    let mut bytes = Vec::new();
    decoder
        .read_to_end(&mut bytes)
        .context("GzDecoder decode err")?;
    let ffmpeg_path = target_dir.join("ffmpeg");
    std::fs::write(&ffmpeg_path, bytes).unwrap();

    #[cfg(target_family = "unix")]
    {
        use std::os::unix::fs::PermissionsExt;

        std::fs::set_permissions(&ffmpeg_path, std::fs::Permissions::from_mode(0o755))?;
    }
    info!("ffmpeg downloaded to {target_dir:?}");
    Ok(ffmpeg_path)
}

#[cfg(test)]
mod tests {
    use super::render_output_basename;

    #[test]
    fn render_output_basename_substitutes_template_placeholders() {
        let cases = [
            ("{name}_{width}x{height}_{fps}", "my_scene_1920x1080_60"),
            (
                "{name}_{fps}fps_{width}x{height}",
                "my_scene_60fps_1920x1080",
            ),
            ("{name}", "my_scene"),
        ];
        for (template, expected) in cases {
            assert_eq!(
                render_output_basename(template, "my_scene", 1920, 1080, 60),
                expected,
                "template {template}"
            );
        }
    }
}