dear-imgui-rs 0.14.0

//! Thread-safe rendering snapshot.
//!
//! This module provides `Send + Sync` data structures which capture everything a renderer backend
//! needs to render a frame, without retaining any pointers into ImGui-owned memory.

use crate::render::draw_data::{
    DrawData, DrawIdx, DrawList, DrawVert, StandardDrawCallback, classify_standard_draw_callback,
};
use crate::sys;
pub use crate::texture::ManagedTextureId;
use crate::texture::{TextureFormat, TextureId, TextureRect, TextureStatus};
use thiserror::Error;

#[cfg(feature = "multi-viewport")]
const IMGUI_VIEWPORT_DEFAULT_ID: u32 = 0x1111_1111;

/// How a draw command binds its texture.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum TextureBinding {
    /// Legacy texture binding (plain handle).
    Legacy(TextureId),
    /// ImGui-managed texture binding, resolved by `ManagedTextureId`.
    Managed(ManagedTextureId),
}

/// A thread-safe snapshot of everything needed to render a frame.
#[derive(Clone, Debug)]
pub struct FrameSnapshot {
    pub draw: DrawDataSnapshot,
    pub viewports: Vec<ViewportDrawDataSnapshot>,
    pub texture_requests: Vec<TextureRequest>,
}

/// Thread-safe draw data snapshot for one Dear ImGui viewport.
#[derive(Clone, Debug)]
pub struct ViewportDrawDataSnapshot {
    pub viewport_id: crate::Id,
    pub draw: DrawDataSnapshot,
}

/// Options controlling snapshot behavior.
#[derive(Copy, Clone, Debug)]
pub struct SnapshotOptions {
    pub user_callback_policy: UserCallbackPolicy,
    pub capture_texture_requests: bool,
}

impl Default for SnapshotOptions {
    fn default() -> Self {
        Self {
            user_callback_policy: UserCallbackPolicy::Error,
            capture_texture_requests: true,
        }
    }
}

/// Policy for `ImDrawCmd::UserCallback` commands when snapshotting.
#[derive(Copy, Clone, Debug, Default)]
pub enum UserCallbackPolicy {
    /// Return an error when encountering a user callback (default).
    #[default]
    Error,
    /// Drop callback commands from the snapshot.
    Drop,
}

/// Errors that can occur when building a snapshot.
#[derive(Error, Debug)]
pub enum SnapshotError {
    #[error("user callback commands are not supported in the snapshot path")]
    UserCallbackUnsupported,

    #[error("managed texture {id:?} has status {status:?} but no pixel buffer is available")]
    TexturePixelsMissing {
        id: ManagedTextureId,
        status: TextureStatus,
    },

    #[error(
        "managed texture {id:?} has invalid dimensions/format (width={width}, height={height}, bpp={bpp})"
    )]
    TextureInvalidLayout {
        id: ManagedTextureId,
        width: i32,
        height: i32,
        bpp: i32,
    },
}

/// Thread-safe draw data snapshot.
#[derive(Clone, Debug)]
pub struct DrawDataSnapshot {
    pub display_pos: [f32; 2],
    pub display_size: [f32; 2],
    pub framebuffer_scale: [f32; 2],
    pub draw_lists: Vec<DrawListSnapshot>,
}

/// Thread-safe draw list snapshot.
#[derive(Clone, Debug)]
pub struct DrawListSnapshot {
    pub vtx: Vec<DrawVert>,
    pub idx: Vec<DrawIdx>,
    pub commands: Vec<DrawCmdSnapshot>,
}

/// Thread-safe draw command snapshot.
#[derive(Clone, Debug)]
pub enum DrawCmdSnapshot {
    Elements {
        count: usize,
        clip_rect: [f32; 4],
        texture: TextureBinding,
        vtx_offset: usize,
        idx_offset: usize,
    },
    ResetRenderState,
    SetSamplerLinear,
    SetSamplerNearest,
}

/// A thread-safe managed texture request (ImGui 1.92+).
#[derive(Clone, Debug)]
pub struct TextureRequest {
    pub id: ManagedTextureId,
    pub op: TextureOp,
}

