Faster mersenne-twister with specialized-arrays

basis/random/mersenne-twister/mersenne-twister.factor

@@ -2,48 +2,54 @@
 ! See http://factorcode.org/license.txt for BSD license.
 ! mersenne twister based on 
 ! http://www.math.sci.hiroshima-u.ac.jp/~m-mat/MT/MT2002/CODES/mt19937ar.c
-USING: kernel math namespaces sequences system init
-accessors math.ranges random circular math.bitwise
-combinators specialized-arrays.uint ;
+USING: kernel math namespaces sequences sequences.private system
+init accessors math.ranges random math.bitwise combinators
+specialized-arrays.uint fry ;
 IN: random.mersenne-twister
 
 <PRIVATE
 
-TUPLE: mersenne-twister seq i ;
+TUPLE: mersenne-twister { seq uint-array } { i fixnum } ;
 
 : mt-n 624 ; inline
 : mt-m 397 ; inline
 : mt-a HEX: 9908b0df ; inline
 
+: wrap-nth ( n seq -- obj )
+    [ length mod ] keep nth-unsafe ; inline
+
+: set-wrap-nth ( obj n seq -- )
+    [ length mod ] keep set-nth-unsafe ; inline
+
 : calculate-y ( n seq -- y )
-    [ nth 31 mask-bit ]
-    [ [ 1+ ] [ nth ] bi* 31 bits ] 2bi bitor ; inline
+    [ wrap-nth 31 mask-bit ]
+    [ [ 1+ ] [ wrap-nth ] bi* 31 bits ] 2bi bitor ; inline
 
 : (mt-generate) ( n seq -- next-mt )
     [
         calculate-y
         [ 2/ ] [ odd? mt-a 0 ? ] bi bitxor
     ] [
-        [ mt-m + ] [ nth ] bi*
-    ] 2bi bitxor ;
+        [ mt-m + ] [ wrap-nth ] bi*
+    ] 2bi bitxor ; inline
 
 : mt-generate ( mt -- )
     [
-        mt-n swap seq>> [
-            [ (mt-generate) ] [ set-nth ] 2bi
-        ] curry each
-    ] [ 0 >>i drop ] bi ;
+        mt-n swap seq>> '[
+            _ [ (mt-generate) ] [ set-wrap-nth ] 2bi
+        ] each
+    ] [ 0 >>i drop ] bi ; inline
 
 : init-mt-formula ( i seq -- f(seq[i]) )
-    dupd nth dup -30 shift bitxor 1812433253 * + 1+ 32 bits ;
+    dupd wrap-nth dup -30 shift bitxor 1812433253 * + 1+ 32 bits ; inline
 
 : init-mt-rest ( seq -- )
     mt-n 1- swap [
-        [ init-mt-formula ] [ [ 1+ ] dip set-nth ] 2bi
-    ] curry each ;
+        [ init-mt-formula ] [ [ 1+ ] dip set-wrap-nth ] 2bi
+    ] curry each ; inline
 
 : init-mt-seq ( seed -- seq )
-    32 bits mt-n <uint-array> <circular>
+    32 bits mt-n <uint-array>
     [ set-first ] [ init-mt-rest ] [ ] tri ;
 
 : mt-temper ( y -- yt )
@@ -53,7 +59,7 @@ TUPLE: mersenne-twister seq i ;
     dup -18 shift bitxor ; inline
 
 : next-index  ( mt -- i )
-    dup i>> dup mt-n < [ nip ] [ drop mt-generate 0 ] if ;
+    dup i>> dup mt-n < [ nip ] [ drop mt-generate 0 ] if ; inline
 
 PRIVATE>
 
@@ -66,7 +72,7 @@ M: mersenne-twister seed-random ( mt seed -- )
 
 M: mersenne-twister random-32* ( mt -- r )
     [ next-index ]
-    [ seq>> nth mt-temper ]
+    [ seq>> wrap-nth mt-temper ]
     [ [ 1+ ] change-i drop ] tri ;
 
 USE: init