metal2vulkan 0.1.0

//! Depth sampling and comparison lowering.

use super::*;

/// Lower `air.sample_depth_<dim>.f32`: sample through the existing color-image path and return the
/// first component as the depth scalar. The harness currently seeds captured depth textures as
/// RGBA8_UNORM images, so this intentionally preserves the sampled-image contract instead of
/// introducing true depth-image formats before the cross-host evidence needs them.
pub(in crate::passes) fn lower_sample_depth(
    ctx: &mut Ctx,
    _name: &str,
    res: Option<Word>,
    rty: Option<Word>,
    args: &[Word],
    v4: Word,
) -> Result<Vec<Instruction>, String> {
    let (res, rty) = match (res, rty) {
        (Some(r), Some(t)) => (r, t),
        _ => return Err("air.sample_depth has no result".into()),
    };
    if args.len() < 4 {
        return Err("air.sample_depth missing texture/sampler/coord".into());
    }
    let (mut img, samp, coord) = (resolve_image_value(ctx, args[0]), args[1], args[3]);
    if texture_operand_is_private_pointer(ctx, img) {
        if let Some(sampled_img) = single_sampled_image_for_private_read(ctx, img) {
            img = sampled_img;
        } else {
            return lower_null_texture_result(ctx, res, rty);
        }
    }
    let (dim, arrayed) = ctx
        .image_dims
        .get(&img)
        .copied()
        .unwrap_or((Dim::Dim2D, false));
    let comp = ctx
        .image_comp
        .get(&img)
        .copied()
        .unwrap_or(crate::passes::ImageComp::Float);
    if arrayed {
        return Err("air.sample_depth array textures are unsupported".into());
    }
    if comp != crate::passes::ImageComp::Float {
        return Err("air.sample_depth on non-float texture".into());
    }
    let img_ty = ctx.ty_image(dim, arrayed, comp);
    let si_ty = ctx.ty_sampled_image(img_ty);
    let si = ctx.module.fresh_id();
    let color = ctx.module.fresh_id();
    let depth = ctx.module.fresh_id();
    let mut out = vec![Instruction::new(
        Op::SampledImage,
        Some(si_ty),
        Some(si),
        vec![Operand::IdRef(img), Operand::IdRef(samp)],
    )];
    let coord_for_sample = build_sample_coord(ctx, dim, arrayed, coord, args, &mut out)?;
    push_image_sample(
        ctx,
        &mut out,
        v4,
        color,
        si,
        coord_for_sample,
        None,
        false,
        None,
    );
    out.push(Instruction::new(
        Op::CompositeExtract,
        Some(ctx.ty_float()),
        Some(depth),
        vec![Operand::IdRef(color), Operand::LiteralBit32(0)],
    ));

    let rdef = type_def_of(ctx, rty);
    let is_struct = rdef
        .as_ref()
        .map(|d| d.class.opcode == Op::TypeStruct)
        .unwrap_or(false);
    if is_struct {
        let i8u = ctx.ty_int8();
        let undef8 = ctx.module.fresh_id();
        out.push(Instruction::new(Op::Undef, Some(i8u), Some(undef8), vec![]));
        out.push(Instruction::new(
            Op::CompositeConstruct,
            Some(rty),
            Some(res),
            vec![Operand::IdRef(depth), Operand::IdRef(undef8)],
        ));
    } else {
        out.push(Instruction::new(
            Op::CopyObject,
            Some(rty),
            Some(res),
            vec![Operand::IdRef(depth)],
        ));
    }
    Ok(out)
}