/// Feedback produced by the renderer thread, to be applied on the UI thread.
#[derive(Copy, Clone, Debug)]
#[non_exhaustive]
pub struct TextureFeedback {
    pub id: ManagedTextureId,
    pub status: TextureStatus,
    pub tex_id: Option<TextureId>,
    pub backend_user_data: Option<usize>,
}

impl TextureFeedback {
    /// Feedback for a created/updated GPU texture handle.
    pub fn with_tex_id(id: ManagedTextureId, status: TextureStatus, tex_id: TextureId) -> Self {
        Self {
            id,
            status,
            tex_id: Some(tex_id),
            backend_user_data: None,
        }
    }

    /// Feedback for status-only transitions such as `Destroyed`.
    pub fn status(id: ManagedTextureId, status: TextureStatus) -> Self {
        Self {
            id,
            status,
            tex_id: None,
            backend_user_data: None,
        }
    }

    /// Attach renderer-owned backend user data to apply alongside status/TexID.
    #[must_use]
    pub fn backend_user_data(mut self, backend_user_data: usize) -> Self {
        self.backend_user_data = Some(backend_user_data);
        self
    }
}

/// A managed texture operation requested by ImGui.
#[derive(Clone, Debug)]
pub enum TextureOp {
    Create {
        format: TextureFormat,
        width: u32,
        height: u32,
        row_pitch: usize,
        pixels: Vec<u8>,
    },
    Update {
        format: TextureFormat,
        width: u32,
        height: u32,
        rects: Vec<TextureUploadRect>,
    },
    Destroy,
}

/// A tightly-packed pixel upload for a sub-rectangle of a managed texture.
#[derive(Clone, Debug)]
pub struct TextureUploadRect {
    pub rect: TextureRect,
    pub row_pitch: usize,
    pub data: Vec<u8>,
}

impl FrameSnapshot {
    /// Build a `FrameSnapshot` from ImGui draw data.
    ///
    /// This must be called on the UI thread while the `Context` is current.
    pub fn from_draw_data(
        draw_data: &DrawData,
        options: SnapshotOptions,
    ) -> Result<Self, SnapshotError> {
        let draw = snapshot_draw_data(draw_data, options)?;
        let texture_requests = if options.capture_texture_requests {
            snapshot_texture_requests(draw_data)?
        } else {
            Vec::new()
        };
        let viewports = viewport_draw_data_snapshots_from_draw_data(draw_data, &draw);
        Ok(Self {
            draw,
            viewports,
            texture_requests,
        })
    }

    /// Build a `FrameSnapshot` from all current Dear ImGui platform viewport draw data.
    ///
    /// This must be called on the UI thread while the `Context` is current and after
    /// `Context::render()` has produced valid draw data.
    #[cfg(feature = "multi-viewport")]
    pub(crate) fn from_platform_io(
        platform_io: &crate::platform_io::PlatformIo,
        options: SnapshotOptions,
    ) -> Result<Self, SnapshotError> {
        let mut viewports = Vec::new();
        let mut main_draw_index = None;
        for viewport in platform_io.viewports_iter() {
            let Some(draw_data) = viewport.draw_data_ref() else {
                continue;
            };
            let draw_data = draw_data_from_sys(draw_data);
            if !draw_data.valid() {
                continue;
            }
            if main_draw_index.is_none() && is_main_platform_viewport(viewport.id(), draw_data) {
                main_draw_index = Some(viewports.len());
            }
            viewports.push(ViewportDrawDataSnapshot {
                viewport_id: viewport.id(),
                draw: snapshot_draw_data(draw_data, options)?,
            });
        }

        let Some(main_draw_index) = main_draw_index else {
            return Ok(Self {
                draw: DrawDataSnapshot {
                    display_pos: [0.0, 0.0],
                    display_size: [0.0, 0.0],
                    framebuffer_scale: [1.0, 1.0],
                    draw_lists: Vec::new(),
                },
                viewports: Vec::new(),
                texture_requests: Vec::new(),
            });
        };

        let texture_requests = if options.capture_texture_requests {
            let main_draw_data = platform_io
                .viewports_iter()
                .filter_map(|viewport| viewport.draw_data_ref())
                .map(draw_data_from_sys)
                .filter(|draw_data| draw_data.valid())
                .nth(main_draw_index)
                .expect("main viewport draw data exists");
            snapshot_texture_requests(main_draw_data)?
        } else {
            Vec::new()
        };

        Ok(Self {
            draw: viewports[main_draw_index].draw.clone(),
            viewports,
            texture_requests,
        })
    }

