Rust使用DX11进行截图并保存

写的比较糙，基本上是对照着C++转过来的。先放着，回头记不起来怎么用再回来看吧。

use image::{ImageBuffer, Rgba};
use windows::{
    core::Interface,
    Win32::Graphics::{
        Direct3D::{D3D_DRIVER_TYPE_HARDWARE, D3D_FEATURE_LEVEL},
        Direct3D11::{
            D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D,
            D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_BGRA_SUPPORT,
            D3D11_CREATE_DEVICE_SINGLETHREADED, D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ,
            D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
        },
        Dxgi::{
            Common::{DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC},
            IDXGIDevice, IDXGIOutput1, IDXGIResource, IDXGIResource1, DXGI_OUTDUPL_FRAME_INFO,
        },
    },
};

fn create_d3d_device() -> Result<(ID3D11Device, ID3D11DeviceContext), String> {
    let mut dervie: Option<ID3D11Device> = None;
    let mut feature_level = D3D_FEATURE_LEVEL::default();
    let mut context = None;
    let result = unsafe {
        D3D11CreateDevice(
            None,
            D3D_DRIVER_TYPE_HARDWARE,
            None,
            D3D11_CREATE_DEVICE_BGRA_SUPPORT | D3D11_CREATE_DEVICE_SINGLETHREADED,
            None,
            D3D11_SDK_VERSION,
            Some(&mut dervie),
            Some(&mut feature_level),
            Some(&mut context),
        )
    };
    match result {
        Ok(()) => Ok((dervie.unwrap(), context.unwrap())),
        Err(err_msg) => Err(format!("Failed to create D3D11 device: {}", err_msg)),
    }
}

fn main() {
    // 1. 创建 D3D11 设备
    let (device, context) = create_d3d_device().unwrap();

    let dx_device = device.cast::<IDXGIDevice>().unwrap();
    // 适配器
    let adapter = unsafe {
        match dx_device.GetAdapter() {
            Ok(adapter) => adapter,
            Err(err) => {
                println!("Failed to get adapter: {:?}", err);
                panic!();
            }
        }
    };
    // 输出
    let output = unsafe {
        match adapter.EnumOutputs(0) {
            Ok(output) => output,
            Err(err) => {
                println!("Failed to get output: {:?}", err);
                panic!();
            }
        }
    };
    // 输出1
    let output1 = output.cast::<IDXGIOutput1>().unwrap();

    let result = unsafe { output1.DuplicateOutput(&device) };
    let dupl_output = match result {
        Ok(dupl_output) => dupl_output,
        Err(err) => {
            println!("Failed to duplicate output: {:?}", err);
            panic!();
        }
    };

    // 获取帧信息
    let mut frame_info = DXGI_OUTDUPL_FRAME_INFO::default();
    let mut desktop_image: Option<IDXGIResource> = None;
    loop {
        let result =
            unsafe { dupl_output.AcquireNextFrame(500, &mut frame_info, &mut desktop_image) };
        if result.is_err() {
            println!("Failed to acquire next frame: {:?}", result.err());
            return;
        }
        if frame_info.LastPresentTime != 0 {
            break;
        }
        // 释放桌面图像
        unsafe { dupl_output.ReleaseFrame().unwrap() };
    }

    frame_info.LastPresentTime;
    if desktop_image.is_none() {
        panic!("desktop_image is none!")
    }

    let desktop_image = desktop_image.unwrap();
    let desktop_image = desktop_image.cast::<IDXGIResource1>().unwrap();

    let start_time = std::time::Instant::now();

    save_frame_to_file(&desktop_image, &device, &context).unwrap();

    let elapsed = start_time.elapsed();
    println!("Time elapsed: {:?}", elapsed);

    // 释放资源
    unsafe { dupl_output.ReleaseFrame().unwrap() };
}

fn save_frame_to_file(
    desktop_image: &IDXGIResource1,
    device: &ID3D11Device,
    context: &ID3D11DeviceContext,
) -> Result<(), Box<dyn std::error::Error>> {
    let desktop_image1 = desktop_image.cast::<IDXGIResource1>()?;
    let desktop_texture: ID3D11Texture2D = desktop_image1.cast()?;
    let mut desc = D3D11_TEXTURE2D_DESC::default();
    unsafe { desktop_texture.GetDesc(&mut desc) };

    // 创建用于CPU读取的纹理
    desc.Usage = D3D11_USAGE_STAGING;
    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ.0 as u32;
    desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    desc.SampleDesc = DXGI_SAMPLE_DESC {
        Count: 1,
        Quality: 0,
    };
    desc.BindFlags = 0;
    desc.MiscFlags = 0;
    desc.MipLevels = 1;
    desc.ArraySize = 1;

    let mut staging_texture = None;
    let result = unsafe { device.CreateTexture2D(&desc, None, Some(&mut staging_texture)) };
    if let Err(err) = result {
        println!("Failed to create staging texture: {:?}", err);
        return Err("Failed to create staging texture".into());
    }

    let staging_texture = match staging_texture {
        Some(staging_texture) => staging_texture,
        None => {
            println!("Failed to create staging texture");
            return Err("Failed to create staging texture".into());
        }
    };
    // 将桌面图像复制到CPU可访问的纹理
    unsafe {
        context.CopyResource(&staging_texture, &desktop_texture);
    }
    let mut mapped_resource: D3D11_MAPPED_SUBRESOURCE = Default::default();
    unsafe {
        context.Map(
            &staging_texture,
            0,
            D3D11_MAP_READ,
            0,
            Some(&mut mapped_resource),
        )?
    };

    // 将图像数据复制到 buffer 中
    let width = desc.Width as usize;
    let height = desc.Height as usize;
    let row_pitch = mapped_resource.RowPitch as usize;
    let data = unsafe {
        std::slice::from_raw_parts(mapped_resource.pData as *const u8, row_pitch * height)
    };

    // 创建 image crate 的缓冲区并保存
    let mut img_buffer = ImageBuffer::<Rgba<u8>, Vec<u8>>::new(desc.Width, desc.Height);
    for y in 0..height {
        for x in 0..width {
            let offset = y * row_pitch + x * 4;
            // 交换 R 和 B 通道
            let pixel = Rgba([
                data[offset + 2],
                data[offset + 1],
                data[offset],
                data[offset + 3],
            ]);
            img_buffer.put_pixel(x as u32, y as u32, pixel);
        }
    }
    img_buffer.save("output.png")?;
    unsafe { context.Unmap(&staging_texture, 0) };
    Ok(())
}

[package]
edition = "2021"
name = "test_image"
version = "0.1.0"

[dependencies]
image = "0.25.5"

[dependencies.windows]
features = [
  "Win32_UI_WindowsAndMessaging",
  "Win32_Graphics_Gdi",
  "Win32_Storage_Xps",
  "Win32_Graphics_Direct3D",
  "Win32_Graphics_Direct3D11",
  "Win32_Graphics_Dxgi",
  "Win32_Graphics_Dxgi_Common",
  "Win32_Graphics_Imaging",
]
version = "0.59.0"