make aesni and ars operate on public types

Author: jw3126

src/aesni.jl

@@ -201,16 +201,14 @@ copy(src::AESNI4x) = copyto!(AESNI4x(), src)
 ==(r1::AESNI4x, r2::AESNI4x) = unsafe_compare(r1, r2, UInt128, 2) &&
     r1.key == r2.key && r1.p == r2.p
 
-function get_key(o::Union{AESNI1x, AESNI4x})::NTuple{11, __m128i}
+function get_key__m128i(o::Union{AESNI1x, AESNI4x})::NTuple{11, __m128i}
     k = o.key
     (k.key1,k.key2,k.key3,k.key4,k.key5,k.key6,k.key7,k.key8,k.key9,k.key10,k.key11)
 end
-function get_ctr(o::AESNI4x)::Tuple{__m128i}
-    (o.ctr1,)
-end
-function get_ctr(o::AESNI1x)::Tuple{__m128i}
-    (o.ctr,)
-end
+get_ctr__m128i(o::AESNI4x)::Tuple{__m128i} = (o.ctr1,)
+get_ctr__m128i(o::AESNI1x)::Tuple{__m128i} = (o.ctr,)
+get_key(o::Union{AESNI1x, AESNI4x})::NTuple{11,UInt128} = map(UInt128, get_key__m128i(o))
+get_ctr(o::Union{AESNI1x, AESNI4x})::Tuple{UInt128} = map(UInt128, get_ctr__m128i(o))
 
 @inline function aesni(key::NTuple{11,__m128i}, ctr::Tuple{__m128i})::Tuple{__m128i}
     key1, key2, key3, key4, key5, key6, key7, key8, key9, key10, key11 = key
@@ -229,12 +227,25 @@ end
     (x,)
 end
 
+"""
+    aesni(key::NTuple{11,UInt128}, ctr::Tuple{UInt128})::Tuple{UInt128}
+
+Functional variant of [`AESNI1x`](@ref) and [`AESNI4x`](@ref). 
+This function if free of mutability and side effects.
+"""
+@inline function aesni(key::NTuple{11,UInt128}, ctr::Tuple{UInt128})::Tuple{UInt128}
+    k = map(__m128i, key)
+    c = map(__m128i, ctr)
+    map(UInt128,aesni(k,c))
+end
+
+
 @inline function random123_r(r::AESNI1x)
-    r.x = only(aesni(get_key(r), get_ctr(r)))
+    r.x = only(aesni(get_key__m128i(r), get_ctr__m128i(r)))
     (UInt128(r.x),)
 end
 
 @inline function random123_r(r::AESNI4x)
-    r.x = only(aesni(get_key(r), get_ctr(r)))
+    r.x = only(aesni(get_key__m128i(r), get_ctr__m128i(r)))
     split_uint(UInt128(r.x), UInt32)
 end

src/ars.jl

@@ -98,7 +98,7 @@ copy(src::ARS4x{R}) where R = ARS4x{R}(src.x, src.ctr1, src.key, src.p)
 
 ==(r1::ARS4x{R}, r2::ARS4x{R}) where R = unsafe_compare(r1, r2, UInt128, 3) && r1.p ≡ r2.p
 
-function expr_ars1xm128i(expr_key::Expr, expr_ctr::Expr, R)
+function expr_ars1xm128i(expr_key, expr_ctr, R)
     @assert R isa Int && 1 ≤ R ≤ 10
     rounds = [quote
         kk += kweyl
@@ -130,14 +130,25 @@ end
     expr_ars1xm128i(expr_key, expr_ctr, R)
 end
 
-@generated function ars(key, ctr, ::Val{R}) where {R}
-    expr_key = :(only(key))
-    expr_ctr = :(only(ctr))
-    expr_ars1xm128i(expr_key, expr_ctr, R)
+@generated function ars(key::Tuple{__m128i}, ctr::Tuple{__m128i}, ::Val{R})::Tuple{__m128i} where {R}
+    :(($(expr_ars1xm128i(:(only(key)), :(only(ctr)), R)),))
 end
-get_key(r::Union{ARS1x, ARS4x})::Tuple{__m128i} = (r.key,)
-get_ctr(r::ARS1x)::Tuple{__m128i} = (r.ctr,)
-get_ctr(r::ARS4x)::Tuple{__m128i} = (r.ctr1,)
+
+"""
+    ars(key::Tuple{UInt128}, ctr::Tuple{UInt128}, rounds::Val{R})::Tuple{UInt128} where {R}
+
+Functional variant of [`ARS1x`](@ref) and [`ARS4x`](@ref). 
+This function if free of mutability and side effects.
+"""
+function ars(key::Tuple{UInt128}, ctr::Tuple{UInt128}, rounds::Val{R})::Tuple{UInt128} where {R}
+    k = map(__m128i, key)
+    c = map(__m128i, ctr)
+    map(UInt128,ars(k,c,rounds))
+end
+
+get_key(r::Union{ARS1x, ARS4x}) = (UInt128(r.key),)
+get_ctr(r::ARS1x) = (UInt128(r.ctr),)
+get_ctr(r::ARS4x) = (UInt128(r.ctr1),)
 
 @inline function random123_r(r::ARS1x{R}) where R
     r.x = ars1xm128i(r)

test/runtests.jl

@@ -36,14 +36,23 @@ using Printf: @printf
         @test isbitstype(typeof(ctr))
         @test key isa Tuple
         @test ctr isa Tuple
+        @test eltype(key) <: Union{UInt32, UInt64, UInt128}
+        @test eltype(ctr) <: Union{UInt32, UInt64, UInt128}
         val1 = @inferred alg(key, ctr, options...)
         val2 = @inferred alg(key, ctr, options...)
         @test val1 === val2
         @test val1 isa Tuple
         @test isbitstype(typeof(val1))
+        @test eltype(val1) <: Union{UInt32, UInt64, UInt128}
     end
 end
 @testset "functional consistency" begin
+    threefry = Random123.threefry
+    philox = Random123.philox
+    aesni = Random123.aesni
+    ars = Random123.ars
+    get_key = Random123.get_key
+    get_ctr = Random123.get_ctr
     for T in [UInt32, UInt64]
         for (rng, alg, option) in [
                 (Threefry2x(T, (T(123), T(456))), threefry, Val(20)),
@@ -87,6 +96,19 @@ end
         @test x8 === y8
         @test x9 === y9
     end
+
+    rng = ARS1x(1)
+    @test (rand(rng, UInt128),) === ars(get_key(rng), get_ctr(rng), Val(7))
+    @test (rand(rng, UInt128),) === ars(get_key(rng), get_ctr(rng), Val(7))
+    @test (rand(rng, UInt128),) === ars(get_key(rng), get_ctr(rng), Val(7))
+    @test (rand(rng, UInt128),) === ars(get_key(rng), get_ctr(rng), Val(7))
+
+    rng = AESNI1x(1)
+    @test (rand(rng, UInt128),) === aesni(get_key(rng), get_ctr(rng))
+    @test (rand(rng, UInt128),) === aesni(get_key(rng), get_ctr(rng))
+    @test (rand(rng, UInt128),) === aesni(get_key(rng), get_ctr(rng))
+    @test (rand(rng, UInt128),) === aesni(get_key(rng), get_ctr(rng))
+
 end