    /// Draw data for a specific viewport, if it was captured in this snapshot.
    #[must_use]
    pub fn viewport_draw(&self, viewport_id: crate::Id) -> Option<&DrawDataSnapshot> {
        self.viewports
            .iter()
            .find(|viewport| viewport.viewport_id == viewport_id)
            .map(|viewport| &viewport.draw)
    }
}

fn viewport_draw_data_snapshots_from_draw_data(
    draw_data: &DrawData,
    draw: &DrawDataSnapshot,
) -> Vec<ViewportDrawDataSnapshot> {
    let Some(viewport_id) = owner_viewport_id(draw_data) else {
        return Vec::new();
    };
    vec![ViewportDrawDataSnapshot {
        viewport_id,
        draw: draw.clone(),
    }]
}

fn owner_viewport_id(draw_data: &DrawData) -> Option<crate::Id> {
    let owner_viewport = draw_data.owner_viewport();
    if owner_viewport.is_null() {
        return None;
    }
    let raw = unsafe { (*owner_viewport).ID };
    (raw != 0).then_some(crate::Id::from(raw))
}

#[cfg(feature = "multi-viewport")]
fn draw_data_from_sys(draw_data: &sys::ImDrawData) -> &DrawData {
    unsafe { <DrawData as crate::internal::RawCast<sys::ImDrawData>>::from_raw(draw_data) }
}

#[cfg(feature = "multi-viewport")]
fn is_main_platform_viewport(viewport_id: crate::Id, draw_data: &DrawData) -> bool {
    viewport_id.raw() == IMGUI_VIEWPORT_DEFAULT_ID
        || owner_viewport_id(draw_data)
            .is_some_and(|owner_id| owner_id.raw() == IMGUI_VIEWPORT_DEFAULT_ID)
}

#[cfg(all(test, feature = "multi-viewport"))]
mod tests {
    use super::*;

    fn empty_draw_data(
        viewport: *mut sys::ImGuiViewport,
        display_pos: [f32; 2],
        display_size: [f32; 2],
    ) -> *mut sys::ImDrawData {
        let draw_data = unsafe { sys::ImDrawData_ImDrawData() };
        assert!(!draw_data.is_null());
        unsafe {
            (*draw_data).Valid = true;
            (*draw_data).DisplayPos = sys::ImVec2 {
                x: display_pos[0],
                y: display_pos[1],
            };
            (*draw_data).DisplaySize = sys::ImVec2 {
                x: display_size[0],
                y: display_size[1],
            };
            (*draw_data).FramebufferScale = sys::ImVec2 { x: 1.0, y: 1.0 };
            (*draw_data).OwnerViewport = viewport;
            (*draw_data).Textures = std::ptr::null_mut();
        }
        draw_data
    }

    fn viewport(id: u32, draw_data: *mut sys::ImDrawData) -> *mut sys::ImGuiViewport {
        let viewport = unsafe { sys::ImGuiViewport_ImGuiViewport() };
        assert!(!viewport.is_null());
        unsafe {
            (*viewport).ID = id;
            (*viewport).DrawData = draw_data;
        }
        viewport
    }

