Support of std::array as top level type for a TTree branch

stdarraytoplevel.C

/**
The idea behind this code is to treat std::array as a C array also for top level
TTree branches. This is achieved with template code and works for nested
std::arrays too.
Only the types which can be put in a leaflist are supported, e.g. no std::array
of classes.
The idea is to intervene on two template methods of TTree:
- `Branch` for writing: when the user inputs the address of a std::array, the
code
automatically creates a branch for a C style array of the size of the array via
the
leaflist syntax.
- `SetBranchAddress` for reading: when the user tries to read into a std::array,
the
code reads into the std::array::data() pointer.
 */

//------------------------------------------------------------------------------
// The code which should go to ROOT

// Transform the name of a type to the char which is to be used in the leaflist.
// If the type name is not supported, an exception is thrown.
char GetTypeForLeafList(std::string_view tname) {
  if (tname == "char")
    return 'B';
  if (tname == "unsigned char")
    return 'b';
  if (tname == "short int")
    return 'S';
  if (tname == "unsigned short int")
    return 's';
  if (tname == "int")
    return 'I';
  if (tname == "unsigned int")
    return 'i';
  if (tname == "float")
    return 'F';
  if (tname == "double")
    return 'D';
  if (tname == "Float16_t")
    return 'f';
  if (tname == "Double32_t")
    return 'd';
  if (tname == "Long64_t")
    return 'L';
  if (tname == "ULong64_t")
    return 'l';
  if (tname == "bool")
    return 'O';

  std::string msg(tname);
  msg += ": type not supported in an array by TTree.";
  throw std::runtime_error(msg);
  return ' ';
}

// This template invokes `TTree::Branch` on the TTree instance
// passed for the address `arr` to create a branch called `name`.
template <typename T>
TBranch *BranchHelper(TTree &t, T *arr, const char *name) {
  return t.Branch(name, arr);
}

// Overload of BranchHelper which treats the std::array case
template <typename T, unsigned long N>
TBranch *BranchHelper(TTree &t, std::array<T, N> *arr, const char *name) {
  // We have an std::array as top level branch. We now treat it as
  // if it were a C array of fixed size.

  // We get all the info we need about this std::array
  const auto cl = TClass::GetClass<std::array<T, N>>();
  std::string tname;
  std::array<int, 5> maxIndices;
  int ndim;
  TClassEdit::GetStdArrayProperties(cl->GetName(), tname, maxIndices, ndim);

  // We need to build now the string for the branch
  const auto tchar = GetTypeForLeafList(tname);
  std::string leaflist(name);
  for (auto idim : ROOT::TSeqI(ndim)) {
    leaflist += "[" + std::to_string(maxIndices[idim]) + "]";
  }
  leaflist += "/";
  leaflist += tchar;
  return t.Branch(name, arr->data(), leaflist.c_str());
}

template <typename T>
int SetBranchAddressHelper(TTree &t, const char *name, T *&addr) {
  return t.SetBranchAddress(name, addr);
}

template <typename T, unsigned long N>
const T *GetDataFromArray(const std::array<T, N> &arr) {
  return arr.data();
}

template <typename T, unsigned long N>
int SetBranchAddressHelper(TTree &t, const char *name, std::array<T, N> *arrp) {
  auto dataPtr = (void *)GetDataFromArray(*arrp);
  return t.SetBranchAddress(name, dataPtr);
}

//------------------------------------------------------------------------------
// The code for the test. This code can be easily inserted in a gtest
// RTestArrayFiller is there to fill the data structure of vector<std::array<T,
// N>>, where N can be nesting
// further std::arrays

template <typename ARR, unsigned int I> struct RTestArrayFiller {
  RTestArrayFiller(std::vector<ARR> &, unsigned int) {
    static_assert(I > 5, "Not supported dimension.");
  }
};

template <typename ARR> struct RTestArrayFiller<ARR, 1> {
  RTestArrayFiller(std::vector<ARR> &varr, unsigned int nEntries) {
    ARR init;
    for (auto iEntry : ROOT::TSeqI(nEntries)) {
      for (auto &v : init)
        v = gRandom->Gaus(0, 4);
      varr.emplace_back(init);
    }
  }
};
template <typename ARR> struct RTestArrayFiller<ARR, 2> {
  RTestArrayFiller(std::vector<ARR> &varr, unsigned int nEntries) {
    ARR init;
    for (auto iEntry : ROOT::TSeqI(nEntries)) {
      for (auto &arr0 : init) {
        for (auto &v : arr0) {
          v = gRandom->Gaus(0, 4);
        }
      }
      varr.emplace_back(init);
    }
  }
};

