cocount.f90 Source File

! SPDX-License-Identifier: GPL-3.0-or-later

! Copyright (C) 2025  Marco Origlia

! This program is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.

! This program is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.

! You should have received a copy of the GNU General Public License
! along with this program.  If not, see <https://www.gnu.org/licenses/>.
module cocount
    !! author: Marco Origlia
    !! license: GPL-3.0-or-later
    !!
    !! Provides a cross-image counter based on Coarrays
    implicit none
    private

    integer :: N
    !! Number of ranges
    integer, allocatable :: cranges(:)
    !! Counter ranges
    integer, allocatable :: partial_products(:)
    !! Partial products of ranges of counters varying faster
    !! than than each index, i.e., of the counters
    !! with lower index
    integer :: top
    !! Total count
    integer :: current[*]
    !! Current count, only first image has the valid value
    integer :: im_first
    !! Index of the reference image, with the valid `current` value

    public :: initialize, next
contains

    subroutine dealloc
        if (allocated(cranges)) deallocate(cranges)
        if (allocated(partial_products)) deallocate(partial_products)
    end subroutine dealloc

    subroutine alloc
        call dealloc

        allocate(cranges(N))
        allocate(partial_products(N))
    end subroutine alloc

    subroutine initialize(ranges)
        !! Set the ranges for the counter
        integer, intent(in) :: ranges(:)
        !! Set of ranges:
        !! size is the number of counters,
        !! counter `i` will run from 0 to `ranges(i)`

        integer :: i

        if (any(ranges <= 0)) stop "Ranges must be positive"

        N = size(ranges)
        call alloc

        cranges = ranges
        partial_products(1) = 1
        do i = 2, N
            partial_products(i) = partial_products(i-1) * ranges(i-1)
        end do

        top = partial_products(N) * ranges(N)
        current = 0

        im_first = lcobound(current, 1)
    end subroutine initialize

    function next(next_count) result(valid)
        !! Provide the next value for all counters
        integer, intent(out) :: next_count(N)
        !! Counters values
        logical :: valid
        !! True as long as there is no counter overflow

        integer :: current_copy, i

        critical
            current_copy = current[im_first]
            current[im_first] = current[im_first] + 1
        end critical

        valid = current_copy < top

        if (.not. valid) return

        do i = 1, N
            next_count(i) = mod(current_copy / partial_products(i), cranges(i))
        end do
    end function next
end module cocount