#!/usr/bin/env python3

# Copyright 2022 The Clspv Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from string import Template
import os

TEMPLATE_CL=Template("""
// RUN: clspv %target %s -o %t.spv${clspv_additional_arg}
// RUN: spirv-dis -o %t2.spvasm %t.spv
// RUN: FileCheck %s < %t2.spvasm
// RUN: spirv-val --target-env vulkan1.0 %t.spv

// AUTO-GENERATED TEST FILE
// This test was generated by ${test_gen_file}.
// Please modify that file and regenate the tests to make changes.

void kernel order(global int *dst, ${type} ${access} image) {
    *dst = get_image_channel_order(image);
}

void kernel data_type(global int *dst, ${type} ${access} image) {
    *dst = get_image_channel_data_type(image);
}

// CHECK: OpMemberDecorate [[struct:%[^ ]+]] 0 Offset 0
// CHECK: OpMemberDecorate [[pc_struct:%[^ ]+]] 0 Offset 0
// CHECK: OpDecorate [[pc_struct]] Block

// CHECK: [[struct]] = OpTypeStruct %uint
// CHECK: [[pc_struct]] = OpTypeStruct [[struct]]
// CHECK: [[ptr_pc_struct:%[^ ]+]] = OpTypePointer PushConstant [[pc_struct]]
// CHECK: [[pc:%[^ ]+]] = OpVariable [[ptr_pc_struct]] PushConstant

// CHECK: [[order_fct:%[^ ]+]] = OpFunction
// CHECK: [[gep:%[^ ]+]] = OpAccessChain {{.*}} [[pc]] %uint_0 %uint_${order_offset}
// CHECK: [[load:%[^ ]+]] = OpLoad %uint [[gep]]
// CHECK: OpStore {{.*}} [[load]]
// CHECK: OpFunctionEnd

// CHECK: [[data_type_fct:%[^ ]+]] = OpFunction
// CHECK: [[gep:%[^ ]+]] = OpAccessChain {{.*}} [[pc]] %uint_0 %uint_${order_offset}
// CHECK: [[load:%[^ ]+]] = OpLoad %uint [[gep]]
// CHECK: OpStore {{.*}} [[load]]
// CHECK: OpFunctionEnd

// CHECK: [[order_kernel:%[^ ]+]] = OpExtInst %void {{.*}} Kernel [[order_fct]]
// CHECK: ImageArgumentInfoChannelOrderPushConstant [[order_kernel]] %uint_${image_ordinal} %uint_${order_offset} %uint_${size}
// CHECK: [[data_type_kernel:%[^ ]+]] = OpExtInst %void {{.*}} Kernel [[data_type_fct]]
// CHECK: ImageArgumentInfoChannelDataTypePushConstant [[data_type_kernel]] %uint_${image_ordinal} %uint_${data_type_offset} %uint_${size}
""")

TEMPLATE_LL=Template("""
; RUN: clspv-opt %s -o %t.ll --passes=set-image-metadata
; RUN: FileCheck %s < %t.ll

; AUTO-GENERATED TEST FILE
; This test was generated by ${test_gen_file}.
; Please modify that file and regenate the tests to make changes.

target datalayout = "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
target triple = "spir-unknown-unknown"

declare spir_func i32 ${order_function}(ptr addrspace(1) %0)

declare spir_func i32 ${data_type_function}(ptr addrspace(1) %0)

define spir_kernel void @order(ptr addrspace(1) writeonly align 4 %dst, ptr addrspace(1) %image) {
entry:
  %0 = call ptr addrspace(1) @_Z14clspv.resource.0(i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, { [0 x i32] } zeroinitializer)
  %1 = getelementptr { [0 x i32] }, ptr addrspace(1) %0, i32 0, i32 0, i32 0
  %2 = call ptr addrspace(1) @_Z14clspv.resource.1(i32 0, i32 1, i32 6, i32 1, i32 1, i32 0, ${image_type} undef)
  %call = tail call spir_func i32 ${order_function}(ptr addrspace(1) %2)
  store i32 %call, ptr addrspace(1) %1, align 4
  ret void
}

define spir_kernel void @data_type(ptr addrspace(1) writeonly align 4 %dst, ptr addrspace(1) %image) {
entry:
  %0 = call ptr addrspace(1) @_Z14clspv.resource.0(i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, { [0 x i32] } zeroinitializer)
  %1 = getelementptr { [0 x i32] }, ptr addrspace(1) %0, i32 0, i32 0, i32 0
  %2 = call ptr addrspace(1) @_Z14clspv.resource.1(i32 0, i32 1, i32 6, i32 1, i32 1, i32 0, ${image_type} undef)
  %call = tail call spir_func i32 ${data_type_function}(ptr addrspace(1) %2)
  store i32 %call, ptr addrspace(1) %1, align 4
  ret void
}

declare ptr addrspace(1) @_Z14clspv.resource.0(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, { [0 x i32] } %6)

declare ptr addrspace(1) @_Z14clspv.resource.1(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, ${image_type} %6)

; CHECK: @__push_constants = addrspace(9) global %0 zeroinitializer, !push_constants [[pc:![^ ]+]]

; CHECK: define spir_kernel void @order(ptr addrspace(1) writeonly align 4 {{.*}}, ptr addrspace(1) {{.*}}) !push_constants_image_channel [[order_kernel:![^ ]+]]
; CHECK: tail call spir_func i32 ${order_function}(ptr addrspace(1) {{.*}}), !image_getter_push_constant_offset [[call:![^ ]+]]

; CHECK: define spir_kernel void @data_type(ptr addrspace(1) writeonly align 4 {{.*}}, ptr addrspace(1) {{.*}}) !push_constants_image_channel [[data_type_kernel:![^ ]+]]
; CHECK: tail call spir_func i32 ${data_type_function}(ptr addrspace(1) {{.*}}), !image_getter_push_constant_offset [[call:![^ ]+]]

; CHECK: [[pc]] = !{i32 {{.*}}}
; CHECK: [[order_kernel]] = !{i32 [[ordinal:[^ ]+]], i32 [[order_offset:[^ ]+]], i32 [[order_pc:[^ ]+]]}
; CHECK: [[call]] = !{i32 ${offset}}
; CHECK: [[data_type_kernel]] = !{i32 [[ordinal:[^ ]+]], i32 ${offset}, i32 [[data_type_pc:[^ ]+]]}
""")

