load.hpp

// Copyright 2025 Katherine Whitlock
//
// 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
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// 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.

 
#pragma once
#include <ranges>
#include "argon/argon_full.hpp"
#include "argon/helpers/size.hpp"
#include "arm_simd/helpers/vec128.hpp"
 
#ifdef __ARM_FEATURE_MVE
#define simd mve
#else
#define simd neon
#endif
 
namespace argon::vectorize {

template <typename ScalarType>
struct load : std::ranges::view_interface<load<ScalarType>> {
  using intrinsic_type = simd::Vec128_t<ScalarType>;  
  static constexpr size_t lanes =
      sizeof(intrinsic_type) / sizeof(ScalarType);  
 
 public:
  struct LoadIterator {
    using iterator_category = std::bidirectional_iterator_tag;  
    using difference_type = std::ptrdiff_t;                     
    using value_type = helpers::ArgonFor_t<intrinsic_type>;     
    using reference_type = value_type&;
 
    LoadIterator() = default;

    LoadIterator(const ScalarType* ptr) : ptr_{ptr} {}

    Argon<ScalarType> operator*() const { return value_type::Load(ptr_); }

    LoadIterator& operator+=(int n) {
      ptr_ += n * lanes;
      return *this;
    }

    LoadIterator& operator-=(int n) {
      ptr_ -= n * lanes;
      return *this;
    }

    LoadIterator operator++(int) {
      LoadIterator tmp = *this;
      ++(*this);
      return tmp;
    }

    LoadIterator& operator++() {
      ptr_ += lanes;
      return *this;
    }

    LoadIterator operator--(int) {
      LoadIterator tmp = *this;
      --(*this);
      return tmp;
    }

    LoadIterator& operator--() {
      ptr_ -= lanes;
      return *this;
    }

    LoadIterator operator+(int n) const {
      LoadIterator it = *this;
      it += n;
      return it;
    }

    LoadIterator operator-(int n) const {
      LoadIterator it = *this;
      it -= n;
      return it;
    }

    LoadIterator& operator[](int n) const { return *(*this + n); }

    difference_type operator-(const LoadIterator& other) const { return ptr_ - other.ptr_; }

    friend LoadIterator operator+(const int n, const LoadIterator& it) { return it + n; }

    friend bool operator==(const LoadIterator& a, const LoadIterator& b) { return a.ptr_ == b.ptr_; }

    friend bool operator==(const LoadIterator& a, const ScalarType* ptr) { return a.ptr_ == ptr; }

    friend bool operator!=(const LoadIterator& a, const LoadIterator& b) { return a.ptr_ != b.ptr_; }

    friend bool operator!=(const LoadIterator& a, const ScalarType* ptr) { return a.ptr_ != ptr; }

    friend auto operator<=>(const LoadIterator& a, const LoadIterator& b) { return a.ptr_ <=> b.ptr_; }
 
   private:
    const ScalarType* ptr_;
  };
  static_assert(std::sized_sentinel_for<LoadIterator, LoadIterator>);
  static_assert(std::bidirectional_iterator<LoadIterator>);
  static_assert(std::input_iterator<LoadIterator>);
 
  using iterator = LoadIterator;       
  using sentinel = const ScalarType*;  

  iterator begin() { return start_; }

  const ScalarType* end() { return start_ + (size_ * lanes); }

  size_t size() const { return size_; }

  template <std::ranges::contiguous_range R>
  load(R&& r)
      : start_{&*std::ranges::begin(r)}, size_{helpers::vectorizeable_size<ScalarType>(std::ranges::size(r)) / lanes} {}

  load(const std::span<ScalarType> span)
      : start_{span.data()}, size_{helpers::vectorizeable_size<ScalarType>(span.size()) / lanes} {}
 
 private:
  const ScalarType* start_;
  size_t size_;
};
 
static_assert(std::ranges::range<load<int32_t>>);
static_assert(std::ranges::view<load<int32_t>>);
static_assert(std::movable<load<int32_t>>);
static_assert(std::ranges::viewable_range<load<int32_t>>);
 
template <std::ranges::contiguous_range R>
load(R&& r) -> load<std::remove_cv_t<std::ranges::range_value_t<R>>>;
 
template <typename ScalarType>
load(const std::span<ScalarType>) -> load<std::remove_cv_t<ScalarType>>;
 
}  // namespace argon::vectorize
#undef simd