mirror of
https://github.com/exaloop/codon.git
synced 2025-06-03 15:03:52 +08:00
262 lines
17 KiB
Python
262 lines
17 KiB
Python
# Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
|
|
|
|
from .seed import SeedSequence
|
|
|
|
N: Static[int] = 624
|
|
M: Static[int] = 397
|
|
|
|
MATRIX_A = u32(0x9908b0df) # constant vector a
|
|
UPPER_MASK = u32(0x80000000) # most significant w-r bits
|
|
LOWER_MASK = u32(0x7fffffff) # least significant r bits
|
|
|
|
# parameters for computing jump
|
|
W_SIZE: Static[int] = 32 # size of unsigned long
|
|
MEXP : Static[int] = 19937
|
|
P_SIZE: Static[int] = ((MEXP // W_SIZE) + 1)
|
|
LSB : Static[int] = 0x00000001
|
|
QQ : Static[int] = 7
|
|
LL : Static[int] = 128 # LL = 2^(QQ)
|
|
|
|
def get_coef(pf: Ptr[u64], deg: int):
|
|
return bool(pf[deg >> 5] & u64(LSB << (deg & 0x1f)))
|
|
|
|
@tuple
|
|
class MT19937:
|
|
data: Ptr[u32]
|
|
seed: SeedSequence
|
|
|
|
def __new__(seed, legacy: Static[int] = False):
|
|
if not isinstance(seed, SeedSequence):
|
|
return MT19937(SeedSequence(seed))
|
|
else:
|
|
g = MT19937(Ptr[u32](N + 1), seed)
|
|
|
|
if legacy:
|
|
if not isinstance(seed, int):
|
|
compile_error("'seed' must be an int for legacy seeding")
|
|
if seed < 0 or seed > 0xffffffff:
|
|
raise ValueError("Seed must be between 0 and 2**32 - 1")
|
|
g.seed_legacy(seed)
|
|
else:
|
|
g.data[0] = u32(N - 1)
|
|
val = seed.generate_state(N, u32)
|
|
key = g.state
|
|
|
|
key[0] = u32(0x80000000)
|
|
for i in range(1, N):
|
|
key[i] = val[i]
|
|
|
|
return g
|
|
|
|
def seed_legacy(self, seed: int):
|
|
seed = u32(seed)
|
|
seed &= u32(0xffffffff)
|
|
|
|
for pos in range(N):
|
|
self.state[pos] = u32(seed)
|
|
seed = (u32(1812433253) * (seed ^ (seed >> u32(30))) + u32(pos + 1)) & u32(0xffffffff)
|
|
|
|
self.data[0] = u32(N)
|
|
|
|
@property
|
|
def pos(self):
|
|
return int(self.data[0])
|
|
|
|
@property
|
|
def state(self):
|
|
return self.data + 1
|
|
|
|
def __get_state__(self):
|
|
from internal.gc import sizeof
|
|
p = Ptr[u32](N + 1)
|
|
str.memcpy(p.as_byte(), self.data.as_byte(), (N + 1) * sizeof(u32))
|
|
return (p,)
|
|
|
|
def __set_state__(self, state):
|
|
from internal.gc import sizeof
|
|
p = state[0]
|
|
str.memcpy(self.data.as_byte(), p.as_byte(), (N + 1) * sizeof(u32))
|
|
|
|
def copy_state(self, state: MT19937):
|
|
str.memcpy(self.data.as_byte(), state.data.as_byte(), (N + 1) * sizeof(u32))
|
|
|
|
def genrand_int32(self):
|
|
mag01t = (u32(0), MATRIX_A)
|
|
mag01 = Ptr[u32](__ptr__(mag01t).as_byte())
|
|
mt = self.state
|
|
|
|
if self.pos >= N:
|
|
kk = 0
|
|
|
|
while kk < int(N - M):
|
|
y = (mt[kk] & UPPER_MASK) | (mt[kk+1] & LOWER_MASK)
|
|
mt[kk] = mt[kk + M] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
kk += 1
|
|
|
|
while kk < int(N - 1):
|
|
y = (mt[kk] & UPPER_MASK) | (mt[kk+1] & LOWER_MASK)
|
|
mt[kk] = mt[kk+(M-N)] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
kk += 1
|
|
|
|
y = (mt[N-1] & UPPER_MASK) | (mt[0] & LOWER_MASK)
|
|
mt[N-1] = mt[M-1] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
self.data[0] = u32(0)
|
|
|
|
i = self.pos
|
|
y = mt[i]
|
|
self.data[0] = u32(i + 1)
|
|
y ^= (y >> u32(11))
|
|
y ^= (y << u32(7)) & u32(0x9d2c5680)
|
|
y ^= (y << u32(15)) & u32(0xefc60000)
|
|
y ^= (y >> u32(18))
|
|
return y
|
|
|
|
def genrand_res53(self):
|
|
a = self.genrand_int32() >> u32(5)
|
|
b = self.genrand_int32() >> u32(6)
|
|
return (int(a) * 67108864.0 + int(b)) * (1.0 / 9007199254740992.0)
|
|
|
|
def random(self):
|
|
return self.genrand_res53()
|
|
|
|
def init_u32(self, s: u32):
|
|
mt = self.state
|
|
mt[0] = s
|
|
for mti in range(1, N):
|
|
mt[mti] = (u32(1812433253) * (mt[mti-1] ^ (mt[mti-1] >> u32(30))) + u32(mti))
|
|
self.data[0] = u32(N)
|
|
|
|
def init_array(self, init_key: Ptr[u32], key_length: int):
|
|
mt = self.state
|
|
self.init_u32(u32(19650218))
|
|
i = 1
|
|
j = 0
|
|
|
|
k = N if N > key_length else key_length
|
|
while k:
|
|
mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> u32(30))) * u32(1664525))) + init_key[j] + u32(j)
|
|
i += 1
|
|
j += 1
|
|
if i >= N:
|
|
mt[0] = mt[N - 1]
|
|
i = 1
|
|
if j >= key_length:
|
|
j = 0
|
|
k -= 1
|
|
|
|
k = N - 1
|
|
while k:
|
|
mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> u32(30))) * u32(1566083941))) - u32(i)
|
|
i += 1
|
|
if i >= N:
|
|
mt[0] = mt[N - 1]
|
|
i = 1
|
|
k -= 1
|
|
|
|
mt[0] = u32(0x80000000)
|
|
|
|
def seed_cpython(self, s: int):
|
|
init_key = (u32(s & ((1 << 32) - 1)), u32(s >> 32))
|
|
self.init_array(Ptr[u32](__ptr__(init_key).as_byte()), 2 if init_key[1] else 1)
|
|
|
|
# jump
|
|
|
|
def gen_next(self):
|
|
mag01t = (u32(0), MATRIX_A)
|
|
mag01 = Ptr[u32](__ptr__(mag01t).as_byte())
|
|
mt = self.state
|
|
|
|
kk = self.pos
|
|
|
|
if kk < N - M:
|
|
y = (mt[kk] & UPPER_MASK) | (mt[kk+1] & LOWER_MASK)
|
|
mt[kk] = mt[kk + M] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
self.data[0] += u32(1)
|
|
kk += 1
|
|
elif kk < N - 1:
|
|
y = (mt[kk] & UPPER_MASK) | (mt[kk+1] & LOWER_MASK)
|
|
mt[kk] = mt[kk+(M-N)] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
self.data[0] += u32(1)
|
|
kk += 1
|
|
elif kk == N - 1:
|
|
y = (mt[N-1] & UPPER_MASK) | (mt[0] & LOWER_MASK)
|
|
mt[N-1] = mt[M-1] ^ (y >> u32(1)) ^ mag01[int(y & u32(1))]
|
|
self.data[0] = u32(0)
|
|
|
|
def add_state(self, state2: MT19937):
|
|
pt1 = self.pos
|
|
pt2 = state2.pos
|
|
mt1 = self.state
|
|
mt2 = state2.state
|
|
i = 0
|
|
|
|
if pt2 - pt1 >= 0:
|
|
while i < N - pt2:
|
|
mt1[i + pt1] ^= mt2[i + pt2]
|
|
i += 1
|
|
while i < N - pt1:
|
|
mt1[i + pt1] ^= mt2[i + (pt2 - N)]
|
|
i += 1
|
|
while i < N:
|
|
mt1[i + (pt1 - N)] ^= mt2[i + (pt2 - N)]
|
|
i += 1
|
|
else:
|
|
while i < N - pt1:
|
|
mt1[i + pt1] ^= mt2[i + pt2]
|
|
i += 1
|
|
while i < N - pt2:
|
|
mt1[i + (pt1 - N)] ^= mt2[i + pt2]
|
|
i += 1
|
|
while i < N:
|
|
mt1[i + (pt1 - N)] ^= mt2[i + (pt2 - N)]
|
|
i += 1
|
|
|
|
def horner1(self, pf: Ptr[u64]):
|
|
from internal.gc import sizeof
|
|
i = MEXP - 1
|
|
temp_data = __array__[u32](N + 1)
|
|
str.memset(temp_data.ptr.as_byte(), byte(0), (N + 1) * sizeof(u32))
|
|
temp = MT19937(temp_data.ptr)
|
|
|
|
while not get_coef(pf, i):
|
|
i -= 1
|
|
|
|
if i > 0:
|
|
temp.copy_state(self)
|
|
temp.gen_next()
|
|
i -= 1
|
|
while i > 0:
|
|
if get_coef(pf, i):
|
|
temp.add_state(self)
|
|
temp.gen_next()
|
|
i -= 1
|
|
if get_coef(pf, 0):
|
|
temp.add_state(self)
|
|
elif i == 0:
|
|
temp.copy_state(self)
|
|
|
|
self.copy_state(temp)
|
|
|
|
def jump_state(self):
|
|
if self.pos >= N:
|
|
self.data[0] = u32(0)
|
|
|
|
self.horner1(MT19937.poly_coef())
|
|
|
|
def jump_inplace(self, jumps: int):
|
|
for _ in range(jumps):
|
|
self.jump_state()
|
|
|
|
def next32(self):
|
|
return self.genrand_int32()
|
|
|
|
def next_double(self):
|
|
return self.genrand_res53()
|
|
|
|
# TODO: somehow add wrapping to this...
|
|
@pure
|
|
@llvm
|
|
def poly_coef() -> Ptr[u64]:
|
|
@pf = private unnamed_addr constant [624 x i64] [i64 1927166307, i64 3044056772, i64 2284297142, i64 2820929765, i64 651705945, i64 69149273, i64 3892165397, i64 2337412983, i64 1219880790, i64 3207074517, i64 3836784057, i64 189286826, i64 1049791363, i64 3916249550, i64 2942382547, i64 166392552, i64 861176918, i64 3246476411, i64 2302311555, i64 4273801148, i64 29196903, i64 1363664063, i64 3802562022, i64 2600400244, i64 3090369801, i64 4040416970, i64 1432485208, i64 3632558139, i64 4015816763, i64 3013316418, i64 551532385, i64 3592224467, i64 3479125595, i64 1195467127, i64 2391032553, i64 2393493419, i64 1482493632, i64 1625159565, i64 748389672, i64 4042774030, i64 2998615036, i64 3393119101, i64 2177492569, i64 2265897321, i64 2507383006, i64 3461498961, i64 2003319700, i64 1942857197, i64 1455226044, i64 4097545580, i64 529653268, i64 3204756480, i64 2486748289, i64 495294513, i64 3396001954, i64 2643963605, i64 2655404568, i64 3881604377, i64 624710790, i64 3443737948, i64 1941294296, i64 2139259604, i64 3368734020, i64 422436761, i64 3602810182, i64 1384691081, i64 3035786407, i64 2551797119, i64 537227499, i64 65486120, i64 642436100, i64 2023822537, i64 2515598203, i64 1122953367, i64 2882306242, i64 1743213032, i64 321965189, i64 336496623, i64 2436602518, i64 3556266590, i64 1055117829, i64 463541647, i64 743234441, i64 527083645, i64 2606668346, i64 2274046499, i64 2761475053, i64 2760669048, i64 2538258534, i64 487125077, i64 3365962306, i64 3604906217, i64 2714700608, i64 680709708, i64 2217161159, i64 1614899374, i64 3710119533, i64 3201300658, i64 3752620679, i64 2755041105, i64 3129723037, i64 1247297753, i64 2812642690, i64 4114340845, i64 3485092247, i64 2752814364, i64 3586551747, i64 4073138437, i64 3462966585, i64 2924318358, i64 4061374901, i64 3314086806, i64 2640385723, i64 744590670, i64 3007586513, i64 3959120371, i64 997207767, i64 3420235506, i64 2092400998, i64 3190305685, i64 60965738, i64 549507222, i64 3784354415, i64 3209279509, i64 1238863299, i64 2605037827, i64 178570440, i64 1743491299, i64 4079686640, i64 2136795825, i64 3435430548, i64 1679732443, i64 1835708342, i64 2159367000, i64 1924487218, i64 4059723674, i64 996192116, i64 2308091645, i64 1336281586, i64 674600050, i64 1642572529, i64 1383973289, i64 2202960007, i64 3165481279, i64 3385474038, i64 2501318550, i64 2671842890, i64 3084085109, i64 3475033915, i64 1551329147, i64 4101397249, i64 1205851807, i64 3641536021, i64 3607635071, i64 1609126163, i64 2910426664, i64 3324508658, i64 4244311266, i64 254034382, i64 1258304384, i64 1914048768, i64 1358592011, i64 527610138, i64 3072108727, i64 4289413885, i64 1417001678, i64 2445445945, i64 896462712, i64 339855811, i64 3699378285, i64 2529457297, i64 3049459401, i64 2723472429, i64 2838633181, i64 2520397330, i64 3272339035, i64 1667003847, i64 3742634787, i64 942706520, i64 2301027215, i64 1907791250, i64 2306299096, i64 1021173342, i64 1539334516, i64 2907834628, i64 3199959207, i64 1556251860, i64 3642580275, i64 2355865416, i64 285806145, i64 867932457, i64 1177354172, i64 3291107470, i64 4022765061, i64 1613380116, i64 588147929, i64 650574324, i64 1236855601, i64 1371354511, i64 2085218212, i64 1203081931, i64 420526905, i64 1022192219, i64 2903287064, i64 2470845899, i64 3649873273, i64 2502333582, i64 3972385637, i64 4246356763, i64 199084157, i64 1567178788, i64 2107121836, i64 4293612856, i64 1902910177, i64 332397359, i64 83422598, i64 3614961721, i64 456321943, i64 2277615967, i64 2302518510, i64 3258315116, i64 2521897172, i64 3900282042, i64 4186973154, i64 3146532165, i64 2299685029, i64 3889120948, i64 1293301857, i64 187455105, i64 3395849230, i64 913321567, i64 3093513909, i64 1440944571, i64 1923481911, i64 338680924, i64 1204882963, i64 2739724491, i64 2886241328, i64 2408907774, i64 1299817192, i64 2474012871, i64 45400213, i64 553186784, i64 134558656, i64 2180943666, i64 2870807589, i64 76511085, i64 3053566760, i64 2516601415, i64 4172865902, i64 1751297915, i64 1251975234, i64 2964780642, i64 1412975316, i64 2739978478, i64 2171013719, i64 637935041, i64 975972384, i64 3044407449, i64 3111425639, i64 1938684970, i64 2860857400, i64 13419586, i64 2772079268, i64 3484375614, i64 3184054178, i64 159924837, i64 1386213021, i64 2765617231, i64 2523689118, i64 1283505218, i64 3510789588, i64 4125878259, i64 2990287597, i64 2152014833, i64 3084155970, i64 2815101609, i64 1932985704, i64 114887365, i64 1712687646, i64 2550515629, i64 3299051916, i64 2022747614, i64 2143630992, i64 2244188960, i64 3309469192, i64 3234358520, i64 800720365, i64 3278176634, i64 554357439, i64 2415629802, i64 1620877315, i64 2389462898, i64 2229691332, i64 1007748450, i64 1966873768, i64 2264971043, i64 1214524156, i64 346854700, i64 3471905342, i64 3984889660, i64 4034246840, i64 216712649, i64 4027196762, i64 3754772604, i64 2121785562, i64 2347070732, i64 7457687, i64 1443375102, i64 683948143, i64 2940226032, i64 3211475670, i64 2836507357, i64 774899409, i64 1588968308, i64 780438009, i64 3278878781, i64 2217181540, i64 2184194887, i64 1642129086, i64 69346830, i64 297114710, i64 3841068188, i64 2631265450, i64 4167492314, i64 2613519651, i64 1388582503, i64 2171556668, i64 1201873758, i64 2698772382, i64 207791958, i64 3936134563, i64 3725025702, i64 3306317801, i64 1055730422, i64 4069230694, i64 1767821343, i64 4252407395, i64 2422583118, i64 3158834399, i64 3754582617, i64 1112422556, i64 376187931, i64 3137549150, i64 712221089, i64 3300799453, i64 3868250200, i64 1165257666, i64 2494837767, i64 131304831, i64 1619349427, i64 1958236644, i64 3678218946, i64 3651007751, i64 2261987899, i64 1567368524, i64 2193599522, i64 3034394674, i64 2994602555, i64 3072727647, i64 889094521, i64 1089692095, i64 1822324824, i64 3876999182, i64 1703361286, i64 902229515, i64 4213728487, i64 3838170364, i64 672727494, i64 2240733828, i64 3858539469, i64 1149254245, i64 4166055926, i64 4193525313, i64 1709921593, i64 2278290377, i64 3190784116, i64 2919588882, i64 1012709717, i64 3640562031, i64 2931984863, i64 3515665246, i64 250577343, i64 1147230194, i64 1183856202, i64 3734511989, i64 3243867808, i64 3499383067, i64 2985115159, i64 2036821626, i64 3298159553, i64 2726542838, i64 1686910320, i64 1778823772, i64 965412224, i64 233509772, i64 3843098861, i64 1312622954, i64 500855830, i64 2950562091, i64 1915683607, i64 3405781138, i64 596073719, i64 2195150546, i64 3381728478, i64 546426436, i64 3527890868, i64 2324975353, i64 2241074266, i64 3992514859, i64 2576108287, i64 4077653225, i64 2632319392, i64 3127212632, i64 917000669, i64 2498161805, i64 3980835128, i64 2259526768, i64 1083920509, i64 1187452089, i64 97018536, i64 3056075838, i64 2059706760, i64 2373335692, i64 182196406, i64 2136713111, i64 1762080153, i64 1572125803, i64 1145919955, i64 1023966754, i64 3921694345, i64 1632005969, i64 1418372326, i64 354407429, i64 2438288265, i64 1620072033, i64 1586320921, i64 1044153697, i64 969324572, i64 613487980, i64 4230993062, i64 397726764, i64 2194259193, i64 735511759, i64 2066049260, i64 88093248, i64 1562536153, i64 2114157419, i64 3630951546, i64 589238503, i64 3120654384, i64 2521793793, i64 2746692127, i64 2557723425, i64 889897693, i64 2778878177, i64 643269509, i64 3342389831, i64 19218890, i64 3442706236, i64 3314581273, i64 3503147052, i64 1546343434, i64 1448529060, i64 529038801, i64 2748942264, i64 2213019208, i64 111314040, i64 2488697563, i64 1180642808, i64 2605272289, i64 4207476668, i64 1502558669, i64 2972370981, i64 4204339995, i64 1046225278, i64 992840610, i64 3847290298, i64 2387673094, i64 2221565747, i64 1045901716, i64 3997739302, i64 1556952765, i64 1103336648, i64 279418400, i64 2711316466, i64 2336215718, i64 2317900806, i64 974624729, i64 909575434, i64 1675610631, i64 1922393214, i64 2054896570, i64 3197007361, i64 3932554569, i64 1008619802, i64 3349254938, i64 113511461, i64 932630384, i64 2098759268, i64 3436837432, i64 3119972401, i64 1612590197, i64 2281609013, i64 4174211248, i64 4016332246, i64 2097525539, i64 1398632760, i64 1543697535, i64 2419227174, i64 1676465074, i64 2882923045, i64 23216933, i64 808195649, i64 3690720147, i64 484419260, i64 2254772642, i64 2975434733, i64 288528113, i64 204598404, i64 589968818, i64 3021152400, i64 2463155141, i64 1397846755, i64 157285579, i64 4230258857, i64 2469135246, i64 625357422, i64 3435224647, i64 465239124, i64 1022535736, i64 2823317040, i64 274194469, i64 2214966446, i64 3661001613, i64 518802547, i64 2293436304, i64 1335881988, i64 2247010176, i64 1856732584, i64 1088028094, i64 1877563709, i64 1015352636, i64 1700817932, i64 2960695857, i64 1882229300, i64 1666906557, i64 1838841022, i64 3983797810, i64 1667630361, i64 385998221, i64 241341791, i64 403550441, i64 2629200403, i64 3552759102, i64 2029750442, i64 2247999048, i64 2726665298, i64 2507798776, i64 2419064129, i64 1266444923, i64 526255242, i64 2384866697, i64 1886200981, i64 3954956408, i64 2171436866, i64 2295200753, i64 1047315850, i64 1967809707, i64 2860382973, i64 3918334466, i64 3057439479, i64 952682588, i64 1925559679, i64 3112119050, i64 3833190964, i64 1430139895, i64 2089165610, i64 3009202424, i64 3989186157, i64 3395807230, i64 347600520, i64 120428923, i64 3017004655, i64 1384933954, i64 303039929, i64 234010146, i64 2278760249, i64 315514836, i64 3987659575, i64 1239335668, i64 2387869477, i64 3885908826, i64 1983922602, i64 698609264, i64 3009002846, i64 1520611399, i64 809159940, i64 3089771783, i64 374838722, i64 2789914419, i64 2500831937, i64 3751970335, i64 4279852547, i64 2362894437, i64 1588814060, i64 1671213155, i64 434218829, i64 2126587176, i64 2002526422, i64 2756464095, i64 141700479, i64 2965974322, i64 2211530172, i64 992085992, i64 1943691492, i64 2705131817, i64 2519208889, i64 1938768395, i64 3949294294, i64 354046666, i64 2158272751, i64 602858583, i64 0], align 16
|
|
ret ptr @pf
|