    #[test]
    fn platform_io_snapshot_captures_draw_data_per_viewport() {
        let main = viewport(IMGUI_VIEWPORT_DEFAULT_ID, std::ptr::null_mut());
        let secondary = viewport(0x222, std::ptr::null_mut());
        let main_draw = empty_draw_data(main, [0.0, 0.0], [640.0, 360.0]);
        let secondary_draw = empty_draw_data(secondary, [100.0, 50.0], [320.0, 200.0]);
        unsafe {
            (*main).DrawData = main_draw;
            (*secondary).DrawData = secondary_draw;
        }

        let mut viewport_ptrs = [main, secondary];
        let mut raw = sys::ImGuiPlatformIO {
            Viewports: sys::ImVector_ImGuiViewportPtr {
                Size: 2,
                Capacity: 2,
                Data: viewport_ptrs.as_mut_ptr(),
            },
            ..Default::default()
        };
        let platform_io = unsafe {
            crate::platform_io::PlatformIo::from_raw(
                (&mut raw as *mut sys::ImGuiPlatformIO).cast_const(),
            )
        };

        let snapshot = FrameSnapshot::from_platform_io(&platform_io, SnapshotOptions::default())
            .expect("valid viewport draw data should snapshot");

        assert_eq!(snapshot.draw.display_size, [640.0, 360.0]);
        assert_eq!(snapshot.viewports.len(), 2);
        assert_eq!(
            snapshot
                .viewport_draw(crate::Id::from(0x222))
                .expect("secondary viewport should be captured")
                .display_pos,
            [100.0, 50.0]
        );

        unsafe {
            sys::ImDrawData_destroy(main_draw);
            sys::ImDrawData_destroy(secondary_draw);
            sys::ImGuiViewport_destroy(main);
            sys::ImGuiViewport_destroy(secondary);
        }
    }

    #[test]
    fn platform_io_snapshot_uses_default_viewport_as_main_even_when_ordered_later() {
        let secondary = viewport(0x222, std::ptr::null_mut());
        let main = viewport(IMGUI_VIEWPORT_DEFAULT_ID, std::ptr::null_mut());
        let secondary_draw = empty_draw_data(secondary, [100.0, 50.0], [320.0, 200.0]);
        let main_draw = empty_draw_data(main, [0.0, 0.0], [640.0, 360.0]);
        unsafe {
            (*secondary).DrawData = secondary_draw;
            (*main).DrawData = main_draw;
        }

        let mut viewport_ptrs = [secondary, main];
        let mut raw = sys::ImGuiPlatformIO {
            Viewports: sys::ImVector_ImGuiViewportPtr {
                Size: 2,
                Capacity: 2,
                Data: viewport_ptrs.as_mut_ptr(),
            },
            ..Default::default()
        };
        let platform_io = unsafe {
            crate::platform_io::PlatformIo::from_raw(
                (&mut raw as *mut sys::ImGuiPlatformIO).cast_const(),
            )
        };

        let snapshot = FrameSnapshot::from_platform_io(&platform_io, SnapshotOptions::default())
            .expect("valid viewport draw data should snapshot");

        assert_eq!(snapshot.draw.display_size, [640.0, 360.0]);
        assert_eq!(
            snapshot.viewports[0].draw.display_size,
            [320.0, 200.0],
            "captured viewport order should be preserved independently from the compatibility main draw"
        );

        unsafe {
            sys::ImDrawData_destroy(secondary_draw);
            sys::ImDrawData_destroy(main_draw);
            sys::ImGuiViewport_destroy(secondary);
            sys::ImGuiViewport_destroy(main);
        }
    }
}

fn snapshot_draw_data(
    draw_data: &DrawData,
    options: SnapshotOptions,
) -> Result<DrawDataSnapshot, SnapshotError> {
    let mut draw_lists = Vec::with_capacity(draw_data.draw_lists_count());
    for draw_list in draw_data.draw_lists() {
        draw_lists.push(snapshot_draw_list(draw_list, options)?);
    }

    Ok(DrawDataSnapshot {
        display_pos: draw_data.display_pos(),
        display_size: draw_data.display_size(),
        framebuffer_scale: draw_data.framebuffer_scale(),
        draw_lists,
    })
}

