nx_kv_arena_test.nx source
↩ module page · 206 lines · 9461 B
1// nx_kv_arena_test.nx -- smoke for PagedAttention page pool.
2//
3// Exercises:
4// 1. Allocation + canary + initial state (all pages free)
5// 2. Sequence creation + page appends from empty pool
6// 3. Pool exhaustion: more requests than pages -> POOL_EXHAUSTED
7// 4. Sequence free returns pages to the pool
8// 5. Newly-free pages can be allocated by a different sequence
9// 6. Multi-sequence independence (each owns its own indices)
10// 7. Page byte-pointer arithmetic: indices map to distinct byte
11// offsets in page_storage
12// 8. Canary tamper on arena and sequence
13// 9. Bad-input gates
14// 10. Sealed-enum verdict gate
15// 11. Idempotent re-free (calling free twice on same seq is no-op)
16
17import "nx_syscalls.nx"
18import "nx_kv_arena.nx"
19
20func main() -> i64 {
21 // 4-page arena, 256 bytes per page (small for test).
22 let a: *NxKvArena = nx_kv_arena_new(4, 256)
23 if nx_kv_arena_is_valid(a) != 1 { return 1 }
24 if nx_kv_arena_n_pages(a) != 4 { return 2 }
25 if nx_kv_arena_n_free(a) != 4 { return 3 }
26 if nx_kv_arena_n_in_use(a) != 0 { return 4 }
27
28 // ----- 2. Sequence + appends -----
29 let s1: *NxKvSequence = nx_kv_sequence_new(a, 1001)
30 if nx_kv_sequence_is_valid(s1) != 1 { return 5 }
31 if nx_kv_sequence_n_pages(s1) != 0 { return 6 }
32
33 let p0: i64 = nx_kv_sequence_append_page(s1)
34 if p0 < 0 { return 7 }
35 if nx_kv_sequence_n_pages(s1) != 1 { return 8 }
36 if nx_kv_arena_n_in_use(a) != 1 { return 9 }
37 if nx_kv_arena_n_free(a) != 3 { return 10 }
38
39 let p1: i64 = nx_kv_sequence_append_page(s1)
40 if p1 < 0 { return 11 }
41 if p1 == p0 { return 12 } // distinct page indices
42 if nx_kv_sequence_n_pages(s1) != 2 { return 13 }
43
44 // ----- 3. Pool exhaustion -----
45 // Burn the remaining 2 pages with s2 + s3, then try a 5th alloc.
46 let s2: *NxKvSequence = nx_kv_sequence_new(a, 1002)
47 let p2: i64 = nx_kv_sequence_append_page(s2)
48 if p2 < 0 { return 14 }
49 let p3: i64 = nx_kv_sequence_append_page(s2)
50 if p3 < 0 { return 15 }
51 if nx_kv_arena_n_free(a) != 0 { return 16 }
52 if nx_kv_arena_n_in_use(a) != 4 { return 17 }
53
54 let s3: *NxKvSequence = nx_kv_sequence_new(a, 1003)
55 let p_exhaust: i64 = nx_kv_sequence_append_page(s3)
56 if p_exhaust != (0 - NX_KV_POOL_EXHAUSTED) { return 18 }
57 if nx_kv_sequence_n_pages(s3) != 0 { return 19 } // unchanged
58
59 // ----- 4. Free s1 -> 2 pages back to pool -----
60 if nx_kv_sequence_free(s1) != NX_KV_OK { return 20 }
61 if nx_kv_arena_n_free(a) != 2 { return 21 }
62 if nx_kv_arena_n_in_use(a) != 2 { return 22 }
63 if nx_kv_sequence_n_pages(s1) != 0 { return 23 } // reset
64
65 // ----- 5. s3 can now allocate from the freed pages -----
66 let p4: i64 = nx_kv_sequence_append_page(s3)
67 if p4 < 0 { return 24 }
68 let p5: i64 = nx_kv_sequence_append_page(s3)
69 if p5 < 0 { return 25 }
70 if nx_kv_sequence_n_pages(s3) != 2 { return 26 }
71 if nx_kv_arena_n_free(a) != 0 { return 27 }
72
73 // ----- 6. Multi-sequence independence -----
74 // s2 still has its 2 pages. s3 has 2 pages. Page indices are
75 // distinct across the two sequences (no overlap).
76 if s2.page_indices[0] == s3.page_indices[0] { return 28 }
77 if s2.page_indices[0] == s3.page_indices[1] { return 29 }
78 if s2.page_indices[1] == s3.page_indices[0] { return 30 }
79 if s2.page_indices[1] == s3.page_indices[1] { return 31 }
80
81 // ----- 7. Page byte-pointer maps distinct offsets -----
82 let ptr_0: *u8 = nx_kv_arena_page_bytes(a, 0)
83 let ptr_1: *u8 = nx_kv_arena_page_bytes(a, 1)
84 if (ptr_0 as i64) == 0 { return 32 }
85 if (ptr_1 as i64) == 0 { return 33 }
86 let delta: i64 = (ptr_1 as i64) - (ptr_0 as i64)
87 if delta != 256 { return 34 } // page_bytes apart
88 // Writing through one page doesn't disturb another.
89 ptr_0[0] = 0xAA as u8
90 ptr_1[0] = 0xBB as u8
91 if (ptr_0[0] as i64) != 0xAA { return 35 }
92 if (ptr_1[0] as i64) != 0xBB { return 36 }
93
94 // Out-of-range index -> NULL.
95 let ptr_bad: *u8 = nx_kv_arena_page_bytes(a, 999)
96 if (ptr_bad as i64) != 0 { return 37 }
97 let ptr_neg: *u8 = nx_kv_arena_page_bytes(a, -1)
98 if (ptr_neg as i64) != 0 { return 38 }
99
100 // ----- 8. Canary tamper -----
101 let tamper_a: *NxKvArena = nx_kv_arena_new(2, 64)
102 let tamper_s: *NxKvSequence = nx_kv_sequence_new(tamper_a, 9999)
103 tamper_a.canary_post = 0xDEADBEEF
104 if nx_kv_arena_is_valid(tamper_a) != 0 { return 39 }
105 if nx_kv_arena_n_in_use(tamper_a) != -1 { return 40 }
106 if nx_kv_sequence_append_page(tamper_s) != (0 - NX_KV_TAMPER) { return 41 }
107
108 let tamper_a2: *NxKvArena = nx_kv_arena_new(2, 64)
109 let tamper_s2: *NxKvSequence = nx_kv_sequence_new(tamper_a2, 9998)
110 tamper_s2.canary_post = 0xDEADBEEF
111 if nx_kv_sequence_is_valid(tamper_s2) != 0 { return 42 }
112 if nx_kv_sequence_append_page(tamper_s2) != (0 - NX_KV_TAMPER) { return 43 }
113 if nx_kv_sequence_n_pages(tamper_s2) != -1 { return 44 }
114
115 // ----- 9. Bad-input gates -----
116 if (nx_kv_arena_new(0, 256) as i64) != 0 { return 45 }
117 if (nx_kv_arena_new(-1, 256) as i64) != 0 { return 46 }
118 if (nx_kv_arena_new(4, 0) as i64) != 0 { return 47 }
119 if (nx_kv_arena_new(4, -1) as i64) != 0 { return 48 }
120 if (nx_kv_arena_new(NX_KV_ARENA_MAX_PAGES + 1, 256) as i64) != 0 { return 49 }
121
122 // ----- 10. Sealed-enum verdict gate -----
123 if nx_kv_verdict_is_valid(NX_KV_OK) != 1 { return 50 }
124 if nx_kv_verdict_is_valid(NX_KV_POOL_EXHAUSTED) != 1 { return 51 }
125 if nx_kv_verdict_is_valid(-1) != 0 { return 52 }
126 if nx_kv_verdict_is_valid(NX_KV_N_VERDICTS) != 0 { return 53 }
127
128 // ----- 11. Idempotent re-free -----
129 // s2 currently has 2 pages.
130 if nx_kv_sequence_free(s2) != NX_KV_OK { return 54 }
131 if nx_kv_sequence_n_pages(s2) != 0 { return 55 }
132 // Free again -- should be a no-op, not a corruption.
133 if nx_kv_sequence_free(s2) != NX_KV_OK { return 56 }
134 if nx_kv_sequence_n_pages(s2) != 0 { return 57 }
135 // Arena free count unchanged across the re-free (only the
136 // first free's 2 pages came back; re-free returned 0 more).
137 // Total free count: 0 (start) + s1 freed earlier returned 2 +
138 // s3 took 2 of those + s2 just freed returned 2 = 2.
139 if nx_kv_arena_n_free(a) != 2 { return 58 }
140
141 // ----- 12. H4 PREFIX CACHING: share existing page -----
142 // Fresh 4-page arena to test sharing cleanly.
143 let ah4: *NxKvArena = nx_kv_arena_new(4, 128)
144 let seq_A: *NxKvSequence = nx_kv_sequence_new(ah4, 4001)
145 let seq_B: *NxKvSequence = nx_kv_sequence_new(ah4, 4002)
146 let seq_C: *NxKvSequence = nx_kv_sequence_new(ah4, 4003)
147
148 // seq_A appends a page (representing e.g. shared system prompt).
149 let prompt_page: i64 = nx_kv_sequence_append_page(seq_A)
150 if prompt_page < 0 { return 59 }
151 if nx_kv_arena_page_refcount(ah4, prompt_page) != 1 { return 60 }
152 if nx_kv_arena_n_in_use(ah4) != 1 { return 61 }
153 if nx_kv_arena_n_free(ah4) != 3 { return 62 }
154
155 // seq_B SHARES that page (instead of duplicating K/V data).
156 let shared_b: i64 = nx_kv_sequence_share_page(seq_B, prompt_page)
157 if shared_b != prompt_page { return 63 }
158 if nx_kv_arena_page_refcount(ah4, prompt_page) != 2 { return 64 }
159 // CRUCIAL: n_in_use unchanged -- sharing doesn't consume a slot.
160 if nx_kv_arena_n_in_use(ah4) != 1 { return 65 }
161 if nx_kv_arena_n_free(ah4) != 3 { return 66 }
162 if nx_kv_sequence_n_pages(seq_B) != 1 { return 67 }
163
164 // seq_C ALSO shares the page (3 sequences -> one page).
165 let shared_c: i64 = nx_kv_sequence_share_page(seq_C, prompt_page)
166 if shared_c != prompt_page { return 68 }
167 if nx_kv_arena_page_refcount(ah4, prompt_page) != 3 { return 69 }
168 if nx_kv_arena_n_in_use(ah4) != 1 { return 70 }
169 if nx_kv_arena_n_free(ah4) != 3 { return 71 }
170
171 // ----- 13. Free decrements refcount, page stays in_use until last -----
172 if nx_kv_sequence_free(seq_A) != NX_KV_OK { return 72 }
173 if nx_kv_arena_page_refcount(ah4, prompt_page) != 2 { return 73 }
174 if nx_kv_arena_n_in_use(ah4) != 1 { return 74 } // still in use
175 if nx_kv_arena_n_free(ah4) != 3 { return 75 }
176
177 if nx_kv_sequence_free(seq_B) != NX_KV_OK { return 76 }
178 if nx_kv_arena_page_refcount(ah4, prompt_page) != 1 { return 77 }
179 if nx_kv_arena_n_in_use(ah4) != 1 { return 78 }
180 if nx_kv_arena_n_free(ah4) != 3 { return 79 }
181
182 if nx_kv_sequence_free(seq_C) != NX_KV_OK { return 80 }
183 if nx_kv_arena_page_refcount(ah4, prompt_page) != 0 { return 81 }
184 if nx_kv_arena_n_in_use(ah4) != 0 { return 82 }
185 if nx_kv_arena_n_free(ah4) != 4 { return 83 }
186
187 // ----- 14. Share-a-free-page rejected -----
188 let seq_D: *NxKvSequence = nx_kv_sequence_new(ah4, 4004)
189 let rc_free_share: i64 = nx_kv_sequence_share_page(seq_D, prompt_page)
190 if rc_free_share != (0 - NX_KV_NOT_FOUND) { return 84 }
191 if nx_kv_sequence_n_pages(seq_D) != 0 { return 85 }
192
193 // ----- 15. Bad-input on share -----
194 let p_real: i64 = nx_kv_sequence_append_page(seq_D)
195 if p_real < 0 { return 86 }
196 if nx_kv_sequence_share_page(seq_D, -1) != (0 - NX_KV_BAD_INPUT) { return 87 }
197 if nx_kv_sequence_share_page(seq_D, 999) != (0 - NX_KV_BAD_INPUT) { return 88 }
198
199 // ----- 16. Refcount accessor bad-input -----
200 if nx_kv_arena_page_refcount(ah4, -1) != -1 { return 89 }
201 if nx_kv_arena_page_refcount(ah4, 999) != -1 { return 90 }
202 let null_a: *NxKvArena = (0 as i64) as *NxKvArena
203 if nx_kv_arena_page_refcount(null_a, 0) != -1 { return 91 }
204
205 return 0
206}