radiance · Radiance self-hosting compiler

test/tests/prog.rbtree.rad 8.6 KiB raw

//! returns: 0
//! Red-black tree.
//! Implement a red-black tree (balanced BST) using a stack-allocated node pool.

constant POOL_SIZE: u32 = 128;
constant NIL: u32 = 0;

constant RED: u32 = 0;
constant BLACK: u32 = 1;

record RBNode {
    key: i32,
    color: u32,
    left: u32,
    right: u32,
    parent: u32,
}

record RBTree {
    pool: *mut [RBNode],
    poolNext: u32,
    root: u32,
    inorder: *mut [i32],
    inorderCount: u32,
}

unsafe fn allocNode(t: *mut RBTree, key: i32) -> u32 {
    let idx: u32 = t.poolNext;
    set t.poolNext += 1;
    set t.pool[idx] = RBNode { key, color: RED, left: NIL, right: NIL, parent: NIL };
    return idx;
}

unsafe fn rotateLeft(t: *mut RBTree, x: u32) {
    let y: u32 = t.pool[x].right;
    set t.pool[x].right = t.pool[y].left;
    if t.pool[y].left <> NIL {
        set t.pool[t.pool[y].left].parent = x;
    }
    set t.pool[y].parent = t.pool[x].parent;
    if t.pool[x].parent == NIL {
        set t.root = y;
    } else if x == t.pool[t.pool[x].parent].left {
        set t.pool[t.pool[x].parent].left = y;
    } else {
        set t.pool[t.pool[x].parent].right = y;
    }
    set t.pool[y].left = x;
    set t.pool[x].parent = y;
}

unsafe fn rotateRight(t: *mut RBTree, x: u32) {
    let y: u32 = t.pool[x].left;
    set t.pool[x].left = t.pool[y].right;
    if t.pool[y].right <> NIL {
        set t.pool[t.pool[y].right].parent = x;
    }
    set t.pool[y].parent = t.pool[x].parent;
    if t.pool[x].parent == NIL {
        set t.root = y;
    } else if x == t.pool[t.pool[x].parent].right {
        set t.pool[t.pool[x].parent].right = y;
    } else {
        set t.pool[t.pool[x].parent].left = y;
    }
    set t.pool[y].right = x;
    set t.pool[x].parent = y;
}

unsafe fn insertFixup(t: *mut RBTree, zArg: u32) {
    let mut z: u32 = zArg;
    while t.pool[t.pool[z].parent].color == RED {
        if t.pool[z].parent == t.pool[t.pool[t.pool[z].parent].parent].left {
            let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].right;
            if t.pool[y].color == RED {
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[y].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                set z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].right {
                    set z = t.pool[z].parent;
                    rotateLeft(t, z);
                }
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                rotateRight(t, t.pool[t.pool[z].parent].parent);
            }
        } else {
            let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].left;
            if t.pool[y].color == RED {
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[y].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                set z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].left {
                    set z = t.pool[z].parent;
                    rotateRight(t, z);
                }
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                rotateLeft(t, t.pool[t.pool[z].parent].parent);
            }
        }
    }
    set t.pool[t.root].color = BLACK;
}

unsafe fn insert(t: *mut RBTree, key: i32) {
    let z: u32 = allocNode(t, key);
    let mut y: u32 = NIL;
    let mut x: u32 = t.root;

    while x <> NIL {
        set y = x;
        if key < t.pool[x].key {
            set x = t.pool[x].left;
        } else {
            set x = t.pool[x].right;
        }
    }

    set t.pool[z].parent = y;
    if y == NIL {
        set t.root = z;
    } else if key < t.pool[y].key {
        set t.pool[y].left = z;
    } else {
        set t.pool[y].right = z;
    }

    insertFixup(t, z);
}

unsafe fn search(t: *RBTree, key: i32) -> bool {
    let mut x: u32 = t.root;
    while x <> NIL {
        if key == t.pool[x].key {
            return true;
        } else if key < t.pool[x].key {
            set x = t.pool[x].left;
        } else {
            set x = t.pool[x].right;
        }
    }
    return false;
}