fn snapshot_draw_list(
    draw_list: &DrawList,
    options: SnapshotOptions,
) -> Result<DrawListSnapshot, SnapshotError> {
    let vtx = draw_list.vtx_buffer().to_vec();
    let idx = draw_list.idx_buffer().to_vec();

    let mut commands = Vec::new();
    for cmd in unsafe { draw_list.cmd_buffer() } {
        if cmd.UserCallback.is_some() {
            match classify_standard_draw_callback(cmd.UserCallback) {
                Some(StandardDrawCallback::ResetRenderState) => {
                    commands.push(DrawCmdSnapshot::ResetRenderState);
                    continue;
                }
                Some(StandardDrawCallback::SetSamplerLinear) => {
                    commands.push(DrawCmdSnapshot::SetSamplerLinear);
                    continue;
                }
                Some(StandardDrawCallback::SetSamplerNearest) => {
                    commands.push(DrawCmdSnapshot::SetSamplerNearest);
                    continue;
                }
                None => match options.user_callback_policy {
                    UserCallbackPolicy::Error => {
                        return Err(SnapshotError::UserCallbackUnsupported);
                    }
                    UserCallbackPolicy::Drop => continue,
                },
            }
        }

        let texture = snapshot_texture_binding(cmd.TexRef);
        commands.push(DrawCmdSnapshot::Elements {
            count: count_from_u32("DrawCmdSnapshot::Elements::count", cmd.ElemCount),
            clip_rect: [
                cmd.ClipRect.x,
                cmd.ClipRect.y,
                cmd.ClipRect.z,
                cmd.ClipRect.w,
            ],
            texture,
            vtx_offset: count_from_u32("DrawCmdSnapshot::Elements::vtx_offset", cmd.VtxOffset),
            idx_offset: count_from_u32("DrawCmdSnapshot::Elements::idx_offset", cmd.IdxOffset),
        });
    }

    Ok(DrawListSnapshot { vtx, idx, commands })
}

fn count_from_u32(caller: &str, raw: u32) -> usize {
    usize::try_from(raw).unwrap_or_else(|_| panic!("{caller} exceeded usize range"))
}

fn snapshot_texture_binding(tex_ref: sys::ImTextureRef) -> TextureBinding {
    if tex_ref._TexID != 0 {
        return TextureBinding::Legacy(TextureId::from(tex_ref._TexID as u64));
    }

    if !tex_ref._TexData.is_null() {
        let unique_id = unsafe { (*tex_ref._TexData).UniqueID };
        return TextureBinding::Managed(ManagedTextureId::from_raw(unique_id));
    }

    TextureBinding::Legacy(TextureId::null())
}

fn snapshot_texture_requests(draw_data: &DrawData) -> Result<Vec<TextureRequest>, SnapshotError> {
    let mut out = Vec::new();

    for tex in draw_data.textures() {
        let status = tex.status();
        if status == TextureStatus::OK || status == TextureStatus::Destroyed {
            continue;
        }

        let id = tex.unique_id();
        let raw_width = tex.raw_width_i32();
        let raw_height = tex.raw_height_i32();
        let raw_bpp = tex.raw_bytes_per_pixel_i32();
        let Some(width) = u32::try_from(raw_width).ok().filter(|value| *value > 0) else {
            return Err(SnapshotError::TextureInvalidLayout {
                id,
                width: raw_width,
                height: raw_height,
                bpp: raw_bpp,
            });
        };
        let Some(height) = u32::try_from(raw_height).ok().filter(|value| *value > 0) else {
            return Err(SnapshotError::TextureInvalidLayout {
                id,
                width: raw_width,
                height: raw_height,
                bpp: raw_bpp,
            });
        };
        let Some(bpp) = usize::try_from(raw_bpp).ok().filter(|value| *value > 0) else {
            return Err(SnapshotError::TextureInvalidLayout {
                id,
                width: raw_width,
                height: raw_height,
                bpp: raw_bpp,
            });
        };
        let format = tex.format();

        match status {
            TextureStatus::WantCreate => {
                let pixels = tex
                    .pixels()
                    .ok_or(SnapshotError::TexturePixelsMissing { id, status })?;
                let expected = usize::try_from(width)
                    .ok()
                    .and_then(|w| usize::try_from(height).ok().and_then(|h| w.checked_mul(h)))
                    .and_then(|px| px.checked_mul(bpp));
                let Some(expected) = expected else {
                    return Err(SnapshotError::TextureInvalidLayout {
                        id,
                        width: raw_width,
                        height: raw_height,
                        bpp: raw_bpp,
                    });
                };
                if pixels.len() < expected {
                    return Err(SnapshotError::TextureInvalidLayout {
                        id,
                        width: raw_width,
                        height: raw_height,
                        bpp: raw_bpp,
                    });
                }
                out.push(TextureRequest {
                    id,
                    op: TextureOp::Create {
                        format,
                        width,
                        height,
                        row_pitch: usize::try_from(width)
                            .ok()
                            .and_then(|w| w.checked_mul(bpp))
                            .ok_or(SnapshotError::TextureInvalidLayout {
                                id,
                                width: raw_width,
                                height: raw_height,
                                bpp: raw_bpp,
                            })?,
                        pixels: pixels[..expected].to_vec(),
                    },
                });
            }
            TextureStatus::WantUpdates => {
                let pixels = tex
                    .pixels()
                    .ok_or(SnapshotError::TexturePixelsMissing { id, status })?;
                let expected = usize::try_from(width)
                    .ok()
                    .and_then(|w| usize::try_from(height).ok().and_then(|h| w.checked_mul(h)))
                    .and_then(|px| px.checked_mul(bpp));
                let Some(expected) = expected else {
                    return Err(SnapshotError::TextureInvalidLayout {
                        id,
                        width: raw_width,
                        height: raw_height,
                        bpp: raw_bpp,
                    });
                };
                if pixels.len() < expected {
                    return Err(SnapshotError::TextureInvalidLayout {
                        id,
                        width: raw_width,
                        height: raw_height,
                        bpp: raw_bpp,
                    });
                }

                let mut rects: Vec<TextureRect> = tex.updates().collect();
                if rects.is_empty() {
                    let r = tex.update_rect();
                    if r.w != 0 && r.h != 0 {
                        rects.push(r);
                    } else {
                        rects.push(TextureRect {
                            x: 0,
                            y: 0,
                            w: width.min(u16::MAX as u32) as u16,
                            h: height.min(u16::MAX as u32) as u16,
                        });
                    }
                }

                let rects = rects
                    .into_iter()
                    .filter_map(|r| copy_texture_rect(pixels, width, height, bpp, r))
                    .collect::<Vec<_>>();

                out.push(TextureRequest {
                    id,
                    op: TextureOp::Update {
                        format,
                        width,
                        height,
                        rects,
                    },
                });
            }
            TextureStatus::WantDestroy => {
                out.push(TextureRequest {
                    id,
                    op: TextureOp::Destroy,
                });
            }
            TextureStatus::OK | TextureStatus::Destroyed => {}
        }
    }

    Ok(out)
}