template <typename ARR> struct RTestArrayFiller<ARR, 3> {
  RTestArrayFiller(std::vector<ARR> &varr, unsigned int nEntries) {
    ARR init;
    for (auto iEntry : ROOT::TSeqI(nEntries)) {
      for (auto &arr1 : init) {
        for (auto &arr0 : arr1) {
          for (auto &v : arr0) {
            v = gRandom->Gaus(0, 4);
          }
        }
      }
      varr.emplace_back(init);
    }
  }
};

template <typename ARR> struct RTestArrayFiller<ARR, 4> {
  RTestArrayFiller(std::vector<ARR> &varr, unsigned int nEntries) {
    ARR init;
    for (auto iEntry : ROOT::TSeqI(nEntries)) {
      for (auto &arr2 : init) {
        for (auto &arr1 : arr2) {
          for (auto &arr0 : arr1) {
            for (auto &v : arr0) {
              v = gRandom->Gaus(0, 4);
            }
          }
        }
      }
      varr.emplace_back(init);
    }
  }
};

template <typename ARR> struct RTestArrayFiller<ARR, 5> {
  RTestArrayFiller(std::vector<ARR> &varr, unsigned int nEntries) {
    ARR init;
    for (auto iEntry : ROOT::TSeqI(nEntries)) {
      for (auto &arr3 : init) {
        for (auto &arr2 : arr3) {
          for (auto &arr1 : arr2) {
            for (auto &arr0 : arr1) {
              for (auto &v : arr0) {
                v = gRandom->Gaus(0, 4);
              }
            }
          }
        }
      }
      varr.emplace_back(init);
    }
  }
};

// WORKAROUND TO PRINT VALUES: value printing is broken in 6.18
template <typename ARR, unsigned int NDIM> struct RArrPrinter {
  RArrPrinter(const ARR &arr) {
    std::cout << "{ ";
    const auto size = arr.size();
    for (auto i : ROOT::TSeqI(size - 1)) {
      RArrPrinter<typename ARR::value_type, NDIM - 1> p(arr[i]);
      std::cout << ", ";
    }
    if (size > 0)
      RArrPrinter<typename ARR::value_type, NDIM - 1> p(arr[size - 1]);
    std::cout << " }" << std::endl;
  }
};

template <typename ARR> struct RArrPrinter<ARR, 1> {
  RArrPrinter(const ARR &arr) {
    std::cout << "{ ";
    const auto size = arr.size();
    for (auto i : ROOT::TSeqI(size - 1)) {
      std::cout << arr[i] << " ,";
    }
    if (size > 0)
      std::cout << arr[size - 1];
    std::cout << " }";
  }
};
// END OF WORKAROUND TO PRINT VALUES

template <typename T, unsigned int NDIM> class RStdArrayTester {
public:
  using testArray_t = T;

private:
  static const auto constexpr fFilename = "stdarraytoplevel.root";
  static const auto constexpr fTreename = "tree";
  unsigned int fSeed;
  std::vector<testArray_t> fTestData;

  void Write() {
    testArray_t arr;
    TFile ofile(fFilename, "RECREATE");
    TTree tree(fTreename, fTreename);
    BranchHelper(tree, &arr, "arr");
    for (const auto &testArr : fTestData) {
      arr = testArr;
      tree.Fill();
    }
    tree.Write();
    ofile.Close();
  }

  void Read() {
    TFile ofile(fFilename);
    auto tree = ofile.Get<TTree>(fTreename);
    const auto nentries = tree->GetEntries();
    testArray_t arr;
    SetBranchAddressHelper(*tree, "arr", &arr);
    for (auto ievt : ROOT::TSeqI(nentries)) {
      tree->GetEntry(ievt);
      std::cout << "Event " << ievt << std::endl;
      if (arr != fTestData[ievt]) {
        std::cerr << "Read array is different from the one which was written!"
                  << std::endl;
      }
      // Here value printing!!
      RArrPrinter<T, NDIM> p(arr);
      std::cout << std::endl;
    }
  }

public:
  RStdArrayTester() : fSeed(gRandom->GetSeed()) {
    gRandom->SetSeed(123);
    const auto nEntries = gRandom->Uniform(1, 16);
    fTestData.reserve(nEntries);
    RTestArrayFiller<testArray_t, NDIM>(fTestData, nEntries);
    Write();
    Read();
  }
  ~RStdArrayTester() {
    // Restore the seed
    gRandom->SetSeed(fSeed);
  }
};

void stdarraytoplevel() {
  RStdArrayTester<std::array<float, 4>, 1> t0;
  RStdArrayTester<std::array<std::array<int, 4>, 8>, 2> t1;
  RStdArrayTester<std::array<std::array<std::array<double, 3>, 3>, 4>, 3> t3;
  RStdArrayTester<
      std::array<std::array<std::array<std::array<float, 4>, 3>, 3>, 4>, 4>
      t4;
}