From bec752361c21bd22726b3c6a2c84f9191063e6c0 Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Thu, 10 Jun 2021 01:49:43 +0200
Subject: [PATCH] Index OK

---
 src/templates_front/templator_front.org | 83 ++++++++++++++++++++-----
 tests/test_f.f90                        | 47 ++++++++++----
 tools/generator_tools.py                |  9 ++-
 3 files changed, 110 insertions(+), 29 deletions(-)

diff --git a/src/templates_front/templator_front.org b/src/templates_front/templator_front.org
index cdd63ba..9b8ac23 100644
--- a/src/templates_front/templator_front.org
+++ b/src/templates_front/templator_front.org
@@ -1129,6 +1129,7 @@ trexio_read_$group_dset$_64 (trexio_t* const file, $group_dset_dtype_double$* co
 
   /* Error handling for this call is added by the generator */
   rc = trexio_read_$group_dset_dim$_64(file, &($group_dset_dim$));
+  if (rc != TREXIO_SUCCESS) return rc;
 
   if ($group_dset_dim$ == 0L) return TREXIO_INVALID_NUM;
 
@@ -1137,22 +1138,36 @@ trexio_read_$group_dset$_64 (trexio_t* const file, $group_dset_dtype_double$* co
 
   assert(file->back_end < TREXIO_INVALID_BACK_END);
 
+  rc = TREXIO_FAILURE;
   switch (file->back_end) {
-
+    
   case TREXIO_TEXT:
-    return trexio_text_read_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_text_read_$group_dset$(file, $group_dset$, rank, dims);
     break;
 
   case TREXIO_HDF5:
-    return trexio_hdf5_read_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_hdf5_read_$group_dset$(file, $group_dset$, rank, dims);
     break;
 /*
   case TREXIO_JSON:
-    return trexio_json_read_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_json_read_$group_dset$(file, $group_dset$, rank, dims);
     break;
 ,*/
   }
-  return TREXIO_FAILURE;
+  if (rc != TREXIO_SUCCESS) return rc;
+  
+  /* Handle index type */
+  if ($is_index$) {
+    uint64_t dim_size = 1;
+    for (uint32_t i=0; i<rank; ++i){
+        dim_size *= dims[i];
+    }
+    for (uint64_t i=0; i<dim_size; ++i){
+       $group_dset$[i] += ($group_dset_dtype_single$) 1;
+    }
+  }
+  
+  return TREXIO_SUCCESS;
 }
      #+end_src
 
@@ -1170,30 +1185,56 @@ trexio_write_$group_dset$_64 (trexio_t* const file, const $group_dset_dtype_doub
 
   /* Error handling for this call is added by the generator */
   rc = trexio_read_$group_dset_dim$_64(file, &($group_dset_dim$));
+  if (rc != TREXIO_SUCCESS) return rc;
 
   if ($group_dset_dim$ == 0L) return TREXIO_INVALID_NUM;
 
   uint32_t rank = $group_dset_rank$;
   uint64_t dims[$group_dset_rank$] = {$group_dset_dim_list$};
 
+  $group_dset_dtype_double$* $group_dset$_p = $group_dset$;
+
+  /* Handle index type */
+  if ($is_index$) {
+
+    uint64_t dim_size = 1;
+    for (uint32_t i=0; i<rank; ++i){
+        dim_size *= dims[i];
+    }
+
+    $group_dset$_p = CALLOC(dim_size, $group_dset_dtype_double$);
+    if ($group_dset$_p == NULL) return TREXIO_ALLOCATION_FAILED;
+
+    for (uint64_t i=0; i<dim_size; ++i){
+       $group_dset$_p[i] = $group_dset$_p[i] - ($group_dset_dtype_single$) 1;
+    }
+  }
+
   assert(file->back_end < TREXIO_INVALID_BACK_END);
 
+  rc = TREXIO_FAILURE;
   switch (file->back_end) {
 
   case TREXIO_TEXT:
-    return trexio_text_write_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_text_write_$group_dset$(file, $group_dset$_p, rank, dims);
     break;
 
   case TREXIO_HDF5:
-    return trexio_hdf5_write_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_hdf5_write_$group_dset$(file, $group_dset$_p, rank, dims);
     break;
 /*
   case TREXIO_JSON:
-    return trexio_json_write_$group_dset$(file, $group_dset$, rank, dims);
+    rc = trexio_json_write_$group_dset$(file, $group_dset$_p, rank, dims);
     break;
 ,*/
   }
-  return TREXIO_FAILURE;
+
+  /* Handle index type */
+  if ($is_index$) {
+    FREE($group_dset$_p);
+  }
+
+  return rc;
 }
      #+end_src
 
@@ -1217,7 +1258,7 @@ trexio_read_$group_dset$_32 (trexio_t* const file, $group_dset_dtype_single$* co
   uint64_t dims[$group_dset_rank$] = {$group_dset_dim_list$};
 
   uint64_t dim_size = 1;
-  for (unsigned int i=0; i<rank; ++i){
+  for (uint32_t i=0; i<rank; ++i){
     dim_size *= dims[i];
   }
 
@@ -1249,8 +1290,14 @@ trexio_read_$group_dset$_32 (trexio_t* const file, $group_dset_dtype_single$* co
     return rc;
   }
 
-  for (uint64_t i=0; i<dim_size; ++i){
-    $group_dset$[i] = ($group_dset_dtype_single$) $group_dset$_64[i];
+  if ($is_index$) {
+    for (uint64_t i=0; i<dim_size; ++i){
+      $group_dset$[i] = ($group_dset_dtype_single$) $group_dset$_64[i] + ($group_dset_dtype_single$) 1;
+    }
+  } else {
+    for (uint64_t i=0; i<dim_size; ++i){
+      $group_dset$[i] = ($group_dset_dtype_single$) $group_dset$_64[i];
+    }
   }
 
   FREE($group_dset$_64);
@@ -1279,7 +1326,7 @@ trexio_write_$group_dset$_32 (trexio_t* const file, const $group_dset_dtype_sing
   uint64_t dims[$group_dset_rank$] = {$group_dset_dim_list$};
 
   uint64_t dim_size = 1;
-  for (unsigned int i=0; i<rank; ++i){
+  for (uint32_t i=0; i<rank; ++i){
     dim_size *= dims[i];
   }
 
@@ -1287,8 +1334,14 @@ trexio_write_$group_dset$_32 (trexio_t* const file, const $group_dset_dtype_sing
   if ($group_dset$_64 == NULL) return TREXIO_ALLOCATION_FAILED;
 
   /* A type conversion from single precision to double reqired since back end only accepts 64-bit data */
-  for (uint64_t i=0; i<dim_size; ++i){
-    $group_dset$_64[i] = ($group_dset_dtype_double$) $group_dset$[i];
+  if ($is_index$) {
+    for (uint64_t i=0; i<dim_size; ++i){
+      $group_dset$_64[i] = ($group_dset_dtype_double$) $group_dset$[i] - ($group_dset_dtype_double$) 1;
+    }
+  } else {
+    for (uint64_t i=0; i<dim_size; ++i){
+      $group_dset$_64[i] = ($group_dset_dtype_double$) $group_dset$[i];
+    }
   }
 
   assert(file->back_end < TREXIO_INVALID_BACK_END);
diff --git a/tests/test_f.f90 b/tests/test_f.f90
index 9ceef27..dcf1526 100644
--- a/tests/test_f.f90
+++ b/tests/test_f.f90
@@ -35,6 +35,7 @@ subroutine test_write(file_name, back_end)
   character*(128) :: str
 
 
+  integer :: basis_nucleus_index(12)
   double precision :: charge(12)
   double precision :: coord(3,12)
 
@@ -55,6 +56,8 @@ subroutine test_write(file_name, back_end)
                        0.00000000d0,  2.47304151d0 ,  0.00000000d0 /), &
                        shape(coord) )
 
+  basis_nucleus_index = (/ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 /)
+
 ! ================= START OF TEST ===================== !
 
   trex_file = trexio_open(file_name, 'w', back_end)
@@ -63,7 +66,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_HAS_NOT) then
     write(*,*) 'SUCCESS HAS NOT 1'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -72,7 +75,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_HAS_NOT) then
     write(*,*) 'SUCCESS HAS NOT 2'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -82,7 +85,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS WRITE NUM'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -91,7 +94,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS WRITE CHARGE'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -100,7 +103,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS WRITE COORD'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -109,7 +112,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS HAS 1'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -118,7 +121,16 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS HAS 2'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
+    print *, trim(str)
+    call exit(1)
+  endif
+
+  rc = trexio_write_basis_nucleus_index(trex_file, basis_nucleus_index)
+  if (rc == TREXIO_SUCCESS) then
+    write(*,*) 'SUCCESS WRITE INDEX'
+  else
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -127,7 +139,7 @@ subroutine test_write(file_name, back_end)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS CLOSE'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -172,6 +184,7 @@ subroutine test_read(file_name, back_end)
   integer :: rc = 1
   integer :: num, num_read
 
+  integer :: basis_nucleus_index(12)
   double precision :: charge(12)
   double precision :: coord(3,12)
 
@@ -188,7 +201,7 @@ subroutine test_read(file_name, back_end)
   if (rc == TREXIO_SUCCESS .and. num_read == num) then
     write(*,*) 'SUCCESS READ NUM'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
@@ -199,7 +212,7 @@ subroutine test_read(file_name, back_end)
   if (rc == TREXIO_SUCCESS .and. (dabs(charge(11) - 1.d0) < 1.0D-8) ) then
     write(*,*) 'SUCCESS READ CHARGE'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(-1)
   endif
@@ -209,16 +222,26 @@ subroutine test_read(file_name, back_end)
   if (rc == TREXIO_SUCCESS .and. (dabs(coord(2,1) - 1.39250319d0) < 1.0D-8) ) then
     write(*,*) 'SUCCESS READ COORD'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(-1)
   endif
 
+  rc = trexio_read_basis_nucleus_index(trex_file, basis_nucleus_index)
+  if (rc == TREXIO_SUCCESS .and. (basis_nucleus_index(12) == 12) ) then
+    write(*,*) 'SUCCESS READ INDEX'
+  else
+    call trexio_string_of_error(rc,str)
+    print *, trim(str)
+    call exit(-1)
+  endif
+
+
   rc = trexio_close(trex_file)
   if (rc == TREXIO_SUCCESS) then
     write(*,*) 'SUCCESS CLOSE'
   else
-    call trexio_string_of_error(TREXIO_READONLY,str)
+    call trexio_string_of_error(rc,str)
     print *, trim(str)
     call exit(1)
   endif
diff --git a/tools/generator_tools.py b/tools/generator_tools.py
index 03adfd9..d9a496e 100644
--- a/tools/generator_tools.py
+++ b/tools/generator_tools.py
@@ -100,7 +100,7 @@ def recursive_populate_file(fname: str, paths: dict, detailed_source: dict) -> N
     fname_new = join('populated',f'pop_{fname}')
     templ_path = get_template_path(fname, paths)
 
-    triggers = ['group_dset_dtype', 'group_dset_h5_dtype', 'default_prec',
+    triggers = ['group_dset_dtype', 'group_dset_h5_dtype', 'default_prec', 'is_index', 
                 'group_dset_f_dtype_default', 'group_dset_f_dtype_double', 'group_dset_f_dtype_single', 
                 'group_dset_dtype_default', 'group_dset_dtype_double', 'group_dset_dtype_single', 
                 'group_dset_rank', 'group_dset_dim_list', 'group_dset_f_dims',
@@ -472,7 +472,7 @@ def split_dset_dict_detailed (datasets: dict) -> tuple:
             default_prec   = '64'
             group_dset_std_dtype_out = '24.16e'
             group_dset_std_dtype_in = 'lf'
-        elif v[0] == 'int':
+        elif v[0] in ['int', 'index']:
             datatype = 'int64_t'
             group_dset_h5_dtype = 'native_int64'
             group_dset_f_dtype_default= 'integer(4)'
@@ -490,6 +490,11 @@ def split_dset_dict_detailed (datasets: dict) -> tuple:
         
         # add the dset name for templates
         tmp_dict['group_dset'] = k
+        # add flag to detect index types
+        if 'index' == v[0]:
+            tmp_dict['is_index'] = 'file->one_based'
+        else:
+            tmp_dict['is_index'] = 'false'
         # add the datatypes for templates
         tmp_dict['dtype'] = datatype
         tmp_dict['group_dset_dtype'] = datatype