bound_check_worker.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./range.hpp"
 
#include <cstdint>
#include <stdexcept>
#include <sstream>
 
namespace nda::detail {
 
  // Check the bounds when accessing single elements or slices of an array/view.
  struct bound_check_worker {
    // Shape of the array/view.
    long const *lengths{};
 
    // Error code to store the positions of the arguments which are out of bounds.
    uint32_t error_code = 0;
 
    // Number of dimensions that are covered by a given nda::ellipsis.
    int ellipsis_loss = 0;
 
    // Current dimension to be checked.
    int N = 0;
 
    // Check if the given index is within the bounds of the array/view.
    void check_current_dim(long idx) {
      if ((idx < 0) or (idx >= lengths[N])) { error_code += 1ul << N; }
      ++N;
    }
 
    // Check if the given nda::range is within the bounds of the array/view.
    void check_current_dim(range const &r) {
      if (r.size() > 0) {
        auto first_idx = r.first();
        auto last_idx  = first_idx + (r.size() - 1) * r.step();
        if (first_idx < 0 or first_idx >= lengths[N] or last_idx < 0 or last_idx >= lengths[N]) error_code += 1ul << N;
      }
      ++N;
    }
 
    // Check the bounds when an nda::range::all_t is encountered (no need to check anything).
    void check_current_dim(range::all_t) { ++N; }
 
    // Check the bounds when an nda::ellipsis is encountered (no need to check anything).
    void check_current_dim(ellipsis) { N += ellipsis_loss + 1; }
 
    // Accumulate an error message for the current dimension and index.
    void accumulate_error_msg(std::stringstream &fs, long idx) {
      if (error_code & (1ull << N)) fs << "Argument " << N << " = " << idx << " is not within [0," << lengths[N] << "[.\n";
      N++;
    }
 
    // Accumulate an error message for the current dimension and nda::range.
    void accumulate_error_msg(std::stringstream &fs, range const &r) {
      if (error_code & (1ull << N)) fs << "Argument " << N << " = " << r << " is not within [0," << lengths[N] << "[.\n";
      ++N;
    }
 
    // Accumulate an error message for the current dimension and nda::range::all_t.
    void accumulate_error_msg(std::stringstream &, range::all_t) { ++N; }
 
    // Accumulate an error message for the current dimension and nda::ellipsis.
    void accumulate_error_msg(std::stringstream &, ellipsis) { N += ellipsis_loss + 1; }
  };
 
} // namespace nda::detail
 
namespace nda {
 
  template <typename... Args>
  void assert_in_bounds(int rank, long const *lengths, Args const &...args) {
    // initialize the bounds checker
    detail::bound_check_worker w{lengths};
 
    // number of dimensions that are covered by an nda::ellipsis
    w.ellipsis_loss = rank - sizeof...(Args);
 
    // check the bounds on each argument/index
    (w.check_current_dim(args), ...);
 
    // if no error, stop here
    if (!w.error_code) return;
 
    // accumulate error message and throw
    w.N = 0;
    std::stringstream fs;
    (w.accumulate_error_msg(fs, args), ...);
    throw std::runtime_error("Index/Range out of bounds:\n" + fs.str());
  }
 
} // namespace nda