fn copy_texture_rect(
    pixels: &[u8],
    width: u32,
    height: u32,
    bpp: usize,
    rect: TextureRect,
) -> Option<TextureUploadRect> {
    let width = usize::try_from(width).ok()?;
    let height = usize::try_from(height).ok()?;
    if width == 0 || height == 0 || bpp == 0 {
        return None;
    }

    let x = usize::from(rect.x);
    let y = usize::from(rect.y);
    let w = usize::from(rect.w);
    let h = usize::from(rect.h);
    if w == 0 || h == 0 {
        return None;
    }

    let x_end = x.saturating_add(w).min(width);
    let y_end = y.saturating_add(h).min(height);
    if x >= x_end || y >= y_end {
        return None;
    }

    let rect_w = x_end - x;
    let rect_h = y_end - y;

    let full_row_pitch = width.checked_mul(bpp)?;
    let rect_row_pitch = rect_w.checked_mul(bpp)?;
    let needed_size = rect_row_pitch.checked_mul(rect_h)?;

    let mut out = vec![0u8; needed_size];
    for row in 0..rect_h {
        let src_row = y.checked_add(row)?;
        let src_off = src_row
            .checked_mul(full_row_pitch)?
            .checked_add(x.checked_mul(bpp)?)?;
        let dst_off = row.checked_mul(rect_row_pitch)?;
        let src_end = src_off.checked_add(rect_row_pitch)?;
        let dst_end = dst_off.checked_add(rect_row_pitch)?;
        if src_end > pixels.len() || dst_end > out.len() {
            return None;
        }
        out[dst_off..dst_end].copy_from_slice(&pixels[src_off..src_end]);
    }

    Some(TextureUploadRect {
        rect: TextureRect {
            x: rect.x,
            y: rect.y,
            w: rect_w.min(u16::MAX as usize) as u16,
            h: rect_h.min(u16::MAX as usize) as u16,
        },
        row_pitch: rect_row_pitch,
        data: out,
    })
}