add tsv-video-buffer

3 files changed, 418 insertions, 54 deletions

diff --git a/src/bin/tsv_video_buffer.rs b/src/bin/tsv_video_buffer.rs
new file mode 100644
index 0000000..7fc94b8
--- /dev/null
+++ b/src/bin/tsv_video_buffer.rs
@@ -0,0 +1,410 @@
+use std::io::Read;
+use std::process::{Command, Stdio};
+
+use ash::vk;
+use wgsl_view::{gpu, hap};
+
+/// Options that ffmpeg supports but ffprobe does not.
+const FFMPEG_ONLY_OPTIONS: &[&str] = &["-stream_loop", "-t"];
+
+fn filter_probe_args(ff_args: &[String]) -> Vec<String> {
+    let mut out = Vec::new();
+    let mut skip_next = false;
+    for arg in ff_args {
+        if skip_next {
+            skip_next = false;
+            continue;
+        }
+        if FFMPEG_ONLY_OPTIONS.contains(&arg.as_str()) {
+            skip_next = true;
+            continue;
+        }
+        out.push(arg.clone());
+    }
+    out
+}
+
+struct VideoInfo {
+    width: u32,
+    height: u32,
+    num_frames: u32,
+    codec_tag: String,
+}
+
+fn probe_video(ff_args: &[String], max_frames: Option<u32>) -> VideoInfo {
+    let probe_args = filter_probe_args(ff_args);
+    let output = Command::new("ffprobe")
+        .args(["-v", "error", "-select_streams", "v:0"])
+        .args([
+            "-show_entries",
+            "stream=width,height,nb_frames,codec_tag_string",
+        ])
+        .args(["-of", "csv=p=0"])
+        .args(&probe_args)
+        .output()
+        .expect("failed to run ffprobe");
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        panic!("ffprobe failed: {stderr}");
+    }
+
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    let line = stdout.trim();
+    let parts: Vec<&str> = line.split(',').collect();
+    if parts.len() < 4 {
+        panic!("unexpected ffprobe output: {line}");
+    }
+
+    let num_frames = parts[3].parse().unwrap_or_else(|_| {
+        max_frames.unwrap_or_else(|| {
+            panic!(
+                "could not determine frame count (nb_frames={}).\n\
+                 Use --frames N or use tsv-video-stream instead.",
+                parts[3]
+            )
+        })
+    });
+
+    VideoInfo {
+        codec_tag: parts[0].to_string(),
+        width: parts[1].parse().expect("parse width"),
+        height: parts[2].parse().expect("parse height"),
+        num_frames,
+    }
+}
+
+fn is_hap_codec(codec_tag: &str) -> Option<hap::HapFormat> {
+    match codec_tag {
+        "Hap1" => Some(hap::HapFormat::Bc1),
+        "Hap5" => Some(hap::HapFormat::Bc3),
+        _ => None,
+    }
+}
+
+fn main() {
+    env_logger::init();
+
+    let args: Vec<String> = std::env::args().skip(1).collect();
+    let sep = args.iter().position(|a| a == "--");
+
+    let (own_args, ff_args) = match sep {
+        Some(i) => (&args[..i], &args[i + 1..]),
+        None => {
+            eprintln!("usage: tsv-video-buffer [--name NAME] [--array] [--frames N] -- <ffmpeg input args>");
+            std::process::exit(1);
+        }
+    };
+
+    let mut name = "tsv-video-buffer".to_string();
+    let mut use_array = false;
+    let mut max_frames: Option<u32> = None;
+
+    let mut i = 0;
+    while i < own_args.len() {
+        match own_args[i].as_str() {
+            "--name" => {
+                name = own_args[i + 1].clone();
+                i += 2;
+            }
+            "--array" => {
+                use_array = true;
+                i += 1;
+            }
+            "--frames" => {
+                max_frames = Some(own_args[i + 1].parse().expect("parse --frames"));
+                i += 2;
+            }
+            other => panic!("unknown argument: {other}"),
+        }
+    }
+
+    let mut info = probe_video(ff_args, max_frames);
+    if let Some(max) = max_frames {
+        info.num_frames = info.num_frames.min(max);
+    }
+    let hap_format = is_hap_codec(&info.codec_tag);
+
+    let (texture_format, tsv_img_format) = match hap_format {
+        Some(hap::HapFormat::Bc1) => (
+            wgpu::TextureFormat::Bc1RgbaUnorm,
+            gpu::ImgFormat::BC1_RGBA,
+        ),
+        Some(hap::HapFormat::Bc3) => (
+            wgpu::TextureFormat::Bc3RgbaUnorm,
+            gpu::ImgFormat::BC3_RGBA,
+        ),
+        None => (wgpu::TextureFormat::Rgba8UnormSrgb, gpu::ImgFormat::R8G8B8A8),
+    };
+
+    log::info!(
+        "{}x{}, {} frames, codec_tag={}, format={:?}, tsv image: {}",
+        info.width,
+        info.height,
+        info.num_frames,
+        info.codec_tag,
+        texture_format,
+        name
+    );
+
+    // Request BC compression features if needed.
+    let features = if hap_format.is_some() {
+        wgpu::Features::TEXTURE_COMPRESSION_BC
+            | if use_array {
+                wgpu::Features::empty()
+            } else {
+                wgpu::Features::TEXTURE_COMPRESSION_BC_SLICED_3D
+            }
+    } else {
+        wgpu::Features::empty()
+    };
+
+    let (tsv_img_type, wgpu_dimension) = if use_array {
+        (gpu::ImgType::D2, wgpu::TextureDimension::D2)
+    } else {
+        (gpu::ImgType::D3, wgpu::TextureDimension::D3)
+    };
+
+    let instance = gpu::create_instance();
+    let adapter = gpu::create_adapter(&instance, None);
+    let (device, queue) = gpu::create_device_with_features(&adapter, features);
+
+    let mut client = gpu::create_tsv_client(&device);
+    client
+        .init_image(
+            &name,
+            info.width,
+            info.height,
+            info.num_frames,
+            tsv_img_format,
+            tsv_img_type,
+            true,
+        )
+        .expect("init tsv image");
+
+    let fence = gpu::create_fence(&client);
+
+    let texture = device.create_texture(&wgpu::TextureDescriptor {
+        label: Some("video_buffer"),
+        size: wgpu::Extent3d {
+            width: info.width,
+            height: info.height,
+            depth_or_array_layers: info.num_frames,
+        },
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu_dimension,
+        format: texture_format,
+        usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC,
+        view_formats: &[],
+    });
+
+    if let Some(hap_fmt) = hap_format {
+        upload_hap_frames(&device, &queue, &texture, ff_args, &info, hap_fmt);
+    } else {
+        upload_raw_frames(&device, &queue, &texture, ff_args, &info);
+    }
+
+    // Submit all pending writes and wait for GPU
+    queue.submit([]);
+    device.poll(wgpu::PollType::wait_indefinitely()).unwrap();
+
+    // Share the complete texture via TSV
+    let raw = unsafe { gpu::raw_image(&texture) };
+    if let Err(e) = client.send_image(
+        &name,
+        raw,
+        vk::ImageLayout::TRANSFER_DST_OPTIMAL,
+        vk::ImageLayout::TRANSFER_DST_OPTIMAL,
+        fence,
+    ) {
+        log::error!("send_image error: {e}");
+        std::process::exit(1);
+    }
+
+    log::info!("all {} frames uploaded and shared, waiting...", info.num_frames);
+
+    // Keep the process alive so the shared image remains available
+    loop {
+        std::thread::sleep(std::time::Duration::from_secs(3600));
+    }
+}
+
+fn upload_hap_frames(
+    _device: &wgpu::Device,
+    queue: &wgpu::Queue,
+    texture: &wgpu::Texture,
+    ff_args: &[String],
+    info: &VideoInfo,
+    hap_fmt: hap::HapFormat,
+) {
+    // For HAP, we need raw packet data (not decoded pixels).
+    // Use ffmpeg-next to read packets directly.
+    ffmpeg_next::init().expect("ffmpeg init");
+
+    let (_, input_url) = parse_ff_input_args(ff_args);
+    let mut ictx = ffmpeg_next::format::input(input_url).expect("open input");
+
+    let video_stream_index = ictx
+        .streams()
+        .best(ffmpeg_next::media::Type::Video)
+        .expect("no video stream")
+        .index();
+
+    let compressed_frame_size = hap_fmt.compressed_size(info.width, info.height);
+    let mut decode_buf = vec![0u8; compressed_frame_size + 1024 * 1024]; // extra margin
+    let mut frame_index: u32 = 0;
+
+    let block_bytes = hap_fmt.block_bytes();
+    let blocks_x = info.width.div_ceil(4);
+
+    for (stream, packet) in ictx.packets() {
+        if stream.index() != video_stream_index {
+            continue;
+        }
+
+        let data = packet.data().expect("packet data");
+        let (_, decoded_len) = hap::decode(data, &mut decode_buf).expect("HAP decode");
+
+        queue.write_texture(
+            wgpu::TexelCopyTextureInfo {
+                texture,
+                mip_level: 0,
+                origin: wgpu::Origin3d {
+                    x: 0,
+                    y: 0,
+                    z: frame_index,
+                },
+                aspect: wgpu::TextureAspect::All,
+            },
+            &decode_buf[..decoded_len],
+            wgpu::TexelCopyBufferLayout {
+                offset: 0,
+                // BC: bytes_per_row = blocks_x * block_bytes
+                bytes_per_row: Some(blocks_x * block_bytes),
+                // rows_per_image = height in blocks
+                rows_per_image: Some(info.height.div_ceil(4)),
+            },
+            wgpu::Extent3d {
+                width: info.width,
+                height: info.height,
+                depth_or_array_layers: 1,
+            },
+        );
+
+        frame_index += 1;
+        if frame_index.is_multiple_of(10) || frame_index == info.num_frames {
+            log::info!("uploaded frame {}/{}", frame_index, info.num_frames);
+        }
+        if frame_index >= info.num_frames {
+            break;
+        }
+    }
+
+    if frame_index < info.num_frames {
+        log::warn!(
+            "expected {} frames but only decoded {frame_index}",
+            info.num_frames
+        );
+    }
+}
+
+fn upload_raw_frames(
+    _device: &wgpu::Device,
+    queue: &wgpu::Queue,
+    texture: &wgpu::Texture,
+    ff_args: &[String],
+    info: &VideoInfo,
+) {
+    let frame_size = (info.width * info.height * 4) as usize;
+
+    let mut ffmpeg = Command::new("ffmpeg")
+        .args(["-v", "quiet"])
+        .args(ff_args)
+        .args(["-f", "rawvideo", "-pix_fmt", "rgba", "pipe:1"])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::inherit())
+        .spawn()
+        .expect("failed to start ffmpeg");
+
+    let stdout = ffmpeg.stdout.take().unwrap();
+    let mut reader = std::io::BufReader::new(stdout);
+    let mut frame_buf = vec![0u8; frame_size];
+    let mut frame_index: u32 = 0;
+
+    while reader.read_exact(&mut frame_buf).is_ok() {
+        queue.write_texture(
+            wgpu::TexelCopyTextureInfo {
+                texture,
+                mip_level: 0,
+                origin: wgpu::Origin3d {
+                    x: 0,
+                    y: 0,
+                    z: frame_index,
+                },
+                aspect: wgpu::TextureAspect::All,
+            },
+            &frame_buf,
+            wgpu::TexelCopyBufferLayout {
+                offset: 0,
+                bytes_per_row: Some(info.width * 4),
+                rows_per_image: Some(info.height),
+            },
+            wgpu::Extent3d {
+                width: info.width,
+                height: info.height,
+                depth_or_array_layers: 1,
+            },
+        );
+
+        frame_index += 1;
+        if frame_index.is_multiple_of(10) || frame_index == info.num_frames {
+            log::info!("uploaded frame {}/{}", frame_index, info.num_frames);
+        }
+        if frame_index >= info.num_frames {
+            break;
+        }
+    }
+
+    // Drop the reader/stdout to close the pipe — this signals ffmpeg to exit
+    // (important when --frames truncates before ffmpeg finishes)
+    drop(reader);
+    ffmpeg.wait().ok();
+
+    if frame_index < info.num_frames {
+        log::warn!(
+            "expected {} frames but only decoded {frame_index}",
+            info.num_frames
+        );
+    }
+}
+
+/// Parse ffmpeg-style input args to extract -f <format> and the input URL.
+/// Returns (Option<format>, input_url).
+fn parse_ff_input_args(ff_args: &[String]) -> (Option<String>, &str) {
+    let mut format = None;
+    let mut input_url = None;
+    let mut i = 0;
+
+    while i < ff_args.len() {
+        match ff_args[i].as_str() {
+            "-f" => {
+                format = Some(ff_args[i + 1].clone());
+                i += 2;
+            }
+            "-i" => {
+                input_url = Some(ff_args[i + 1].as_str());
+                i += 2;
+            }
+            _ => {
+                // Last positional arg is input URL if no -i was given
+                if i == ff_args.len() - 1 && input_url.is_none() {
+                    input_url = Some(ff_args[i].as_str());
+                }
+                i += 1;
+            }
+        }
+    }
+
+    (format, input_url.expect("no input file specified in ffmpeg args"))
+}
diff --git a/src/bin/tsv_video_stream.rs b/src/bin/tsv_video_stream.rs
index d68f4b4..22b38f8 100644
--- a/src/bin/tsv_video_stream.rs
+++ b/src/bin/tsv_video_stream.rs
@@ -98,7 +98,7 @@ fn main() {
 
     let mut client = gpu::create_tsv_client(&device);
     client
-        .init_image(&name, width, height, gpu::TSV_FORMAT, true)
+        .init_image(&name, width, height, 1, gpu::TSV_FORMAT, gpu::ImgType::D2, true)
         .expect("init tsv image");
 
     let fence = gpu::create_fence(&client);
diff --git a/src/bin/wgsl_render.rs b/src/bin/wgsl_render.rs
index 1860b7c..29b5d5d 100644
--- a/src/bin/wgsl_render.rs
+++ b/src/bin/wgsl_render.rs
@@ -33,11 +33,15 @@ fn main() {
 
     let instance = gpu::create_instance();
     let adapter = gpu::create_adapter(&instance, None);
-    let (device, queue) = gpu::create_device(&adapter);
+    let (device, queue) = gpu::create_device_with_features(
+        &adapter,
+        wgpu::Features::TEXTURE_COMPRESSION_BC
+            | wgpu::Features::TEXTURE_COMPRESSION_BC_SLICED_3D,
+    );
 
     let mut client = gpu::create_tsv_client(&device);
     client
-        .init_image(&name, width, height, gpu::TSV_FORMAT, true)
+        .init_image(&name, width, height, 1, gpu::TSV_FORMAT, gpu::ImgType::D2, true)
         .expect("init tsv image");
 
     let fence = gpu::create_fence(&client);
@@ -114,7 +118,7 @@ fn main() {
         };
 
         if dirty {
-            refresh_textures(renderer.uniforms(), &device, &mut client, fence);
+            gpu::refresh_textures(renderer.uniforms(), &device, &mut client, fence);
             renderer.render(&device, &queue);
             device.poll(wgpu::PollType::wait_indefinitely()).unwrap();
 
@@ -228,54 +232,4 @@ fn configure_sampler(
     Ok(())
 }
 
-/// Refresh all texture inputs from TSV shared images.
-fn refresh_textures(
-    cache: &mut wgsl_view::uniform::UniformCache,
-    device: &wgpu::Device,
-    client: &mut texture_share_vk_client::VkClient,
-    fence: vk::Fence,
-) {
-    let mut needs_rebind = false;
-
-    for slot in cache.texture_slots_mut() {
-        let tsv_name = match slot.tsv_name() {
-            Some(n) => n.to_string(),
-            None => continue,
-        };
-
-        if !slot.tsv_registered() {
-            if let Err(e) = client.find_image(&tsv_name, true) {
-                log::debug!("tsv find '{tsv_name}': {e}");
-                continue;
-            }
-
-            match client.find_image_data(&tsv_name, true) {
-                Ok(Some((_lock, data))) => {
-                    if slot.resize(device, data.width, data.height) {
-                        needs_rebind = true;
-                    }
-                }
-                _ => continue,
-            }
-
-            slot.set_tsv_registered();
-        }
 
-        let raw = unsafe { gpu::raw_image(slot.texture()) };
-        match client.recv_image(
-            &tsv_name,
-            raw,
-            vk::ImageLayout::UNDEFINED,
-            vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
-            fence,
-        ) {
-            Ok(Some(())) => log::trace!("texture '{}' updated from '{tsv_name}'", slot.name()),
-            Ok(None) => {}
-            Err(e) => log::warn!("recv_image '{tsv_name}': {e}"),
-        }
-    }
-
-    if needs_rebind {
-        cache.rebuild_bind_group(device);
-    }
-}