Fix SGD image desc indexing for 3D UASTC HDR 6x6i textures

lib/src/basis_transcode.cpp

@@ -571,11 +571,11 @@ ktxTexture2_transcodeLzEtc1s(ktxTexture2* This,
     // Temporary invariant value
     uint32_t layersFaces = This->numLayers * This->numFaces;
     firstImages[0] = 0;
-    for (uint32_t level = 1; level <= This->numLevels; level++) {
+    for (uint32_t level = 0; level < This->numLevels; level++) {
         // NOTA BENE: numFaces * depth is only reasonable because they can't
         // both be > 1. I.e there are no 3d cubemaps.
-        firstImages[level] = firstImages[level - 1]
-                           + layersFaces * MAX(This->baseDepth >> (level - 1), 1);
+        firstImages[level + 1] = firstImages[level]
+                           + layersFaces * MAX(This->baseDepth >> level, 1);
     }
     uint32_t& imageCount = firstImages[This->numLevels];
 
@@ -928,6 +928,30 @@ transcodeUastcHDR6x6_intermediate(ktxTexture2* This, alpha_content_e alphaConten
         reinterpret_cast<ktxUASTCHDR6x6IntermediateImageDesc*>(This->_private->_supercompressionGlobalData);
     const uint64_t totalImageDescs = This->_private->_sgdByteLength / sizeof(ktxUASTCHDR6x6IntermediateImageDesc);
 
+    // The image descriptions are stored in level order, level 0 first, with
+    // each level contributing numLayers * numFaces * depth(level) images (see
+    // the writer in basis_encode.cpp). level * levelImageCount only equals
+    // the index of a level's first description while every level has the
+    // same image count; for 3D textures depth halves with each level, so the
+    // first-image index of each level must be accumulated, as transcodeEtc1s
+    // does with its firstImages table.
+    std::vector<uint64_t> firstImages(This->numLevels + 1);
+    firstImages[0] = 0;
+    for (uint32_t l = 0; l < This->numLevels; l++) {
+        firstImages[l + 1] = firstImages[l]
+                             + (uint64_t)This->numLayers * This->numFaces
+                               * MAX(This->baseDepth >> l, 1);
+    }
+
+    // firstImages[numLevels] has the total image count for the texture's
+    // dimensions so a descriptor table whose size does not match exactly is
+    // corrupt; reject it before processing any level.
+    if (This->_private->_sgdByteLength
+            % sizeof(ktxUASTCHDR6x6IntermediateImageDesc) != 0
+        || firstImages[This->numLevels] != totalImageDescs) {
+        return KTX_FILE_DATA_ERROR;
+    }
+
     for (ktx_int32_t level = This->numLevels - 1; level >= 0; level--) {
         ktx_uint32_t depth;
         uint64_t writeOffset = levelOffsetWrite;
@@ -973,8 +997,8 @@ transcodeUastcHDR6x6_intermediate(ktxTexture2* This, alpha_content_e alphaConten
             // See comment before same lines in transcodeEtc1s.
             if (++stateIndex == xcoderStates.size()) stateIndex = 0;
 
-            // Compute the start index into the image seek table.
-            const uint32_t sgdImageDescIndex = (level * levelImageCount) + image;
+            // Compute the index into the image seek table.
+            const uint64_t sgdImageDescIndex = firstImages[level] + image;
 
             // Sanity check the SGD image desc index
             if (sgdImageDescIndex >= totalImageDescs) {

tests/transcodetests/transcodetests.cc

@@ -20,12 +20,15 @@ extern "C" {
   #include "memstream.h"
 }
 #include "platform_utils.h"
+#include "vkformat_enum.h"
 #include "gtest/gtest.h"
 
 #include <version>
 #include <vector>
 #include <cstring>
 #include <filesystem>
+#include <algorithm>
+#include <memory>
 #if defined(__cpp_lib_format)
   #include <format>
 #else
@@ -218,6 +221,100 @@ TEST_P(TextureCombinationsTest, Basic) {
     FormatFeature format = get<1>(GetParam());
     test_texture_set(ts,format);
 }
+
+//////////////////////////////
+// UASTC HDR 6x6 intermediate SGD image description indexing
+//////////////////////////////
+
+// The image description table in the supercompression global data is written
+// in level order, level 0 first, each level contributing
+// numLayers * numFaces * depth(level) descriptions. For 3D textures depth
+// halves with each level, so indexing the table with
+// level * levelImageCount selected descriptions belonging to other levels
+// and transcoding failed. The 2D case covers the constant-image-count path.
+
+static ktx_uint16_t
+floatToHalf(float value) {
+    ktx_uint32_t bits;
+    std::memcpy(&bits, &value, sizeof(bits));
+    const ktx_uint32_t sign = (bits >> 16) & 0x8000u;
+    const ktx_int32_t exponent = (ktx_int32_t)((bits >> 23) & 0xFFu) - 127 + 15;
+    const ktx_uint32_t mantissa = (bits >> 13) & 0x3FFu;
+    if (exponent <= 0)
+        return (ktx_uint16_t)sign;
+    if (exponent >= 31)
+        return (ktx_uint16_t)(sign | 0x7C00u);
+    return (ktx_uint16_t)(sign | ((ktx_uint32_t)exponent << 10) | mantissa);
+}
+
+static void
+roundTripUastcHdr6x6i(ktx_uint32_t numDimensions) {
+    ktxTextureCreateInfo createInfo = {};
+    createInfo.vkFormat = VK_FORMAT_R16G16B16A16_SFLOAT;
+    createInfo.baseWidth = 24;
+    createInfo.baseHeight = 24;
+    createInfo.baseDepth = numDimensions == 3 ? 8 : 1;
+    createInfo.numDimensions = numDimensions;
+    createInfo.numLevels = 4;
+    createInfo.numLayers = 1;
+    createInfo.numFaces = 1;
+    createInfo.isArray = KTX_FALSE;
+    createInfo.generateMipmaps = KTX_FALSE;
+
+    ktxTexture2* texture = nullptr;
+    KTX_error_code result = ktxTexture2_Create(&createInfo,
+                                              KTX_TEXTURE_CREATE_ALLOC_STORAGE,
+                                              &texture);
+    ASSERT_EQ(result, KTX_SUCCESS) << ktxErrorString(result);
+    std::unique_ptr<ktxTexture2, void(*)(ktxTexture2*)> texture_raii(
+        texture, [](ktxTexture2* t) { ktxTexture_Destroy(ktxTexture(t)); });
+
+    for (ktx_uint32_t level = 0; level < createInfo.numLevels; level++) {
+        const ktx_uint32_t width = std::max(1u, createInfo.baseWidth >> level);
+        const ktx_uint32_t height = std::max(1u, createInfo.baseHeight >> level);
+        const ktx_uint32_t depth = std::max(1u, createInfo.baseDepth >> level);
+        for (ktx_uint32_t slice = 0; slice < depth; slice++) {
+            std::vector<ktx_uint16_t> pixels((size_t)width * height * 4);
+            for (ktx_uint32_t y = 0; y < height; y++) {
+                for (ktx_uint32_t x = 0; x < width; x++) {
+                    const size_t i = ((size_t)y * width + x) * 4;
+                    pixels[i + 0] = floatToHalf(0.1f + 2.0f * x / width + level);
+                    pixels[i + 1] = floatToHalf(0.2f + 1.5f * y / height + slice);
+                    pixels[i + 2] = floatToHalf(0.4f + 0.5f * level);
+                    pixels[i + 3] = floatToHalf(1.0f);
+                }
+            }
+            result = ktxTexture_SetImageFromMemory(
+                         ktxTexture(texture), level, 0, slice,
+                         reinterpret_cast<const ktx_uint8_t*>(pixels.data()),
+                         pixels.size() * sizeof(ktx_uint16_t));
+            ASSERT_EQ(result, KTX_SUCCESS) << ktxErrorString(result);
+        }
+    }
+
+    ktxBasisParams cparams = {};
+    cparams.structSize = sizeof(cparams);
+    cparams.threadCount = 1;
+    cparams.codec = KTX_BASIS_CODEC_UASTC_HDR_6x6_INTERMEDIATE;
+    result = ktxTexture2_CompressBasisEx(texture, &cparams);
+    ASSERT_EQ(result, KTX_SUCCESS) << ktxErrorString(result);
+    ASSERT_EQ(texture->supercompressionScheme, KTX_SS_UASTC_HDR_6x6_INTERMEDIATE);
+
+    result = ktxTexture2_TranscodeBasis(texture, KTX_TTF_ASTC_HDR_6x6_RGBA, 0);
+    ASSERT_EQ(result, KTX_SUCCESS) << ktxErrorString(result);
+    EXPECT_EQ(texture->vkFormat,
+              static_cast<ktx_uint32_t>(VK_FORMAT_ASTC_6x6_SFLOAT_BLOCK));
+    EXPECT_EQ(texture->supercompressionScheme, KTX_SS_NONE);
+    EXPECT_NE(texture->pData, nullptr);
+}
+
+TEST(TranscodeUastcHdr6x6i, RoundTrip2D) {
+    roundTripUastcHdr6x6i(2);
+}
+
+TEST(TranscodeUastcHdr6x6i, RoundTrip3DMipLevels) {
+    roundTripUastcHdr6x6i(3);
+}
 }  // namespace
 
 int main(int argc, char **argv) {