Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  Zarr driver
       5             :  * Author:   Even Rouault <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2025, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "zarr.h"
      14             : 
      15             : #include "cpl_conv.h"
      16             : 
      17             : #include <limits>
      18             : 
      19             : /************************************************************************/
      20             : /*                       ZarrShuffleCompressor()                        */
      21             : /************************************************************************/
      22             : 
      23           1 : static bool ZarrShuffleCompressor(const void *input_data, size_t input_size,
      24             :                                   void **output_data, size_t *output_size,
      25             :                                   CSLConstList options,
      26             :                                   void * /* compressor_user_data */)
      27             : {
      28             :     // 4 is the default of the shuffle numcodecs:
      29             :     // https://numcodecs.readthedocs.io/en/v0.10.0/shuffle.html
      30           1 :     const int eltSize = atoi(CSLFetchNameValueDef(options, "ELEMENTSIZE", "4"));
      31           1 :     if (eltSize != 2 && eltSize != 4 && eltSize != 8)
      32             :     {
      33           0 :         CPLError(CE_Failure, CPLE_AppDefined,
      34             :                  "Only ELEMENTSIZE=2,4,8 is supported");
      35           0 :         if (output_size)
      36           0 :             *output_size = 0;
      37           0 :         return false;
      38             :     }
      39           1 :     if ((input_size % eltSize) != 0)
      40             :     {
      41           0 :         CPLError(CE_Failure, CPLE_AppDefined,
      42             :                  "input_size should be a multiple of ELEMENTSIZE");
      43           0 :         if (output_size)
      44           0 :             *output_size = 0;
      45           0 :         return false;
      46             :     }
      47           1 :     if (output_data != nullptr && *output_data != nullptr &&
      48           1 :         output_size != nullptr && *output_size != 0)
      49             :     {
      50           1 :         if (*output_size < input_size)
      51             :         {
      52           0 :             CPLError(CE_Failure, CPLE_AppDefined, "Too small output size");
      53           0 :             *output_size = input_size;
      54           0 :             return false;
      55             :         }
      56             : 
      57           1 :         const size_t nElts = input_size / eltSize;
      58             :         // Put at the front of the output buffer all the least significant
      59             :         // bytes of each word, then then 2nd least significant byte, etc.
      60           3 :         for (size_t i = 0; i < nElts; ++i)
      61             :         {
      62           6 :             for (int j = 0; j < eltSize; j++)
      63             :             {
      64           4 :                 (static_cast<uint8_t *>(*output_data))[j * nElts + i] =
      65           4 :                     (static_cast<const uint8_t *>(input_data))[i * eltSize + j];
      66             :             }
      67             :         }
      68             : 
      69           1 :         *output_size = input_size;
      70           1 :         return true;
      71             :     }
      72             : 
      73           0 :     if (output_data == nullptr && output_size != nullptr)
      74             :     {
      75           0 :         *output_size = input_size;
      76           0 :         return true;
      77             :     }
      78             : 
      79           0 :     if (output_data != nullptr && *output_data == nullptr &&
      80             :         output_size != nullptr)
      81             :     {
      82           0 :         *output_data = VSI_MALLOC_VERBOSE(input_size);
      83           0 :         *output_size = input_size;
      84           0 :         if (*output_data == nullptr)
      85           0 :             return false;
      86           0 :         bool ret = ZarrShuffleCompressor(input_data, input_size, output_data,
      87             :                                          output_size, options, nullptr);
      88           0 :         if (!ret)
      89             :         {
      90           0 :             VSIFree(*output_data);
      91           0 :             *output_data = nullptr;
      92             :         }
      93           0 :         return ret;
      94             :     }
      95             : 
      96           0 :     CPLError(CE_Failure, CPLE_AppDefined, "Invalid use of API");
      97           0 :     return false;
      98             : }
      99             : 
     100             : /************************************************************************/
     101             : /*                       ZarrShuffleDecompressor()                      */
     102             : /************************************************************************/
     103             : 
     104           2 : static bool ZarrShuffleDecompressor(const void *input_data, size_t input_size,
     105             :                                     void **output_data, size_t *output_size,
     106             :                                     CSLConstList options,
     107             :                                     void * /* compressor_user_data */)
     108             : {
     109             :     // 4 is the default of the shuffle numcodecs:
     110             :     // https://numcodecs.readthedocs.io/en/v0.10.0/shuffle.html
     111           2 :     const int eltSize = atoi(CSLFetchNameValueDef(options, "ELEMENTSIZE", "4"));
     112           2 :     if (eltSize != 2 && eltSize != 4 && eltSize != 8)
     113             :     {
     114           0 :         CPLError(CE_Failure, CPLE_AppDefined,
     115             :                  "Only ELEMENTSIZE=2,4,8 is supported");
     116           0 :         if (output_size)
     117           0 :             *output_size = 0;
     118           0 :         return false;
     119             :     }
     120           2 :     if ((input_size % eltSize) != 0)
     121             :     {
     122           0 :         CPLError(CE_Failure, CPLE_AppDefined,
     123             :                  "input_size should be a multiple of ELEMENTSIZE");
     124           0 :         if (output_size)
     125           0 :             *output_size = 0;
     126           0 :         return false;
     127             :     }
     128           2 :     if (output_data != nullptr && *output_data != nullptr &&
     129           2 :         output_size != nullptr && *output_size != 0)
     130             :     {
     131           2 :         if (*output_size < input_size)
     132             :         {
     133           0 :             CPLError(CE_Failure, CPLE_AppDefined, "Too small output size");
     134           0 :             *output_size = input_size;
     135           0 :             return false;
     136             :         }
     137             : 
     138             :         // Reverse of what is done in the compressor function.
     139           2 :         const size_t nElts = input_size / eltSize;
     140           6 :         for (size_t i = 0; i < nElts; ++i)
     141             :         {
     142          12 :             for (int j = 0; j < eltSize; j++)
     143             :             {
     144           8 :                 (static_cast<uint8_t *>(*output_data))[i * eltSize + j] =
     145           8 :                     (static_cast<const uint8_t *>(input_data))[j * nElts + i];
     146             :             }
     147             :         }
     148             : 
     149           2 :         *output_size = input_size;
     150           2 :         return true;
     151             :     }
     152             : 
     153           0 :     if (output_data == nullptr && output_size != nullptr)
     154             :     {
     155           0 :         *output_size = input_size;
     156           0 :         return true;
     157             :     }
     158             : 
     159           0 :     if (output_data != nullptr && *output_data == nullptr &&
     160             :         output_size != nullptr)
     161             :     {
     162           0 :         *output_data = VSI_MALLOC_VERBOSE(input_size);
     163           0 :         *output_size = input_size;
     164           0 :         if (*output_data == nullptr)
     165           0 :             return false;
     166           0 :         bool ret = ZarrShuffleDecompressor(input_data, input_size, output_data,
     167             :                                            output_size, options, nullptr);
     168           0 :         if (!ret)
     169             :         {
     170           0 :             VSIFree(*output_data);
     171           0 :             *output_data = nullptr;
     172             :         }
     173           0 :         return ret;
     174             :     }
     175             : 
     176           0 :     CPLError(CE_Failure, CPLE_AppDefined, "Invalid use of API");
     177           0 :     return false;
     178             : }
     179             : 
     180             : /************************************************************************/
     181             : /*                       ZarrGetShuffleCompressor()                     */
     182             : /************************************************************************/
     183             : 
     184           1 : const CPLCompressor *ZarrGetShuffleCompressor()
     185             : {
     186             :     static const CPLCompressor gShuffleCompressor = {
     187             :         /* nStructVersion = */ 1,
     188             :         /* pszId = */ "shuffle",
     189             :         CCT_FILTER,
     190             :         /* papszMetadata = */ nullptr,
     191             :         ZarrShuffleCompressor,
     192             :         /* user_data = */ nullptr};
     193             : 
     194           1 :     return &gShuffleCompressor;
     195             : }
     196             : 
     197             : /************************************************************************/
     198             : /*                     ZarrGetShuffleDecompressor()                     */
     199             : /************************************************************************/
     200             : 
     201           2 : const CPLCompressor *ZarrGetShuffleDecompressor()
     202             : {
     203             :     static const CPLCompressor gShuffleDecompressor = {
     204             :         /* nStructVersion = */ 1,
     205             :         /* pszId = */ "shuffle",
     206             :         CCT_FILTER,
     207             :         /* papszMetadata = */ nullptr,
     208             :         ZarrShuffleDecompressor,
     209             :         /* user_data = */ nullptr};
     210             : 
     211           2 :     return &gShuffleDecompressor;
     212             : }
     213             : 
     214             : /************************************************************************/
     215             : /*                       ZarrQuantizeDecompressor()                     */
     216             : /************************************************************************/
     217             : 
     218           2 : static bool ZarrQuantizeDecompressor(const void *input_data, size_t input_size,
     219             :                                      void **output_data, size_t *output_size,
     220             :                                      CSLConstList options,
     221             :                                      void * /* compressor_user_data */)
     222             : {
     223           2 :     const char *dtype = CSLFetchNameValue(options, "DTYPE");
     224           2 :     if (!dtype)
     225             :     {
     226           0 :         CPLError(CE_Failure, CPLE_AppDefined, "quantize: DTYPE missing");
     227           0 :         if (output_size)
     228           0 :             *output_size = 0;
     229           0 :         return false;
     230             :     }
     231           2 :     if (!EQUAL(dtype, "<f4") && !EQUAL(dtype, "<f8"))
     232             :     {
     233           0 :         CPLError(CE_Failure, CPLE_AppDefined,
     234             :                  "quantize: Only DTYPE=<f4 or <f8 is supported. Not %s.",
     235             :                  dtype);
     236           0 :         if (output_size)
     237           0 :             *output_size = 0;
     238           0 :         return false;
     239             :     }
     240             : 
     241           2 :     const int outputEltSize = EQUAL(dtype, "<f4") ? 4 : 8;
     242           2 :     const GDALDataType eOutDT = EQUAL(dtype, "<f4") ? GDT_Float32 : GDT_Float64;
     243             : 
     244           2 :     const char *astype = CSLFetchNameValue(options, "ASTYPE");
     245           2 :     if (!astype)
     246             :     {
     247           0 :         CPLError(CE_Failure, CPLE_AppDefined, "quantize: ASTYPE missing");
     248           0 :         if (output_size)
     249           0 :             *output_size = 0;
     250           0 :         return false;
     251             :     }
     252           2 :     if (!EQUAL(astype, "<f4") && !EQUAL(astype, "<f8"))
     253             :     {
     254           0 :         CPLError(CE_Failure, CPLE_AppDefined,
     255             :                  "quantize: Only ASTYPE=<f4 or <f8 is supported. Not %s.",
     256             :                  astype);
     257           0 :         if (output_size)
     258           0 :             *output_size = 0;
     259           0 :         return false;
     260             :     }
     261             : 
     262           2 :     const int inputEltSize = EQUAL(astype, "<f4") ? 4 : 8;
     263           2 :     const GDALDataType eInDT = EQUAL(astype, "<f4") ? GDT_Float32 : GDT_Float64;
     264             : 
     265           2 :     if ((input_size % inputEltSize) != 0)
     266             :     {
     267           0 :         CPLError(CE_Failure, CPLE_AppDefined,
     268             :                  "input_size should be a multiple of sizeof(ASTYPE)");
     269           0 :         if (output_size)
     270           0 :             *output_size = 0;
     271           0 :         return false;
     272             :     }
     273             : 
     274           2 :     const size_t nElts = input_size / inputEltSize;
     275           2 :     const uint64_t required_output_size64 =
     276           2 :         static_cast<uint64_t>(nElts) * outputEltSize;
     277             :     if constexpr (SIZEOF_VOIDP < 8)
     278             :     {
     279             :         if (required_output_size64 >= std::numeric_limits<size_t>::max())
     280             :         {
     281             :             CPLError(CE_Failure, CPLE_AppDefined, "Too large input");
     282             :             if (output_size)
     283             :                 *output_size = 0;
     284             :             return false;
     285             :         }
     286             :     }
     287           2 :     const size_t required_output_size =
     288             :         static_cast<size_t>(required_output_size64);
     289             : 
     290           2 :     if (output_data != nullptr && *output_data != nullptr &&
     291           2 :         output_size != nullptr && *output_size != 0)
     292             :     {
     293           2 :         if (*output_size < required_output_size)
     294             :         {
     295           0 :             CPLError(CE_Failure, CPLE_AppDefined, "Too small output size");
     296           0 :             *output_size = required_output_size;
     297           0 :             return false;
     298             :         }
     299             : 
     300             : #ifdef CPL_MSB
     301             :         std::vector<GByte> abyTmp;
     302             :         try
     303             :         {
     304             :             abyTmp.resize(input_size);
     305             :         }
     306             :         catch (const std::exception &)
     307             :         {
     308             :             CPLError(CE_Failure, CPLE_OutOfMemory,
     309             :                      "ZarrQuantizeDecompressor: Out of memory");
     310             :             return false;
     311             :         }
     312             :         memcpy(abyTmp.data(), input_data, input_size);
     313             :         GDALSwapWordsEx(abyTmp.data(), inputEltSize, nElts, inputEltSize);
     314             :         input_data = abyTmp.data();
     315             : #endif
     316             : 
     317           2 :         GDALCopyWords64(input_data, eInDT, inputEltSize, *output_data, eOutDT,
     318             :                         outputEltSize, nElts);
     319             : 
     320             : #ifdef CPL_MSB
     321             :         GDALSwapWordsEx(*output_data, outputEltSize, nElts, outputEltSize);
     322             : #endif
     323             : 
     324           2 :         *output_size = required_output_size;
     325           2 :         return true;
     326             :     }
     327             : 
     328           0 :     if (output_data == nullptr && output_size != nullptr)
     329             :     {
     330           0 :         *output_size = required_output_size;
     331           0 :         return true;
     332             :     }
     333             : 
     334           0 :     if (output_data != nullptr && *output_data == nullptr &&
     335             :         output_size != nullptr)
     336             :     {
     337           0 :         *output_data = VSI_MALLOC_VERBOSE(required_output_size);
     338           0 :         *output_size = required_output_size;
     339           0 :         if (*output_data == nullptr)
     340           0 :             return false;
     341           0 :         bool ret = ZarrQuantizeDecompressor(input_data, input_size, output_data,
     342             :                                             output_size, options, nullptr);
     343           0 :         if (!ret)
     344             :         {
     345           0 :             VSIFree(*output_data);
     346           0 :             *output_data = nullptr;
     347             :         }
     348           0 :         return ret;
     349             :     }
     350             : 
     351           0 :     CPLError(CE_Failure, CPLE_AppDefined, "Invalid use of API");
     352           0 :     return false;
     353             : }
     354             : 
     355             : /************************************************************************/
     356             : /*                     ZarrGetQuantizeDecompressor()                    */
     357             : /************************************************************************/
     358             : 
     359           2 : const CPLCompressor *ZarrGetQuantizeDecompressor()
     360             : {
     361             :     static const CPLCompressor gQuantizeDecompressor = {
     362             :         /* nStructVersion = */ 1,
     363             :         /* pszId = */ "quantize",
     364             :         CCT_FILTER,
     365             :         /* papszMetadata = */ nullptr,
     366             :         ZarrQuantizeDecompressor,
     367             :         /* user_data = */ nullptr};
     368             : 
     369           2 :     return &gQuantizeDecompressor;
     370             : }