ACCESS_RO="ro"
ACCESS_WO="wo"
ACCESS_RW="rw"

ACCESS=[ACCESS_RO, ACCESS_WO, ACCESS_RW]
ACCESS_STRING={ACCESS_RO: "read_only", ACCESS_WO: "write_only", ACCESS_RW: "read_write"}
SAMPLED={ACCESS_RO: "1", ACCESS_WO: "2", ACCESS_RW: "2"}
SAMPLED_STRING={ACCESS_RO: ".sampled", ACCESS_WO: "", ACCESS_RW: ""}
CLSPV_ADD_ARG={ACCESS_RO: "", ACCESS_WO: "", ACCESS_RW: " -cl-std=CL2.0 -inline-entry-points"}

TYPES=["image1d", "image2d", "image3d", "image1d_buffer", "image1d_array", "image2d_array"]
DIM={"image1d": "0", "image2d": "1", "image3d": "2", "image1d_buffer": "5", "image1d_array": "0", "image2d_array": "1"}
ARRAYED={"image1d": "0", "image2d": "0", "image3d": "0", "image1d_buffer": "0", "image1d_array": "1", "image2d_array": "1"}


def generate_one_cl(ty, acc):
    template = TEMPLATE_CL.substitute(test_gen_file = os.path.basename(__file__),
                                      type = ty + "_t",
                                      clspv_additional_arg = CLSPV_ADD_ARG[acc],
                                      access = ACCESS_STRING[acc],
                                      size = "4",
                                      image_ordinal = "1",
                                      order_offset = "0",
                                      data_type_offset = "0")
    filename = "get_image_channel_" + ty + "_" + ACCESS_STRING[acc] + ".cl"
    with open(filename, "w") as file:
        file.write(template)

def generate_cl():
    for ty in TYPES:
        for acc in ACCESS:
            generate_one_cl(ty, acc)

def generate_one_ll(ty, acc):
    image_ty_str = "opencl." + ty + "_" + acc + "_t.float" + SAMPLED_STRING[acc]
    image_ty = "target(\"spirv.Image\", " + DIM[ty] + ", 0, " + ARRAYED[ty] + ", 0, " + SAMPLED[acc] + ", 0, 0, 0)"
    template = TEMPLATE_LL.substitute(test_gen_file = os.path.basename(__file__),
                                      image_type = image_ty,
                                      order_function = "@_Z23get_image_channel_order" + str(len(image_ty)) + image_ty_str ,
                                      data_type_function = "@_Z27get_image_channel_data_type" + str(len(image_ty)) + image_ty_str,
                                      offset = "0")
    filename = "get_image_channel_" + ty + "_" + ACCESS_STRING[acc] + ".ll"
    with open(filename, "w") as file:
        file.write(template)

def generate_ll():
    for ty in TYPES:
        for acc in ACCESS:
            generate_one_ll(ty, acc)

generate_ll()
generate_cl()