unsafe fn inorderWalk(t: *mut RBTree, x: u32) {
    if x == NIL {
        return;
    }
    inorderWalk(t, t.pool[x].left);
    set t.inorder[t.inorderCount] = t.pool[x].key;
    set t.inorderCount += 1;
    inorderWalk(t, t.pool[x].right);
}

unsafe fn countNodes(t: *RBTree, x: u32) -> u32 {
    if x == NIL {
        return 0;
    }
    return 1 + countNodes(t, t.pool[x].left) + countNodes(t, t.pool[x].right);
}

unsafe fn blackHeight(t: *RBTree, x: u32) -> i32 {
    if x == NIL {
        return 1;
    }
    let leftBH: i32 = blackHeight(t, t.pool[x].left);
    let rightBH: i32 = blackHeight(t, t.pool[x].right);

    if leftBH == -1 or rightBH == -1 {
        return -1;
    }
    if leftBH <> rightBH {
        return -1;
    }

    if t.pool[x].color == BLACK {
        return leftBH + 1;
    }
    return leftBH;
}

unsafe fn noRedRed(t: *RBTree, x: u32) -> bool {
    if x == NIL {
        return true;
    }
    if t.pool[x].color == RED {
        if t.pool[t.pool[x].left].color == RED {
            return false;
        }
        if t.pool[t.pool[x].right].color == RED {
            return false;
        }
    }
    if not noRedRed(t, t.pool[x].left) {
        return false;
    }
    return noRedRed(t, t.pool[x].right);
}

unsafe fn resetTree(t: *mut RBTree) {
    let mut i: u32 = 0;
    while i < POOL_SIZE {
        set t.pool[i] = RBNode { key: 0, color: BLACK, left: NIL, right: NIL, parent: NIL };
        set i += 1;
    }
    set t.poolNext = 1;
    set t.root = NIL;
    set t.inorderCount = 0;
}