/// Lower `air.sample_compare_depth_<dim>.f32`: sample the RGBA8-backed depth fixture, extract
/// component 0, compare it against the AIR reference value, and return 1.0 or 0.0 in the AIR
/// `{float, i8}` shape. This mirrors `lower_sample_depth`'s current harness contract; it is not a
/// Vulkan Dref/comparison-sampler lowering.
pub(in crate::passes) fn lower_sample_compare_depth(
    ctx: &mut Ctx,
    _name: &str,
    res: Option<Word>,
    rty: Option<Word>,
    args: &[Word],
    v4: Word,
) -> Result<Vec<Instruction>, String> {
    let (res, rty) = match (res, rty) {
        (Some(r), Some(t)) => (r, t),
        _ => return Err("air.sample_compare_depth has no result".into()),
    };
    if args.len() < 5 {
        return Err("air.sample_compare_depth missing texture/sampler/coord/reference".into());
    }
    let (mut img, samp, coord, reference) =
        (resolve_image_value(ctx, args[0]), args[1], args[3], args[4]);
    if texture_operand_is_private_pointer(ctx, img) {
        if let Some(sampled_img) = single_sampled_image_for_private_read(ctx, img) {
            img = sampled_img;
        } else {
            return lower_null_texture_result(ctx, res, rty);
        }
    }
    let (dim, arrayed) = ctx
        .image_dims
        .get(&img)
        .copied()
        .unwrap_or((Dim::Dim2D, false));
    let comp = ctx
        .image_comp
        .get(&img)
        .copied()
        .unwrap_or(crate::passes::ImageComp::Float);
    if arrayed {
        return Err("air.sample_compare_depth array textures are unsupported".into());
    }
    if comp != crate::passes::ImageComp::Float {
        return Err("air.sample_compare_depth on non-float texture".into());
    }
    let img_ty = ctx.ty_image(dim, arrayed, comp);
    let si_ty = ctx.ty_sampled_image(img_ty);
    let float_ty = ctx.ty_float();
    let bool_ty = ctx.ty_bool();
    let si = ctx.module.fresh_id();
    let color = ctx.module.fresh_id();
    let depth = ctx.module.fresh_id();
    let passed = ctx.module.fresh_id();
    let shadow = ctx.module.fresh_id();
    let one = ctx.const_float(1.0);
    let zero = ctx.const_float(0.0);
    let mut out = vec![Instruction::new(
        Op::SampledImage,
        Some(si_ty),
        Some(si),
        vec![Operand::IdRef(img), Operand::IdRef(samp)],
    )];
    let coord_for_sample = build_sample_coord(ctx, dim, arrayed, coord, args, &mut out)?;
    push_image_sample(
        ctx,
        &mut out,
        v4,
        color,
        si,
        coord_for_sample,
        None,
        false,
        None,
    );
    out.push(Instruction::new(
        Op::CompositeExtract,
        Some(float_ty),
        Some(depth),
        vec![Operand::IdRef(color), Operand::LiteralBit32(0)],
    ));
    out.push(Instruction::new(
        Op::FOrdLessThanEqual,
        Some(bool_ty),
        Some(passed),
        vec![Operand::IdRef(depth), Operand::IdRef(reference)],
    ));
    out.push(Instruction::new(
        Op::Select,
        Some(float_ty),
        Some(shadow),
        vec![
            Operand::IdRef(passed),
            Operand::IdRef(one),
            Operand::IdRef(zero),
        ],
    ));

    let rdef = type_def_of(ctx, rty);
    let is_struct = rdef
        .as_ref()
        .map(|d| d.class.opcode == Op::TypeStruct)
        .unwrap_or(false);
    if is_struct {
        let i8u = ctx.ty_int8();
        let undef8 = ctx.module.fresh_id();
        out.push(Instruction::new(Op::Undef, Some(i8u), Some(undef8), vec![]));
        out.push(Instruction::new(
            Op::CompositeConstruct,
            Some(rty),
            Some(res),
            vec![Operand::IdRef(shadow), Operand::IdRef(undef8)],
        ));
    } else {
        out.push(Instruction::new(
            Op::CopyObject,
            Some(rty),
            Some(res),
            vec![Operand::IdRef(shadow)],
        ));
    }
    Ok(out)
}

pub(in crate::passes) fn find_sample_lod(
    ctx: &mut Ctx,
    arrayed: bool,
    args: &[Word],
    out: &mut Vec<Instruction>,
) -> Option<Word> {
    // Ordinary AIR sample forms put coord at arg[2]. Arrayed forms consume arg[3] as the layer, so a
    // scalar float after that is the explicit Metal `level(...)` operand. Non-LOD flags are integers.
    let start = if arrayed { 4 } else { 3 };
    args.get(start..)?
        .iter()
        .find_map(|arg| sample_lod_as_f32(ctx, *arg, out))
}

pub(in crate::passes) fn sample_spatial_dims(dim: Dim) -> Option<usize> {
    match dim {
        Dim::Dim1D => Some(1),
        Dim::Dim2D => Some(2),
        Dim::Dim3D => Some(3),
        _ => None,
    }
}

pub(in crate::passes) fn push_image_sample(
    ctx: &mut Ctx,
    out: &mut Vec<Instruction>,
    result_ty: Word,
    result: Word,
    sampled_image: Word,
    coord: Word,
    explicit_lod: Option<Word>,
    force_lod0: bool,
    const_offset: Option<Word>,
) {
    let lod = explicit_lod.or_else(|| {
        (force_lod0 || matches!(ctx.stage, Stage::Kernel)).then(|| ctx.const_float(0.0))
    });
    let mut operands = vec![Operand::IdRef(sampled_image), Operand::IdRef(coord)];
    let mut image_operands = spirv::ImageOperands::empty();
    if lod.is_some() {
        image_operands |= spirv::ImageOperands::LOD;
    }
    if const_offset.is_some() {
        image_operands |= spirv::ImageOperands::CONST_OFFSET;
    }
    if !image_operands.is_empty() {
        operands.push(Operand::ImageOperands(image_operands));
        if let Some(lod) = lod {
            operands.push(Operand::IdRef(lod));
        }
        if let Some(offset) = const_offset {
            operands.push(Operand::IdRef(offset));
        }
    }
    out.push(Instruction::new(
        if lod.is_some() {
            Op::ImageSampleExplicitLod
        } else {
            Op::ImageSampleImplicitLod
        },
        Some(result_ty),
        Some(result),
        operands,
    ));
}

pub(in crate::passes) fn sample_lod_as_f32(
    ctx: &mut Ctx,
    lod: Word,
    out: &mut Vec<Instruction>,
) -> Option<Word> {
    let ty = value_result_type(ctx, lod)?;
    let def = type_def_of(ctx, ty)?;
    if def.class.opcode != Op::TypeFloat {
        return None;
    }
    let width = match def.operands.first() {
        Some(Operand::LiteralBit32(w)) => *w,
        _ => return None,
    };
    if width == 32 {
        return Some(lod);
    }
    if width == 16 {
        let widened = ctx.module.fresh_id();
        out.push(Instruction::new(
            Op::FConvert,
            Some(ctx.ty_float()),
            Some(widened),
            vec![Operand::IdRef(lod)],
        ));
        return Some(widened);
    }
    None
}