Prefix/lib/std/csprng.pre at e10392ca44105f593179b1183e2d426f0dee049d · python-processing-unit/Prefix · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
! ChaCha20-based cryptographically secure pseudorandom number generator

INT MASK32 = SUB( POW(0b10, 0b100000), 1 )  ! 2^32-1

! ChaCha20 constants ("expand 32-byte k")
MAP CH_CONST = < ^
    0b0 = 0b01100001011100000111000001100101, ^
    0b1 = 0b00110011001000000110010001101110, ^
    0b10 = 0b01111001011000100010110100110010, ^
    0b11 = 0b01101011001000000110010101110100 ^
>

TNS ch_key = TNS([0b1000], 0b0)  ! 8 words
TNS ch_nonce = [0b0, 0b0, 0b0] ! use literals for short TNS
INT ch_counter = 0
TNS ch_buf = TNS([0b10000], 0b0) ! 16 words
INT ch_buf_pos = TLEN(ch_buf, 0b1)  ! set to buffer length to force refill

FUNC INT ROTL32(INT x, INT n){
    INT x32 = BAND(x, MASK32)
    INT left = BAND(SHL(x32, n), MASK32)
    INT right = SHR(x32, SUB(0b100000, n))  ! 32 - n
    RETURN( BAND(BOR(left, right), MASK32) )
}

FUNC INT QR(TNS s, INT ia, INT ib, INT ic, INT id){
    INT a = s[ia]
    INT b = s[ib]
    INT c = s[ic]
    INT d = s[id]

    a = BAND( ADD(a, b), MASK32 )
    d = ROTL32( BXOR(d, a), 0b10000 )   ! 16
    c = BAND( ADD(c, d), MASK32 )
    b = ROTL32( BXOR(b, c), 0b1100 )    ! 12
    a = BAND( ADD(a, b), MASK32 )
    d = ROTL32( BXOR(d, a), 0b1000 )    ! 8
    c = BAND( ADD(c, d), MASK32 )
    b = ROTL32( BXOR(b, c), 0b111 )     ! 7

    s[ia] = a
    s[ib] = b
    s[ic] = c
    s[id] = d
    RETURN(0)
}

FUNC INT CHACHA_BLOCK(){
    TNS k = ch_key
    TNS n = ch_nonce
    TNS state = [ ^
        CH_CONST<0b0>, ^
        CH_CONST<0b1>, ^
        CH_CONST<0b10>, ^
        CH_CONST<0b11>, ^
        k[0b1], ^
        k[0b10], ^
        k[0b11], ^
        k[0b100], ^
        k[0b101], ^
        k[0b110], ^
        k[0b111], ^
        k[0b1000], ^
        ch_counter, ^
        n[0b1], ^
        n[0b10], ^
        n[0b11] ^
    ]
    FOR(i, 0b1010){  ! 10 double-rounds
        QR(state, 0b1, 0b101, 0b1001, 0b1101)
        QR(state, 0b10, 0b110, 0b1010, 0b1110)
        QR(state, 0b11, 0b111, 0b1011, 0b1111)
        QR(state, 0b100, 0b1000, 0b1100, 0b10000)

        QR(state, 0b1, 0b110, 0b1011, 0b10000)
        QR(state, 0b10, 0b111, 0b1100, 0b1101)
        QR(state, 0b11, 0b1000, 0b1001, 0b1110)
        QR(state, 0b100, 0b101, 0b1010, 0b1111)
    }

    ch_buf[0b1] = BAND( ADD(state[0b1], CH_CONST<0b0>), MASK32 )
    ch_buf[0b10] = BAND( ADD(state[0b10], CH_CONST<0b1>), MASK32 )
    ch_buf[0b11] = BAND( ADD(state[0b11], CH_CONST<0b10>), MASK32 )
    ch_buf[0b100] = BAND( ADD(state[0b100], CH_CONST<0b11>), MASK32 )
    ch_buf[0b101] = BAND( ADD(state[0b101], k[0b1]), MASK32 )
    ch_buf[0b110] = BAND( ADD(state[0b110], k[0b10]), MASK32 )
    ch_buf[0b111] = BAND( ADD(state[0b111], k[0b11]), MASK32 )
    ch_buf[0b1000] = BAND( ADD(state[0b1000], k[0b100]), MASK32 )
    ch_buf[0b1001] = BAND( ADD(state[0b1001], k[0b101]), MASK32 )
    ch_buf[0b1010] = BAND( ADD(state[0b1010], k[0b110]), MASK32 )
    ch_buf[0b1011] = BAND( ADD(state[0b1011], k[0b111]), MASK32 )
    ch_buf[0b1100] = BAND( ADD(state[0b1100], k[0b1000]), MASK32 )
    ch_buf[0b1101] = BAND( ADD(state[0b1101], ch_counter), MASK32 )
    ch_buf[0b1110] = BAND( ADD(state[0b1110], n[0b1]), MASK32 )
    ch_buf[0b1111] = BAND( ADD(state[0b1111], n[0b10]), MASK32 )
    ch_buf[0b10000] = BAND( ADD(state[0b10000], n[0b11]), MASK32 )

    RETURN(0b0)
}

FUNC INT REFILL_BUF(){
    CHACHA_BLOCK()
    ch_buf_pos = 0b0
    ch_counter = BAND( ADD(ch_counter, 0b1), MASK32 )
    RETURN(0b0)
}

FUNC INT DERIVE_KEY_AND_NONCE(INT seed){
    INT s = BAND(seed, MASK32)
    FOR(i, 0b1000){  ! 8 key words (1000 == 8)
        s = BAND( ADD( MUL(s, ^
            0b01000001110001100100111001101101), ^
            0b00000000000000000011000000111001 ), ^
            MASK32 ^
        )
        ch_key[i] = s
    }
    FOR(j, 0b11){  ! 3 nonce words (11 == 3)
        s = BAND( ADD( MUL(s, ^
            0b01000001110001100100111001101101), ^
            0b00000000000000000011000000111001 ), ^
            MASK32 ^
        )
        ch_nonce[j] = s
    }
    RETURN(0b0)
}

FUNC INT SEED(INT seed){
    DERIVE_KEY_AND_NONCE(seed)
    ch_counter = 0b0
    ch_buf = TNS([0b10000], 0b0)
    ch_buf_pos = TLEN(ch_buf, 0b1)  ! set to buffer length to force refill
    RETURN(ch_counter)
}

FUNC INT NEXT(){
    IF( GTE(ch_buf_pos, TLEN(ch_buf, 0b1)) ){  ! buffer length
        REFILL_BUF()
    }
    INT v = ch_buf[ ADD(ch_buf_pos, 0b1) ]
    ch_buf_pos = ADD(ch_buf_pos, 0b1)
    RETURN(v)
}

FUNC INT RANGE(INT max){
    ASSERT( GT(max, 0b0) )
    RETURN( MOD(NEXT(), max) )
}

FUNC INT RANGE_MIN_MAX(INT min, INT max){
    ASSERT( LTE(min, max) )
    INT range = SUB(max, min)
    IF( EQ(range, 0b0) ){ RETURN(min) }
    ASSERT( GT(range, 0b0) )
    INT offset = RANGE(range)
    RETURN( ADD(offset, min) )
}