unsafe fn testAscending(t: *mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 0;
    while i < 32 {
        insert(t, i);
        set i += 1;
    }

    assert countNodes(t, t.root) == 32;
    assert t.pool[t.root].color == BLACK;

    let bh: i32 = blackHeight(t, t.root);
    assert bh <> -1;

    assert noRedRed(t, t.root);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    assert t.inorderCount == 32;
    let mut j: u32 = 0;
    while j < 32 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testDescending(t: *mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 31;
    while i >= 0 {
        insert(t, i);
        if i == 0 {
            break;
        }
        set i -= 1;
    }

    assert countNodes(t, t.root) == 32;
    assert blackHeight(t, t.root) <> -1;
    assert noRedRed(t, t.root);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    let mut j: u32 = 0;
    while j < 32 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testRandom(t: *mut RBTree) -> i32 {
    resetTree(t);

    let mut inserted: [bool; 48] = [false; 48];
    let mut seed: u32 = 42;
    let mut count: u32 = 0;

    while count < 48 {
        set seed = (seed * 1103515245 + 12345) & 0x7FFFFFFF;
        let val: u32 = seed % 48;
        if not inserted[val] {
            insert(t, val as i32);
            set inserted[val] = true;
            set count += 1;
        }
    }

    assert countNodes(t, t.root) == 48;
    assert blackHeight(t, t.root) <> -1;
    assert noRedRed(t, t.root);

    let mut i: i32 = 0;
    while i < 48 {
        assert search(t, i);
        set i += 1;
    }

    assert not search(t, -1);
    assert not search(t, 48);
    assert not search(t, 100);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    assert t.inorderCount == 48;
    let mut j: u32 = 0;
    while j < 48 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testHeight(t: *mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 0;
    while i < 63 {
        insert(t, i);
        set i += 1;
    }

    let bh: i32 = blackHeight(t, t.root);
    assert bh >= 3;
    assert bh <= 7;

    return 0;
}

@default unsafe fn main() -> i32 {
    let mut pool: [RBNode; 128] = [RBNode { key: 0, color: 1, left: 0, right: 0, parent: 0 }; 128];
    let mut inorder: [i32; 128] = [0; 128];

    let mut t: RBTree = RBTree {
        pool: &mut pool[..],
        poolNext: 1,
        root: NIL,
        inorder: &mut inorder[..],
        inorderCount: 0,
    };

    let r1: i32 = testAscending(&mut t);
    if r1 <> 0 { return 10 + r1; }

    let r2: i32 = testDescending(&mut t);
    if r2 <> 0 { return 20 + r2; }

    let r3: i32 = testRandom(&mut t);
    if r3 <> 0 { return 30 + r3; }

    let r4: i32 = testHeight(&mut t);
    if r4 <> 0 { return 50 + r4; }

    return 0;
}

1	//! returns: 0
2	//! Red-black tree.
3	//! Implement a red-black tree (balanced BST) using a stack-allocated node pool.
4
5	constant POOL_SIZE: u32 = 128;
6	constant NIL: u32 = 0;
7
8	constant RED: u32 = 0;
9	constant BLACK: u32 = 1;
10
11	record RBNode {
12	key: i32,
13	color: u32,
14	left: u32,
15	right: u32,
16	parent: u32,
17	}
18
19	record RBTree {
20	pool: *mut [RBNode],
21	poolNext: u32,
22	root: u32,
23	inorder: *mut [i32],
24	inorderCount: u32,
25	}
26
27	unsafe fn allocNode(t: *mut RBTree, key: i32) -> u32 {
28	let idx: u32 = t.poolNext;
29	set t.poolNext += 1;
30	set t.pool[idx] = RBNode { key, color: RED, left: NIL, right: NIL, parent: NIL };
31	return idx;
32	}
33
34	unsafe fn rotateLeft(t: *mut RBTree, x: u32) {
35	let y: u32 = t.pool[x].right;
36	set t.pool[x].right = t.pool[y].left;
37	if t.pool[y].left <> NIL {
38	set t.pool[t.pool[y].left].parent = x;
39	}
40	set t.pool[y].parent = t.pool[x].parent;
41	if t.pool[x].parent == NIL {
42	set t.root = y;
43	} else if x == t.pool[t.pool[x].parent].left {
44	set t.pool[t.pool[x].parent].left = y;
45	} else {
46	set t.pool[t.pool[x].parent].right = y;
47	}
48	set t.pool[y].left = x;
49	set t.pool[x].parent = y;
50	}
51
52	unsafe fn rotateRight(t: *mut RBTree, x: u32) {
53	let y: u32 = t.pool[x].left;
54	set t.pool[x].left = t.pool[y].right;
55	if t.pool[y].right <> NIL {
56	set t.pool[t.pool[y].right].parent = x;
57	}
58	set t.pool[y].parent = t.pool[x].parent;
59	if t.pool[x].parent == NIL {
60	set t.root = y;
61	} else if x == t.pool[t.pool[x].parent].right {
62	set t.pool[t.pool[x].parent].right = y;
63	} else {
64	set t.pool[t.pool[x].parent].left = y;
65	}
66	set t.pool[y].right = x;
67	set t.pool[x].parent = y;
68	}
69
70	unsafe fn insertFixup(t: *mut RBTree, zArg: u32) {
71	let mut z: u32 = zArg;
72	while t.pool[t.pool[z].parent].color == RED {
73	if t.pool[z].parent == t.pool[t.pool[t.pool[z].parent].parent].left {
74	let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].right;
75	if t.pool[y].color == RED {
76	set t.pool[t.pool[z].parent].color = BLACK;
77	set t.pool[y].color = BLACK;
78	set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
79	set z = t.pool[t.pool[z].parent].parent;
80	} else {
81	if z == t.pool[t.pool[z].parent].right {
82	set z = t.pool[z].parent;
83	rotateLeft(t, z);
84	}
85	set t.pool[t.pool[z].parent].color = BLACK;
86	set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
87	rotateRight(t, t.pool[t.pool[z].parent].parent);
88	}
89	} else {
90	let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].left;
91	if t.pool[y].color == RED {
92	set t.pool[t.pool[z].parent].color = BLACK;
93	set t.pool[y].color = BLACK;
94	set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
95	set z = t.pool[t.pool[z].parent].parent;
96	} else {
97	if z == t.pool[t.pool[z].parent].left {
98	set z = t.pool[z].parent;
99	rotateRight(t, z);
100	}
101	set t.pool[t.pool[z].parent].color = BLACK;
102	set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
103	rotateLeft(t, t.pool[t.pool[z].parent].parent);
104	}
105	}
106	}
107	set t.pool[t.root].color = BLACK;
108	}
109
110	unsafe fn insert(t: *mut RBTree, key: i32) {
111	let z: u32 = allocNode(t, key);
112	let mut y: u32 = NIL;
113	let mut x: u32 = t.root;
114
115	while x <> NIL {
116	set y = x;
117	if key < t.pool[x].key {
118	set x = t.pool[x].left;
119	} else {
120	set x = t.pool[x].right;
121	}
122	}
123
124	set t.pool[z].parent = y;
125	if y == NIL {
126	set t.root = z;
127	} else if key < t.pool[y].key {
128	set t.pool[y].left = z;
129	} else {
130	set t.pool[y].right = z;
131	}
132
133	insertFixup(t, z);
134	}
135
136	unsafe fn search(t: *RBTree, key: i32) -> bool {
137	let mut x: u32 = t.root;
138	while x <> NIL {
139	if key == t.pool[x].key {
140	return true;
141	} else if key < t.pool[x].key {
142	set x = t.pool[x].left;
143	} else {
144	set x = t.pool[x].right;
145	}
146	}
147	return false;
148	}
149
150	unsafe fn inorderWalk(t: *mut RBTree, x: u32) {
151	if x == NIL {
152	return;
153	}
154	inorderWalk(t, t.pool[x].left);
155	set t.inorder[t.inorderCount] = t.pool[x].key;
156	set t.inorderCount += 1;
157	inorderWalk(t, t.pool[x].right);
158	}
159
160	unsafe fn countNodes(t: *RBTree, x: u32) -> u32 {
161	if x == NIL {
162	return 0;
163	}
164	return 1 + countNodes(t, t.pool[x].left) + countNodes(t, t.pool[x].right);
165	}
166
167	unsafe fn blackHeight(t: *RBTree, x: u32) -> i32 {
168	if x == NIL {
169	return 1;
170	}
171	let leftBH: i32 = blackHeight(t, t.pool[x].left);
172	let rightBH: i32 = blackHeight(t, t.pool[x].right);
173
174	if leftBH == -1 or rightBH == -1 {
175	return -1;
176	}
177	if leftBH <> rightBH {
178	return -1;
179	}
180
181	if t.pool[x].color == BLACK {
182	return leftBH + 1;
183	}
184	return leftBH;
185	}
186
187	unsafe fn noRedRed(t: *RBTree, x: u32) -> bool {
188	if x == NIL {
189	return true;
190	}
191	if t.pool[x].color == RED {
192	if t.pool[t.pool[x].left].color == RED {
193	return false;
194	}
195	if t.pool[t.pool[x].right].color == RED {
196	return false;
197	}
198	}
199	if not noRedRed(t, t.pool[x].left) {
200	return false;
201	}
202	return noRedRed(t, t.pool[x].right);
203	}
204
205	unsafe fn resetTree(t: *mut RBTree) {
206	let mut i: u32 = 0;
207	while i < POOL_SIZE {
208	set t.pool[i] = RBNode { key: 0, color: BLACK, left: NIL, right: NIL, parent: NIL };
209	set i += 1;
210	}
211	set t.poolNext = 1;
212	set t.root = NIL;
213	set t.inorderCount = 0;
214	}
215
216	unsafe fn testAscending(t: *mut RBTree) -> i32 {
217	resetTree(t);
218
219	let mut i: i32 = 0;
220	while i < 32 {
221	insert(t, i);
222	set i += 1;
223	}
224
225	assert countNodes(t, t.root) == 32;
226	assert t.pool[t.root].color == BLACK;
227
228	let bh: i32 = blackHeight(t, t.root);
229	assert bh <> -1;
230
231	assert noRedRed(t, t.root);
232
233	set t.inorderCount = 0;
234	inorderWalk(t, t.root);
235	assert t.inorderCount == 32;
236	let mut j: u32 = 0;
237	while j < 32 {
238	assert t.inorder[j] == j as i32;
239	set j += 1;
240	}
241
242	return 0;
243	}
244
245	unsafe fn testDescending(t: *mut RBTree) -> i32 {
246	resetTree(t);
247
248	let mut i: i32 = 31;
249	while i >= 0 {
250	insert(t, i);
251	if i == 0 {
252	break;
253	}
254	set i -= 1;
255	}
256
257	assert countNodes(t, t.root) == 32;
258	assert blackHeight(t, t.root) <> -1;
259	assert noRedRed(t, t.root);
260
261	set t.inorderCount = 0;
262	inorderWalk(t, t.root);
263	let mut j: u32 = 0;
264	while j < 32 {
265	assert t.inorder[j] == j as i32;
266	set j += 1;
267	}
268
269	return 0;
270	}
271
272	unsafe fn testRandom(t: *mut RBTree) -> i32 {
273	resetTree(t);
274
275	let mut inserted: [bool; 48] = [false; 48];
276	let mut seed: u32 = 42;
277	let mut count: u32 = 0;
278
279	while count < 48 {
280	set seed = (seed * 1103515245 + 12345) & 0x7FFFFFFF;
281	let val: u32 = seed % 48;
282	if not inserted[val] {
283	insert(t, val as i32);
284	set inserted[val] = true;
285	set count += 1;
286	}
287	}
288
289	assert countNodes(t, t.root) == 48;
290	assert blackHeight(t, t.root) <> -1;
291	assert noRedRed(t, t.root);
292
293	let mut i: i32 = 0;
294	while i < 48 {
295	assert search(t, i);
296	set i += 1;
297	}
298
299	assert not search(t, -1);
300	assert not search(t, 48);
301	assert not search(t, 100);
302
303	set t.inorderCount = 0;
304	inorderWalk(t, t.root);
305	assert t.inorderCount == 48;
306	let mut j: u32 = 0;
307	while j < 48 {
308	assert t.inorder[j] == j as i32;
309	set j += 1;
310	}
311
312	return 0;
313	}
314
315	unsafe fn testHeight(t: *mut RBTree) -> i32 {
316	resetTree(t);
317
318	let mut i: i32 = 0;
319	while i < 63 {
320	insert(t, i);
321	set i += 1;
322	}
323
324	let bh: i32 = blackHeight(t, t.root);
325	assert bh >= 3;
326	assert bh <= 7;
327
328	return 0;
329	}
330
331	@default unsafe fn main() -> i32 {
332	let mut pool: [RBNode; 128] = [RBNode { key: 0, color: 1, left: 0, right: 0, parent: 0 }; 128];
333	let mut inorder: [i32; 128] = [0; 128];
334
335	let mut t: RBTree = RBTree {
336	pool: &mut pool[..],
337	poolNext: 1,
338	root: NIL,
339	inorder: &mut inorder[..],
340	inorderCount: 0,
341	};
342
343	let r1: i32 = testAscending(&mut t);
344	if r1 <> 0 { return 10 + r1; }
345
346	let r2: i32 = testDescending(&mut t);
347	if r2 <> 0 { return 20 + r2; }
348
349	let r3: i32 = testRandom(&mut t);
350	if r3 <> 0 { return 30 + r3; }
351
352	let r4: i32 = testHeight(&mut t);
353	if r4 <> 0 { return 50 + r4; }
354
355	return 0;
356	}