! gfortran -g -Wall -Wextra -fimplicit-none -fcheck=all -fsanitize=address -c threefry.f90
module ThreefryModule
    use iso_fortran_env, only: int64
    implicit none

    integer, parameter :: NW_THREEFRY = 2, &
                          N_ROUNDS_THREEFRY = 20
    integer(int64), parameter :: C240 = int (z'1BD11BDAA9FC1A22', int64), &
                                 MASK = huge(0_int64) - 1_int64, &
                                 R_THREEFRY (0:*) = [16, 42, 12, 31, 16, 32, 24, 21]

    type, public :: threefry_key
        integer(int64), dimension(0:NW_THREEFRY) :: k
    end type threefry_key

    type, public :: threefry_result
        integer(int64), dimension(0:NW_THREEFRY-1) :: v
    end type threefry_result

    interface threefry
       procedure :: threefry
       procedure :: threefry_e
    end interface threefry

    contains

    pure function key_schedule_threefry(K, s) result(ksi)
        type(threefry_key), intent (in) :: K
        integer, intent (in) :: s
        integer (int64) :: ksi(0:1)

        ksi(0:1) = [ iand (K%k  (mod (s,     NW_THREEFRY + 1)),      MASK), &
                     iand ((K%k (mod (s + 1, NW_THREEFRY + 1)) + s), MASK) ]
    end function key_schedule_threefry

    pure function mix (e, r) result(v)
        integer (int64), intent (in) :: e(0:1), r
        integer (int64), dimension (0:1) :: v

        v(0:1) = [ ieor(e(0), ishftc(e(1), r)), ieor (e(1), ishftc(e(0), 64 - r)) ]
    end function mix

    pure function threefry (p, K0) result (v)
        integer(int64), intent(in) :: p(0:)
        type(threefry_key), intent(in) :: K0
        integer(int64) :: v(0:1)
        type(threefry_key) :: K

        integer :: r
        integer(int64), dimension (0:1) :: e, ksi

        K%k = [K0%k(0), K0%k(1), ieor(C240, ieor(K0%k(0), K0%k(0)))]
        v = p

        ! Perform Threefish rounds
        do r = 0, N_ROUNDS_THREEFRY - 1
            if (mod(r, 4) == 0) then
                ksi = key_schedule_threefry(K, r / 4)
                e = [iand(v(0) + ksi(0), MASK), iand(v(1) + ksi(1), MASK)]
            else
                e = v
            end if
            v = mix (e, R_THREEFRY (mod (r, 8)))
        end do

        ! Final key injection
        ksi = key_schedule_threefry(K, N_ROUNDS_THREEFRY / 4)
        v = iand (v + ksi, MASK)
    end function threefry

    elemental function threefry_e (p, K0) result (v)
        type(threefry_result), intent(in) :: p
        type(threefry_key), intent(in) :: K0
        type(threefry_result) :: v
        type(threefry_key) :: K

        integer :: r
        integer(int64), dimension (0:1) :: e, ksi

        K%k = [K0%k(0), K0%k(1), ieor(C240, ieor(K0%k(0), K0%k(0)))]
        v = p

        ! Perform Threefish rounds
        do r = 0, N_ROUNDS_THREEFRY - 1
            if (mod(r, 4) == 0) then
                ksi = key_schedule_threefry(K, r / 4)
                e = [iand(v%v(0) + ksi(0), MASK), iand(v%v(1) + ksi(1), MASK)]
            else
                e = v%v
            end if
            v%v = mix (e, R_THREEFRY (mod (r, 8)))
        end do

        ! Final key injection
        ksi = key_schedule_threefry(K, N_ROUNDS_THREEFRY / 4)
        v%v = iand (v%v + ksi, MASK)
    end function threefry_e

end module ThreefryModule