From abba3a6ddf07fabc4e1dccc194c43d41f2b99ba7 Mon Sep 17 00:00:00 2001
From: shahramn <shahram.najm@ecmwf.int>
Date: Tue, 26 Nov 2024 13:05:42 +0000
Subject: [PATCH] Element encoding for doubles

---
 src/accessor/grib_accessor_class_element.cc | 47 +++++++++++++++++++++
 src/accessor/grib_accessor_class_element.h  |  1 +
 tests/grib_element.sh                       | 17 +++++++-
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/accessor/grib_accessor_class_element.cc b/src/accessor/grib_accessor_class_element.cc
index ac361dac0..218c246bf 100644
--- a/src/accessor/grib_accessor_class_element.cc
+++ b/src/accessor/grib_accessor_class_element.cc
@@ -123,6 +123,53 @@ int grib_accessor_element_t::pack_long(const long* val, size_t* len)
     return ret;
 }
 
+int grib_accessor_element_t::pack_double(const double* v, size_t* len)
+{
+    int ret               = 0;
+    size_t size           = 0;
+    double* ar            = NULL;
+    const grib_context* c = context_;
+    grib_handle* hand     = grib_handle_of_accessor(this);
+    long index            = element_;
+
+    if (*len < 1) {
+        ret = GRIB_ARRAY_TOO_SMALL;
+        return ret;
+    }
+
+    if ((ret = grib_get_size(hand, array_, &size)) != GRIB_SUCCESS)
+        return ret;
+
+    ar = (double*)grib_context_malloc_clear(c, size * sizeof(double));
+    if (!ar) {
+        grib_context_log(c, GRIB_LOG_ERROR, "Error allocating %zu bytes", size * sizeof(double));
+        return GRIB_OUT_OF_MEMORY;
+    }
+
+    if ((ret = grib_get_double_array_internal(hand, array_, ar, &size)) != GRIB_SUCCESS)
+        return ret;
+
+    // An index of -x means the xth item from the end of the list, so ar[-1] means the last item in ar
+    if (index < 0) {
+        index = size + index;
+    }
+
+    if ((ret = check_element_index(__func__, array_, index, size)) != GRIB_SUCCESS) {
+        goto the_end;
+    }
+
+    Assert(index >= 0);
+    Assert(index < size);
+    ar[index] = *v;
+
+    if ((ret = grib_set_double_array_internal(hand, array_, ar, size)) != GRIB_SUCCESS)
+        goto the_end;
+
+the_end:
+    grib_context_free(c, ar);
+    return ret;
+}
+
 int grib_accessor_element_t::unpack_double(double* val, size_t* len)
 {
     int ret                 = 0;
diff --git a/src/accessor/grib_accessor_class_element.h b/src/accessor/grib_accessor_class_element.h
index 0a93b282d..a19ae8b54 100644
--- a/src/accessor/grib_accessor_class_element.h
+++ b/src/accessor/grib_accessor_class_element.h
@@ -21,6 +21,7 @@ class grib_accessor_element_t : public grib_accessor_long_t
     int pack_long(const long* val, size_t* len) override;
     int unpack_double(double* val, size_t* len) override;
     int unpack_long(long* val, size_t* len) override;
+    int pack_double(const double* val, size_t* len) override;
     void init(const long, grib_arguments*) override;
 
 private:
diff --git a/tests/grib_element.sh b/tests/grib_element.sh
index 8297e09c8..926d48b6a 100755
--- a/tests/grib_element.sh
+++ b/tests/grib_element.sh
@@ -11,6 +11,7 @@
 . ./include.ctest.sh
 
 label="grib_element_test"
+tempGrib=temp.${label}.grib
 tempRef=temp.${label}.ref
 tempText=temp.${label}.txt
 tempFilt=temp.${label}.filt
@@ -72,6 +73,20 @@ status=$?
 set -e
 [ $status -ne 0 ]
 
+# Encode a single double element in an array
+input=$data_dir/sample.grib2
+cat > $tempFilt <<EOF
+    meta elemN element(values, -1);
+    assert( elemN < 301 );
+    print "before: [elemN:d]";
+    set elemN = 311.0;  # value has to be a double
+    print "after:  [elemN:d]";
+    assert( elemN > 310 );
+    write;
+EOF
+${tools_dir}/grib_filter -o $tempGrib $tempFilt $input
+${tools_dir}/grib_get -p avg $input $tempGrib
+
 
 # Clean up
-rm -f $tempRef $tempText $tempFilt
+rm -f $tempRef $tempText $tempFilt $tempGrib