├── MANIFEST.in ├── tox.ini ├── Makefile ├── .travis.yml ├── crbtree ├── compat.py └── __init__.py ├── setup.py ├── UNLICENSE ├── test ├── test_sortedset.py └── test_sorteddict.py ├── .gitignore ├── README.md ├── README.rst ├── rbtree_build.py └── rb_tree ├── rb_tree.h └── rb_tree.c /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include rbtree_build.py 2 | recursive-include rb_tree * 3 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist=py27,py35,py36 3 | [testenv] 4 | usedevelop=true 5 | deps= 6 | nose 7 | cffi 8 | commands=nosetests 9 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | README.rst: README.md 2 | pandoc -f markdown -t rst < README.md > README.rst 3 | 4 | test: 5 | nosetests --with-coverage --cover-html --cover-html-dir htmlcov/ 6 | 7 | .PHONY: test 8 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | 3 | matrix: 4 | include: 5 | - python: "2.7" 6 | env: TOXENV=py27 7 | - python: "3.6" 8 | env: TOXENV=py36 9 | - python: "pypy" 10 | env: TOXENV=py27 11 | - python: "pypy3" 12 | env: TOXENV=py35 13 | 14 | install: 15 | - pip install tox 16 | 17 | script: tox 18 | 19 | notifications: 20 | email: false 21 | -------------------------------------------------------------------------------- /crbtree/compat.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import six 4 | 5 | 6 | def return_list_if_py2(method): 7 | if six.PY3: 8 | return method 9 | @functools.wraps(method) 10 | def list_method(self, *args, **kwargs): 11 | return list(method(self, *args, **kwargs)) 12 | list_method._iterator = method 13 | return list_method 14 | 15 | 16 | if six.PY3: 17 | izip = zip 18 | else: 19 | import itertools 20 | izip = itertools.izip 21 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from codecs import open 2 | from os import path 3 | 4 | from setuptools import setup, find_packages 5 | 6 | 7 | here = path.abspath(path.dirname(__file__)) 8 | 9 | # Get the long description from the README file 10 | with open(path.join(here, 'README.rst'), encoding='utf-8') as f: 11 | readme = f.read() 12 | 13 | 14 | setup( 15 | name='crbtree', 16 | version='0.0.1', 17 | 18 | description='SortedDict and SortedSet implementations, backed by a fast C red-black tree.', 19 | long_description=readme, 20 | url='https://github.com/zacharyvoase/python-crbtree', 21 | 22 | author='Zachary Voase', 23 | author_email='zack@meat.io', 24 | license='UNLICENSE', 25 | 26 | packages=find_packages(exclude=['test']), 27 | 28 | setup_requires=[ 29 | "cffi>=1.4.0", 30 | "six>=1.10.0", 31 | ], 32 | install_requires=[ 33 | "cffi>=1.4.0", 34 | "six>=1.10.0", 35 | ], 36 | 37 | cffi_modules=["rbtree_build.py:FFI_BUILDER"], 38 | ) 39 | -------------------------------------------------------------------------------- /UNLICENSE: -------------------------------------------------------------------------------- 1 | This is free and unencumbered software released into the public domain. 2 | 3 | Anyone is free to copy, modify, publish, use, compile, sell, or 4 | distribute this software, either in source code form or as a compiled 5 | binary, for any purpose, commercial or non-commercial, and by any 6 | means. 7 | 8 | In jurisdictions that recognize copyright laws, the author or authors 9 | of this software dedicate any and all copyright interest in the 10 | software to the public domain. We make this dedication for the benefit 11 | of the public at large and to the detriment of our heirs and 12 | successors. We intend this dedication to be an overt act of 13 | relinquishment in perpetuity of all present and future rights to this 14 | software under copyright law. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 19 | IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | For more information, please refer to 25 | -------------------------------------------------------------------------------- /test/test_sortedset.py: -------------------------------------------------------------------------------- 1 | from nose.tools import assert_equals 2 | from nose.tools import assert_not_equals 3 | from nose.tools import assert_raises 4 | 5 | from crbtree import SortedSet 6 | 7 | 8 | def test_smoke(): 9 | sset = SortedSet() 10 | assert_equals(len(sset), 0) 11 | assert_equals(list(sset), []) 12 | 13 | sset.add('a') 14 | assert_equals(len(sset), 1) 15 | assert_equals(list(sset), ['a']) 16 | 17 | sset.add('b') 18 | assert_equals(len(sset), 2) 19 | assert_equals(list(sset), ['a', 'b']) 20 | 21 | sset.discard('a') 22 | assert_equals(len(sset), 1) 23 | assert_equals(list(sset), ['b']) 24 | 25 | 26 | def test_initialize(): 27 | sset = SortedSet('abcdefghijkl') 28 | assert_equals(len(sset), len('abcdefghijkl')) 29 | assert_equals(list(sset), list('abcdefghijkl')) 30 | 31 | 32 | def test_equals(): 33 | sset1 = SortedSet('abc') 34 | sset2 = SortedSet('abc') 35 | sset3 = SortedSet('abcd') 36 | assert_equals(sset1, sset2) 37 | assert_not_equals(sset1, sset3) 38 | assert_not_equals(sset2, sset3) 39 | sset2.add('d') 40 | assert_not_equals(sset1, sset2) 41 | assert_equals(sset2, sset3) 42 | 43 | 44 | def test_reversed(): 45 | sset = SortedSet(['a', 'b', 'c']) 46 | assert_equals(list(reversed(sset)), ['c', 'b', 'a']) 47 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | crbtree/_rbtree.* 2 | 3 | # Finder files 4 | .DS_Store 5 | 6 | # Byte-compiled / optimized / DLL files 7 | __pycache__/ 8 | *.pyc 9 | *$py.class 10 | 11 | # C extensions 12 | *.so 13 | 14 | # Distribution / packaging 15 | .Python 16 | env/ 17 | build/ 18 | develop-eggs/ 19 | dist/ 20 | downloads/ 21 | eggs/ 22 | .eggs/ 23 | lib/ 24 | lib64/ 25 | parts/ 26 | sdist/ 27 | var/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *,cover 51 | .hypothesis/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | 61 | # Flask stuff: 62 | instance/ 63 | .webassets-cache 64 | 65 | # Scrapy stuff: 66 | .scrapy 67 | 68 | # Sphinx documentation 69 | docs/_build/ 70 | 71 | # PyBuilder 72 | target/ 73 | 74 | # IPython Notebook 75 | .ipynb_checkpoints 76 | 77 | # pyenv 78 | .python-version 79 | 80 | # celery beat schedule file 81 | celerybeat-schedule 82 | 83 | # dotenv 84 | .env 85 | 86 | # virtualenv 87 | .venv/ 88 | venv/ 89 | ENV/ 90 | 91 | # Spyder project settings 92 | .spyderproject 93 | 94 | # Rope project settings 95 | .ropeproject 96 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # cRBTree 2 | 3 | A `SortedDict` and `SortedSet` implementation, backed by a [red-black 4 | tree][rbtree] implementation in C, wrapped with [CFFI][]. 5 | 6 | `SortedDict` and `SortedSet` are collections that always iterate through their 7 | keys/contents in order. Usage is simple: 8 | 9 | >>> sd = SortedDict() 10 | >>> sd['c'] = 789 11 | >>> sd['b'] = 456 12 | >>> sd['a'] = 123 13 | >>> for key, value in sd.iteritems(): 14 | ... print((key, value)) 15 | ('a', 123) 16 | ('b', 456) 17 | ('c', 789) 18 | 19 | You can iterate in reverse order by getting a reversed view of the underlying `SortedDict` or `Set`: 20 | 21 | >>> rsd = reversed(sd) 22 | >>> for key, value in sd.iteritems(): 23 | ... print((key, value)) 24 | ('c', 789) 25 | ('b', 456) 26 | ('a', 123) 27 | 28 | 29 | [rbtree]: https://en.wikipedia.org/wiki/Red%E2%80%93black_tree 30 | [cffi]: https://cffi.readthedocs.org/ 31 | 32 | 33 | ## Installation 34 | 35 | pip install crbtree 36 | 37 | 38 | ## Unlicense 39 | 40 | This is free and unencumbered software released into the public domain. 41 | 42 | Anyone is free to copy, modify, publish, use, compile, sell, or distribute this 43 | software, either in source code form or as a compiled binary, for any purpose, 44 | commercial or non-commercial, and by any means. 45 | 46 | In jurisdictions that recognize copyright laws, the author or authors of this 47 | software dedicate any and all copyright interest in the software to the public 48 | domain. We make this dedication for the benefit of the public at large and to 49 | the detriment of our heirs and successors. We intend this dedication to be an 50 | overt act of relinquishment in perpetuity of all present and future rights to 51 | this software under copyright law. 52 | 53 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 54 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 55 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 56 | AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN 57 | ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 58 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 59 | 60 | For more information, please refer to 61 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | cRBTree 2 | ======= 3 | 4 | A ``SortedDict`` and ``SortedSet`` implementation, backed by a 5 | `red-black 6 | tree `__ 7 | implementation in C, wrapped with 8 | `CFFI `__. 9 | 10 | ``SortedDict`` and ``SortedSet`` are collections that always iterate 11 | through their keys/contents in order. Usage is simple: 12 | 13 | :: 14 | 15 | >>> sd = SortedDict() 16 | >>> sd['c'] = 789 17 | >>> sd['b'] = 456 18 | >>> sd['a'] = 123 19 | >>> for key, value in sd.iteritems(): 20 | ... print((key, value)) 21 | ('a', 123) 22 | ('b', 456) 23 | ('c', 789) 24 | 25 | You can iterate in reverse order by getting a reversed view of the 26 | underlying ``SortedDict`` or ``Set``: 27 | 28 | :: 29 | 30 | >>> rsd = reversed(sd) 31 | >>> for key, value in sd.iteritems(): 32 | ... print((key, value)) 33 | ('c', 789) 34 | ('b', 456) 35 | ('a', 123) 36 | 37 | Installation 38 | ------------ 39 | 40 | :: 41 | 42 | pip install crbtree 43 | 44 | Unlicense 45 | --------- 46 | 47 | This is free and unencumbered software released into the public domain. 48 | 49 | Anyone is free to copy, modify, publish, use, compile, sell, or 50 | distribute this software, either in source code form or as a compiled 51 | binary, for any purpose, commercial or non-commercial, and by any means. 52 | 53 | In jurisdictions that recognize copyright laws, the author or authors of 54 | this software dedicate any and all copyright interest in the software to 55 | the public domain. We make this dedication for the benefit of the public 56 | at large and to the detriment of our heirs and successors. We intend 57 | this dedication to be an overt act of relinquishment in perpetuity of 58 | all present and future rights to this software under copyright law. 59 | 60 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 61 | OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 62 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 63 | IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 64 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 65 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 66 | DEALINGS IN THE SOFTWARE. 67 | 68 | For more information, please refer to http://unlicense.org/ 69 | -------------------------------------------------------------------------------- /rbtree_build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | 4 | from cffi import FFI 5 | import six 6 | 7 | 8 | SRC_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'rb_tree') 9 | 10 | FFI_BUILDER = FFI() 11 | 12 | MACROS = { 13 | 'RB_ITER_MAX_HEIGHT': '64' 14 | } 15 | 16 | HEADER = ''' 17 | extern "Python" int rb_tree_node_compare(struct rb_tree *self, struct rb_node *a, struct rb_node *b); 18 | extern "Python" void rb_tree_node_was_removed(struct rb_tree *self, struct rb_node *node); 19 | typedef int (*rb_tree_node_cmp_f) (struct rb_tree *self, struct rb_node *a, struct rb_node *b); 20 | typedef void (*rb_tree_node_f) (struct rb_tree *self, struct rb_node *node); 21 | 22 | struct rb_node { 23 | int red; // Color red (1), black (0) 24 | struct rb_node *link[2]; // Link left [0] and right [1] 25 | void *value; // User provided, used indirectly via rb_tree_node_cmp_f. 26 | }; 27 | 28 | struct rb_tree { 29 | struct rb_node *root; 30 | rb_tree_node_cmp_f cmp; 31 | size_t size; 32 | void *info; // User provided, not used by rb_tree. 33 | }; 34 | 35 | struct rb_iter { 36 | struct rb_tree *tree; 37 | struct rb_node *node; // Current node 38 | struct rb_node *path[RB_ITER_MAX_HEIGHT]; // Traversal path 39 | size_t top; // Top of stack 40 | void *info; // User provided, not used by rb_iter. 41 | }; 42 | 43 | int rb_tree_node_cmp_ptr_cb (struct rb_tree *self, struct rb_node *a, struct rb_node *b); 44 | void rb_tree_node_dealloc_cb (struct rb_tree *self, struct rb_node *node); 45 | 46 | struct rb_node *rb_node_alloc (void); 47 | struct rb_node *rb_node_create (void *value); 48 | struct rb_node *rb_node_init (struct rb_node *self, void *value); 49 | void rb_node_dealloc (struct rb_node *self); 50 | 51 | struct rb_tree *rb_tree_alloc (void); 52 | struct rb_tree *rb_tree_create (rb_tree_node_cmp_f cmp); 53 | struct rb_tree *rb_tree_init (struct rb_tree *self, rb_tree_node_cmp_f cmp); 54 | void rb_tree_dealloc (struct rb_tree *self, rb_tree_node_f node_cb); 55 | void *rb_tree_find (struct rb_tree *self, void *value); 56 | int rb_tree_insert (struct rb_tree *self, void *value); 57 | int rb_tree_remove (struct rb_tree *self, void *value); 58 | size_t rb_tree_size (struct rb_tree *self); 59 | 60 | int rb_tree_insert_node (struct rb_tree *self, struct rb_node *node); 61 | int rb_tree_remove_with_cb (struct rb_tree *self, void *value, rb_tree_node_f node_cb); 62 | 63 | int rb_tree_test (struct rb_tree *self, struct rb_node *root); 64 | 65 | struct rb_iter *rb_iter_alloc (void); 66 | struct rb_iter *rb_iter_init (struct rb_iter *self); 67 | struct rb_iter *rb_iter_create (void); 68 | void rb_iter_dealloc (struct rb_iter *self); 69 | void *rb_iter_first (struct rb_iter *self, struct rb_tree *tree); 70 | void *rb_iter_last (struct rb_iter *self, struct rb_tree *tree); 71 | void *rb_iter_next (struct rb_iter *self); 72 | void *rb_iter_prev (struct rb_iter *self); 73 | ''' 74 | 75 | for macro_name, macro_value in six.iteritems(MACROS): 76 | HEADER = re.sub(r'\b{}\b'.format(re.escape(macro_name)), re.escape(macro_value), HEADER) 77 | 78 | FFI_BUILDER.cdef(HEADER) 79 | FFI_BUILDER.set_source('crbtree._rbtree', ''' 80 | #define RB_ITER_MAX_HEIGHT 64 81 | #include "rb_tree.c" 82 | ''', include_dirs=[SRC_ROOT]) 83 | 84 | if __name__ == '__main__': 85 | FFI_BUILDER.compile() 86 | -------------------------------------------------------------------------------- /test/test_sorteddict.py: -------------------------------------------------------------------------------- 1 | from nose.tools import assert_equals 2 | from nose.tools import assert_not_equals 3 | from nose.tools import assert_raises 4 | import six 5 | 6 | from crbtree import SortedDict 7 | 8 | 9 | def test_smoke(): 10 | sdict = SortedDict() 11 | assert_equals(len(sdict), 0) 12 | assert_equals(list(sdict), []) 13 | assert 'a' not in sdict 14 | assert 'b' not in sdict 15 | 16 | sdict['a'] = 123 17 | assert_equals(sdict['a'], 123) 18 | assert_equals(len(sdict), 1) 19 | assert 'a' in sdict 20 | assert_equals(list(sdict), ['a']) 21 | 22 | sdict['b'] = 456 23 | assert_equals(sdict['b'], 456) 24 | assert_equals(len(sdict), 2) 25 | assert 'b' in sdict 26 | assert_equals(list(sdict), ['a', 'b']) 27 | 28 | sdict['a'] = 789 29 | assert_equals(sdict['a'], 789) 30 | assert_equals(len(sdict), 2) 31 | assert 'a' in sdict 32 | assert 'b' in sdict 33 | assert_equals(list(sdict), ['a', 'b']) 34 | 35 | del sdict['b'] 36 | assert_equals(len(sdict), 1) 37 | assert 'a' in sdict 38 | assert 'b' not in sdict 39 | assert_equals(list(sdict), ['a']) 40 | with assert_raises(KeyError): 41 | sdict['b'] 42 | with assert_raises(KeyError): 43 | del sdict['b'] 44 | del sdict['b'] 45 | del sdict['b'] 46 | 47 | 48 | def test_initialize(): 49 | sdict = SortedDict(a=123, b=456) 50 | assert_equals(len(sdict), 2) 51 | assert_equals(sdict['a'], 123) 52 | assert_equals(sdict['b'], 456) 53 | 54 | sdict = SortedDict([('a', 123), ('b', 456)]) 55 | assert_equals(len(sdict), 2) 56 | assert_equals(sdict['a'], 123) 57 | assert_equals(sdict['b'], 456) 58 | 59 | sdict = SortedDict([('a', 123), ('b', 456)], c=789) 60 | assert_equals(len(sdict), 3) 61 | assert_equals(sdict['a'], 123) 62 | assert_equals(sdict['b'], 456) 63 | assert_equals(sdict['c'], 789) 64 | 65 | sdict = SortedDict({'a': 123, 'b': 456}) 66 | assert_equals(len(sdict), 2) 67 | assert_equals(sdict['a'], 123) 68 | assert_equals(sdict['b'], 456) 69 | 70 | with assert_raises(TypeError): 71 | SortedDict('abc') 72 | 73 | with assert_raises(TypeError): 74 | SortedDict('key', 'value') 75 | 76 | with assert_raises(TypeError): 77 | SortedDict(123) 78 | 79 | 80 | def test_equals(): 81 | sdict1 = SortedDict(a=123, b=456) 82 | sdict2 = SortedDict(a=123, b=456) 83 | sdict3 = SortedDict(a=123, b=456, c=789) 84 | assert_equals(sdict1, sdict2) 85 | assert_not_equals(sdict1, sdict3) 86 | assert_not_equals(sdict2, sdict3) 87 | assert_not_equals(sdict3, {'a': 123, 'b': 456, 'c': 789}) 88 | 89 | 90 | def test_memory_management(): 91 | sdict = SortedDict() 92 | assert_equals(len(sdict._handles), 0) 93 | sdict['a'] = 123 94 | assert_equals(len(sdict._handles), 1) 95 | del sdict['a'] 96 | assert_equals(len(sdict._handles), 0) 97 | sdict['a'] = 456 98 | first_handles = list(sdict._handles) 99 | sdict['a'] = 789 100 | second_handles = list(sdict._handles) 101 | assert_not_equals(first_handles, second_handles) 102 | 103 | 104 | def test_ordering_maintained(): 105 | sdict = SortedDict() 106 | for i, char in enumerate('abcdefghijkl'): 107 | sdict[char] = i 108 | 109 | assert_equals(list(sdict), list('abcdefghijkl')) 110 | 111 | 112 | def test_reversed(): 113 | sdict = SortedDict(a=123, b=456, c=789) 114 | rsdict = reversed(sdict) 115 | rsdict2 = reversed(sdict) 116 | 117 | assert 'a' in rsdict 118 | assert_equals(rsdict['a'], 123) 119 | 120 | assert_equals(len(rsdict), len(sdict)) 121 | assert_equals(list(rsdict), ['c', 'b', 'a']) 122 | assert_equals(list(six.iterkeys(rsdict)), ['c', 'b', 'a']) 123 | assert_equals(list(six.itervalues(rsdict)), [789, 456, 123]) 124 | assert_equals(list(six.iteritems(rsdict)), [('c', 789), ('b', 456), ('a', 123)]) 125 | 126 | assert_equals(list(rsdict.keys()), ['c', 'b', 'a']) 127 | assert_equals(list(rsdict.values()), [789, 456, 123]) 128 | assert_equals(list(rsdict.items()), [('c', 789), ('b', 456), ('a', 123)]) 129 | 130 | assert_not_equals(sdict, rsdict) 131 | assert_not_equals(rsdict, sdict) 132 | assert_not_equals(rsdict, None) 133 | assert_equals(rsdict, rsdict2) 134 | 135 | assert_equals(reversed(rsdict), sdict) 136 | assert_equals(sdict, reversed(rsdict)) 137 | 138 | 139 | def test_reversed_equal(): 140 | sdict = SortedDict() 141 | rsdict = reversed(sdict) 142 | 143 | assert_equals(sdict, rsdict) 144 | assert_equals(rsdict, sdict) 145 | -------------------------------------------------------------------------------- /rb_tree/rb_tree.h: -------------------------------------------------------------------------------- 1 | // 2 | // Based on Julienne Walker's rb_tree 3 | // implementation. 4 | // 5 | // Modified by Mirek Rusin . 6 | // 7 | // This is free and unencumbered software released into the public domain. 8 | // 9 | // Anyone is free to copy, modify, publish, use, compile, sell, or 10 | // distribute this software, either in source code form or as a compiled 11 | // binary, for any purpose, commercial or non-commercial, and by any 12 | // means. 13 | // 14 | // In jurisdictions that recognize copyright laws, the author or authors 15 | // of this software dedicate any and all copyright interest in the 16 | // software to the public domain. We make this dedication for the benefit 17 | // of the public at large and to the detriment of our heirs and 18 | // successors. We intend this dedication to be an overt act of 19 | // relinquishment in perpetuity of all present and future rights to this 20 | // software under copyright law. 21 | // 22 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 23 | // EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 24 | // MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 25 | // IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 26 | // OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 27 | // ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 28 | // OTHER DEALINGS IN THE SOFTWARE. 29 | // 30 | // For more information, please refer to 31 | // 32 | 33 | #ifndef __RB_TREE_H__ 34 | #define __RB_TREE_H__ 1 35 | 36 | #include 37 | #include 38 | #include 39 | #include 40 | 41 | #ifndef RB_ITER_MAX_HEIGHT 42 | #define RB_ITER_MAX_HEIGHT 64 // Tallest allowable tree to iterate 43 | #endif 44 | 45 | struct rb_node; 46 | struct rb_tree; 47 | 48 | typedef int (*rb_tree_node_cmp_f) (struct rb_tree *self, struct rb_node *a, struct rb_node *b); 49 | typedef void (*rb_tree_node_f) (struct rb_tree *self, struct rb_node *node); 50 | 51 | struct rb_node { 52 | int red; // Color red (1), black (0) 53 | struct rb_node *link[2]; // Link left [0] and right [1] 54 | void *value; // User provided, used indirectly via rb_tree_node_cmp_f. 55 | }; 56 | 57 | struct rb_tree { 58 | struct rb_node *root; 59 | rb_tree_node_cmp_f cmp; 60 | size_t size; 61 | void *info; // User provided, not used by rb_tree. 62 | }; 63 | 64 | struct rb_iter { 65 | struct rb_tree *tree; 66 | struct rb_node *node; // Current node 67 | struct rb_node *path[RB_ITER_MAX_HEIGHT]; // Traversal path 68 | size_t top; // Top of stack 69 | void *info; // User provided, not used by rb_iter. 70 | }; 71 | 72 | int rb_tree_node_cmp_ptr_cb (struct rb_tree *self, struct rb_node *a, struct rb_node *b); 73 | void rb_tree_node_dealloc_cb (struct rb_tree *self, struct rb_node *node); 74 | 75 | struct rb_node *rb_node_alloc (void); 76 | struct rb_node *rb_node_create (void *value); 77 | struct rb_node *rb_node_init (struct rb_node *self, void *value); 78 | void rb_node_dealloc (struct rb_node *self); 79 | 80 | struct rb_tree *rb_tree_alloc (void); 81 | struct rb_tree *rb_tree_create (rb_tree_node_cmp_f cmp); 82 | struct rb_tree *rb_tree_init (struct rb_tree *self, rb_tree_node_cmp_f cmp); 83 | void rb_tree_dealloc (struct rb_tree *self, rb_tree_node_f node_cb); 84 | void *rb_tree_find (struct rb_tree *self, void *value); 85 | int rb_tree_insert (struct rb_tree *self, void *value); 86 | int rb_tree_remove (struct rb_tree *self, void *value); 87 | size_t rb_tree_size (struct rb_tree *self); 88 | 89 | int rb_tree_insert_node (struct rb_tree *self, struct rb_node *node); 90 | int rb_tree_remove_with_cb (struct rb_tree *self, void *value, rb_tree_node_f node_cb); 91 | 92 | int rb_tree_test (struct rb_tree *self, struct rb_node *root); 93 | 94 | struct rb_iter *rb_iter_alloc (void); 95 | struct rb_iter *rb_iter_init (struct rb_iter *self); 96 | struct rb_iter *rb_iter_create (void); 97 | void rb_iter_dealloc (struct rb_iter *self); 98 | void *rb_iter_first (struct rb_iter *self, struct rb_tree *tree); 99 | void *rb_iter_last (struct rb_iter *self, struct rb_tree *tree); 100 | void *rb_iter_next (struct rb_iter *self); 101 | void *rb_iter_prev (struct rb_iter *self); 102 | 103 | #endif 104 | -------------------------------------------------------------------------------- /crbtree/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | A Dict and Set implementation which always iterate in order. 3 | 4 | `SortedDict` and `SortedSet` are red-black tree-based collections, which store 5 | their keys according to their native Python sort order. This means iteration 6 | (i.e. `for key in dictionary:`, or `for value in set:`) always produces the 7 | keys in order. 8 | """ 9 | 10 | import collections 11 | import six 12 | 13 | from crbtree._rbtree import ffi, lib 14 | from crbtree import compat 15 | 16 | 17 | __all__ = ['SortedDict', 'SortedSet'] 18 | 19 | 20 | Item = collections.namedtuple('Item', ('key', 'value')) 21 | 22 | 23 | class SortedDict(collections.MutableMapping): 24 | 25 | "A sorted dictionary, backed by a red-black tree." 26 | 27 | def __init__(self, *args, **kwargs): 28 | self._rbtree = lib.rb_tree_create(lib.rb_tree_node_compare) 29 | # This allows us to get the SortedDict Python object from a node 30 | # removal/dealloc callback. 31 | self._self_handle = ffi.new_handle(self) 32 | self._rbtree.info = self._self_handle 33 | # Track the FFI pointers to Items so they don't get garbage collected. 34 | self._handles = set() 35 | if args: 36 | if len(args) != 1: 37 | raise TypeError("SortedDict() expected exactly 0 or 1 positional args, got {}" 38 | .format(len(args))) 39 | if isinstance(args[0], collections.Mapping): 40 | self.update(args[0]) 41 | elif isinstance(args[0], collections.Iterable): 42 | for item in args[0]: 43 | try: 44 | key, value = item 45 | except ValueError: 46 | raise TypeError("SortedDict expected (key, value) pair, got {!r}".format(item)) 47 | self[key] = value 48 | else: 49 | raise TypeError("Cannot initialize SortedDict from {!r}".format(args[0])) 50 | if kwargs: 51 | self.update(kwargs) 52 | 53 | def __del__(self): 54 | lib.rb_tree_dealloc(self._rbtree, ffi.addressof(lib, 'rb_tree_node_dealloc_cb')) 55 | 56 | def __len__(self): 57 | return lib.rb_tree_size(self._rbtree) 58 | 59 | def _get(self, key): 60 | item = Item(key, None) 61 | item_p = ffi.new_handle(item) 62 | result_p = lib.rb_tree_find(self._rbtree, item_p) 63 | if result_p == ffi.NULL: 64 | return (False, None) 65 | return (True, ffi.from_handle(result_p).value) 66 | 67 | def __contains__(self, key): 68 | return self._get(key)[0] 69 | 70 | def __setitem__(self, key, value): 71 | if key in self: 72 | del self[key] 73 | item = Item(key, value) 74 | item_p = ffi.new_handle(item) 75 | self._handles.add(item_p) 76 | if not lib.rb_tree_insert(self._rbtree, item_p): 77 | raise RuntimeError("Unexpected error inserting key {!r}".format(key)) 78 | 79 | def __getitem__(self, key): 80 | found, item = self._get(key) 81 | if found: 82 | return item 83 | raise KeyError(key) 84 | 85 | def __delitem__(self, key): 86 | if key not in self: 87 | raise KeyError(key) 88 | item = Item(key, None) 89 | item_p = ffi.new_handle(item) 90 | removed = lib.rb_tree_remove_with_cb(self._rbtree, item_p, lib.rb_tree_node_was_removed) 91 | if not removed: 92 | raise RuntimeError("Unexpected error removing key {!r}".format(key)) 93 | 94 | def __iter__(self): 95 | for key, value in six.iteritems(self): 96 | yield key 97 | 98 | def __eq__(self, other): 99 | if isinstance(other, ReversedSortedDictView): 100 | return len(self) < 2 and len(self) == len(other) and sorted_mapping_eq(self, other) 101 | elif not isinstance(other, SortedDict): 102 | return False 103 | return len(self) == len(other) and sorted_mapping_eq(self, other) 104 | 105 | def __reversed__(self): 106 | return ReversedSortedDictView(self) 107 | 108 | @compat.return_list_if_py2 109 | def keys(self): 110 | return iter(self) 111 | 112 | @compat.return_list_if_py2 113 | def values(self): 114 | for _, value in six.iteritems(self): 115 | yield value 116 | 117 | @compat.return_list_if_py2 118 | def items(self): 119 | rb_iter = lib.rb_iter_create() 120 | try: 121 | item_p = lib.rb_iter_first(rb_iter, self._rbtree) 122 | while item_p != ffi.NULL: 123 | item = ffi.from_handle(item_p) 124 | yield (item.key, item.value) 125 | item_p = lib.rb_iter_next(rb_iter) 126 | finally: 127 | lib.rb_iter_dealloc(rb_iter) 128 | 129 | if six.PY2: 130 | iterkeys = keys._iterator 131 | itervalues = values._iterator 132 | iteritems = items._iterator 133 | 134 | 135 | class ReversedSortedDictView(object): 136 | __slots__ = ('sorted_dict',) 137 | 138 | def __init__(self, sorted_dict): 139 | self.sorted_dict = sorted_dict 140 | 141 | def __len__(self): 142 | return len(self.sorted_dict) 143 | 144 | def __contains__(self, key): 145 | return key in self.sorted_dict 146 | 147 | def __getitem__(self, key): 148 | return self.sorted_dict[key] 149 | 150 | def __iter__(self): 151 | for key, _ in six.iteritems(self): 152 | yield key 153 | 154 | def __eq__(self, other): 155 | if isinstance(other, ReversedSortedDictView): 156 | return self.sorted_dict == other.sorted_dict 157 | elif isinstance(other, SortedDict): 158 | return len(self) < 2 and len(self) == len(other) and sorted_mapping_eq(self, other) 159 | return False 160 | 161 | def __reversed__(self): 162 | return self.sorted_dict 163 | 164 | @compat.return_list_if_py2 165 | def keys(self): 166 | return iter(self) 167 | 168 | @compat.return_list_if_py2 169 | def values(self): 170 | for _, value in six.iteritems(self): 171 | yield value 172 | 173 | @compat.return_list_if_py2 174 | def items(self): 175 | rb_iter = lib.rb_iter_create() 176 | try: 177 | item_p = lib.rb_iter_last(rb_iter, self.sorted_dict._rbtree) 178 | while item_p != ffi.NULL: 179 | item = ffi.from_handle(item_p) 180 | yield (item.key, item.value) 181 | item_p = lib.rb_iter_prev(rb_iter) 182 | finally: 183 | lib.rb_iter_dealloc(rb_iter) 184 | 185 | if six.PY2: 186 | iterkeys = keys._iterator 187 | itervalues = values._iterator 188 | iteritems = items._iterator 189 | 190 | 191 | class SortedSet(collections.MutableSet): 192 | 193 | def __init__(self, iterable=None): 194 | if iterable is not None: 195 | self._dict = SortedDict((value, None) for value in iterable) 196 | else: 197 | self._dict = SortedDict() 198 | 199 | def __contains__(self, value): 200 | return value in self._dict 201 | 202 | def __len__(self): 203 | return len(self._dict) 204 | 205 | def __iter__(self): 206 | return six.iterkeys(self._dict) 207 | 208 | def __reversed__(self): 209 | return six.iterkeys(reversed(self._dict)) 210 | 211 | def add(self, value): 212 | self._dict.setdefault(value, None) 213 | 214 | def discard(self, value): 215 | self._dict.pop(value, None) 216 | 217 | 218 | @ffi.def_extern() 219 | def rb_tree_node_compare(rb_tree_p, rb_node_a, rb_node_b): 220 | a, b = ffi.from_handle(rb_node_a.value), ffi.from_handle(rb_node_b.value) 221 | if a.key == b.key: 222 | return 0 223 | if a.key < b.key: 224 | return -1 225 | return 1 226 | 227 | 228 | @ffi.def_extern() 229 | def rb_tree_node_was_removed(rb_tree_p, rb_node_p): 230 | ffi.from_handle(rb_tree_p.info)._handles.discard(rb_node_p.value) 231 | lib.rb_tree_node_dealloc_cb(rb_tree_p, rb_node_p) 232 | 233 | 234 | def sorted_mapping_eq(map1, map2): 235 | return all( 236 | k1 == k2 and v1 == v2 237 | for (k1, v1), (k2, v2) 238 | in compat.izip(six.iteritems(map1), six.iteritems(map2))) 239 | -------------------------------------------------------------------------------- /rb_tree/rb_tree.c: -------------------------------------------------------------------------------- 1 | // 2 | // Based on Julienne Walker's rb_tree 3 | // implementation. 4 | // 5 | // Modified by Mirek Rusin . 6 | // 7 | // This is free and unencumbered software released into the public domain. 8 | // 9 | // Anyone is free to copy, modify, publish, use, compile, sell, or 10 | // distribute this software, either in source code form or as a compiled 11 | // binary, for any purpose, commercial or non-commercial, and by any 12 | // means. 13 | // 14 | // In jurisdictions that recognize copyright laws, the author or authors 15 | // of this software dedicate any and all copyright interest in the 16 | // software to the public domain. We make this dedication for the benefit 17 | // of the public at large and to the detriment of our heirs and 18 | // successors. We intend this dedication to be an overt act of 19 | // relinquishment in perpetuity of all present and future rights to this 20 | // software under copyright law. 21 | // 22 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 23 | // EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 24 | // MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 25 | // IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 26 | // OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 27 | // ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 28 | // OTHER DEALINGS IN THE SOFTWARE. 29 | // 30 | // For more information, please refer to 31 | // 32 | 33 | #include "rb_tree.h" 34 | 35 | // rb_node 36 | 37 | struct rb_node * 38 | rb_node_alloc () { 39 | return malloc(sizeof(struct rb_node)); 40 | } 41 | 42 | struct rb_node * 43 | rb_node_init (struct rb_node *self, void *value) { 44 | if (self) { 45 | self->red = 1; 46 | self->link[0] = self->link[1] = NULL; 47 | self->value = value; 48 | } 49 | return self; 50 | } 51 | 52 | struct rb_node * 53 | rb_node_create (void *value) { 54 | return rb_node_init(rb_node_alloc(), value); 55 | } 56 | 57 | void 58 | rb_node_dealloc (struct rb_node *self) { 59 | if (self) { 60 | free(self); 61 | } 62 | } 63 | 64 | static int 65 | rb_node_is_red (const struct rb_node *self) { 66 | return self ? self->red : 0; 67 | } 68 | 69 | static struct rb_node * 70 | rb_node_rotate (struct rb_node *self, int dir) { 71 | struct rb_node *result = NULL; 72 | if (self) { 73 | result = self->link[!dir]; 74 | self->link[!dir] = result->link[dir]; 75 | result->link[dir] = self; 76 | self->red = 1; 77 | result->red = 0; 78 | } 79 | return result; 80 | } 81 | 82 | static struct rb_node * 83 | rb_node_rotate2 (struct rb_node *self, int dir) { 84 | struct rb_node *result = NULL; 85 | if (self) { 86 | self->link[!dir] = rb_node_rotate(self->link[!dir], !dir); 87 | result = rb_node_rotate(self, dir); 88 | } 89 | return result; 90 | } 91 | 92 | // rb_tree - default callbacks 93 | 94 | int 95 | rb_tree_node_cmp_ptr_cb (struct rb_tree *self, struct rb_node *a, struct rb_node *b) { 96 | return (a->value > b->value) - (a->value < b->value); 97 | } 98 | 99 | void 100 | rb_tree_node_dealloc_cb (struct rb_tree *self, struct rb_node *node) { 101 | if (self) { 102 | if (node) { 103 | rb_node_dealloc(node); 104 | } 105 | } 106 | } 107 | 108 | // rb_tree 109 | 110 | struct rb_tree * 111 | rb_tree_alloc () { 112 | return malloc(sizeof(struct rb_tree)); 113 | } 114 | 115 | struct rb_tree * 116 | rb_tree_init (struct rb_tree *self, rb_tree_node_cmp_f node_cmp_cb) { 117 | if (self) { 118 | self->root = NULL; 119 | self->size = 0; 120 | self->cmp = node_cmp_cb ? node_cmp_cb : rb_tree_node_cmp_ptr_cb; 121 | } 122 | return self; 123 | } 124 | 125 | struct rb_tree * 126 | rb_tree_create (rb_tree_node_cmp_f node_cb) { 127 | return rb_tree_init(rb_tree_alloc(), node_cb); 128 | } 129 | 130 | void 131 | rb_tree_dealloc (struct rb_tree *self, rb_tree_node_f node_cb) { 132 | if (self) { 133 | if (node_cb) { 134 | struct rb_node *node = self->root; 135 | struct rb_node *save = NULL; 136 | 137 | // Rotate away the left links so that 138 | // we can treat this like the destruction 139 | // of a linked list 140 | while (node) { 141 | if (node->link[0] == NULL) { 142 | 143 | // No left links, just kill the node and move on 144 | save = node->link[1]; 145 | node_cb(self, node); 146 | node = NULL; 147 | } else { 148 | 149 | // Rotate away the left link and check again 150 | save = node->link[0]; 151 | node->link[0] = save->link[1]; 152 | save->link[1] = node; 153 | } 154 | node = save; 155 | } 156 | } 157 | free(self); 158 | } 159 | } 160 | 161 | int 162 | rb_tree_test (struct rb_tree *self, struct rb_node *root) { 163 | int lh, rh; 164 | 165 | if ( root == NULL ) 166 | return 1; 167 | else { 168 | struct rb_node *ln = root->link[0]; 169 | struct rb_node *rn = root->link[1]; 170 | 171 | /* Consecutive red links */ 172 | if (rb_node_is_red(root)) { 173 | if (rb_node_is_red(ln) || rb_node_is_red(rn)) { 174 | printf("Red violation"); 175 | return 0; 176 | } 177 | } 178 | 179 | lh = rb_tree_test(self, ln); 180 | rh = rb_tree_test(self, rn); 181 | 182 | /* Invalid binary search tree */ 183 | if ( ( ln != NULL && self->cmp(self, ln, root) >= 0 ) 184 | || ( rn != NULL && self->cmp(self, rn, root) <= 0)) 185 | { 186 | puts ( "Binary tree violation" ); 187 | return 0; 188 | } 189 | 190 | /* Black height mismatch */ 191 | if ( lh != 0 && rh != 0 && lh != rh ) { 192 | puts ( "Black violation" ); 193 | return 0; 194 | } 195 | 196 | /* Only count black links */ 197 | if ( lh != 0 && rh != 0 ) 198 | return rb_node_is_red ( root ) ? lh : lh + 1; 199 | else 200 | return 0; 201 | } 202 | } 203 | 204 | void * 205 | rb_tree_find(struct rb_tree *self, void *value) { 206 | void *result = NULL; 207 | if (self) { 208 | struct rb_node node = { .value = value }; 209 | struct rb_node *it = self->root; 210 | int cmp = 0; 211 | while (it) { 212 | if ((cmp = self->cmp(self, it, &node))) { 213 | 214 | // If the tree supports duplicates, they should be 215 | // chained to the right subtree for this to work 216 | it = it->link[cmp < 0]; 217 | } else { 218 | break; 219 | } 220 | } 221 | result = it ? it->value : NULL; 222 | } 223 | return result; 224 | } 225 | 226 | // Creates (malloc'ates) 227 | int 228 | rb_tree_insert (struct rb_tree *self, void *value) { 229 | return rb_tree_insert_node(self, rb_node_create(value)); 230 | } 231 | 232 | // Returns 1 on success, 0 otherwise. 233 | int 234 | rb_tree_insert_node (struct rb_tree *self, struct rb_node *node) { 235 | int result = 0; 236 | if (self && node) { 237 | if (self->root == NULL) { 238 | self->root = node; 239 | result = 1; 240 | } else { 241 | struct rb_node head = { 0 }; // False tree root 242 | struct rb_node *g, *t; // Grandparent & parent 243 | struct rb_node *p, *q; // Iterator & parent 244 | int dir = 0, last = 0; 245 | 246 | // Set up our helpers 247 | t = &head; 248 | g = p = NULL; 249 | q = t->link[1] = self->root; 250 | 251 | // Search down the tree for a place to insert 252 | while (1) { 253 | if (q == NULL) { 254 | 255 | // Insert node at the first null link. 256 | p->link[dir] = q = node; 257 | } else if (rb_node_is_red(q->link[0]) && rb_node_is_red(q->link[1])) { 258 | 259 | // Simple red violation: color flip 260 | q->red = 1; 261 | q->link[0]->red = 0; 262 | q->link[1]->red = 0; 263 | } 264 | 265 | if (rb_node_is_red(q) && rb_node_is_red(p)) { 266 | 267 | // Hard red violation: rotations necessary 268 | int dir2 = t->link[1] == g; 269 | if (q == p->link[last]) { 270 | t->link[dir2] = rb_node_rotate(g, !last); 271 | } else { 272 | t->link[dir2] = rb_node_rotate2(g, !last); 273 | } 274 | } 275 | 276 | // Stop working if we inserted a node. This 277 | // check also disallows duplicates in the tree 278 | if (self->cmp(self, q, node) == 0) { 279 | break; 280 | } 281 | 282 | last = dir; 283 | dir = self->cmp(self, q, node) < 0; 284 | 285 | // Move the helpers down 286 | if (g != NULL) { 287 | t = g; 288 | } 289 | 290 | g = p, p = q; 291 | q = q->link[dir]; 292 | } 293 | 294 | // Update the root (it may be different) 295 | self->root = head.link[1]; 296 | } 297 | 298 | // Make the root black for simplified logic 299 | self->root->red = 0; 300 | ++self->size; 301 | } 302 | 303 | return 1; 304 | } 305 | 306 | // Returns 1 if the value was removed, 0 otherwise. Optional node callback 307 | // can be provided to dealloc node and/or user data. Use rb_tree_node_dealloc 308 | // default callback to deallocate node created by rb_tree_insert(...). 309 | int 310 | rb_tree_remove_with_cb (struct rb_tree *self, void *value, rb_tree_node_f node_cb) { 311 | if (self->root != NULL) { 312 | struct rb_node head = {0}; // False tree root 313 | struct rb_node node = { .value = value }; // Value wrapper node 314 | struct rb_node *q, *p, *g; // Helpers 315 | struct rb_node *f = NULL; // Found item 316 | int dir = 1; 317 | 318 | // Set up our helpers 319 | q = &head; 320 | g = p = NULL; 321 | q->link[1] = self->root; 322 | 323 | // Search and push a red node down 324 | // to fix red violations as we go 325 | while (q->link[dir] != NULL) { 326 | int last = dir; 327 | 328 | // Move the helpers down 329 | g = p, p = q; 330 | q = q->link[dir]; 331 | dir = self->cmp(self, q, &node) < 0; 332 | 333 | // Save the node with matching value and keep 334 | // going; we'll do removal tasks at the end 335 | if (self->cmp(self, q, &node) == 0) { 336 | f = q; 337 | } 338 | 339 | // Push the red node down with rotations and color flips 340 | if (!rb_node_is_red(q) && !rb_node_is_red(q->link[dir])) { 341 | if (rb_node_is_red(q->link[!dir])) { 342 | p = p->link[last] = rb_node_rotate(q, dir); 343 | } else if (!rb_node_is_red(q->link[!dir])) { 344 | struct rb_node *s = p->link[!last]; 345 | if (s) { 346 | if (!rb_node_is_red(s->link[!last]) && !rb_node_is_red(s->link[last])) { 347 | 348 | // Color flip 349 | p->red = 0; 350 | s->red = 1; 351 | q->red = 1; 352 | } else { 353 | int dir2 = g->link[1] == p; 354 | if (rb_node_is_red(s->link[last])) { 355 | g->link[dir2] = rb_node_rotate2(p, last); 356 | } else if (rb_node_is_red(s->link[!last])) { 357 | g->link[dir2] = rb_node_rotate(p, last); 358 | } 359 | 360 | // Ensure correct coloring 361 | q->red = g->link[dir2]->red = 1; 362 | g->link[dir2]->link[0]->red = 0; 363 | g->link[dir2]->link[1]->red = 0; 364 | } 365 | } 366 | } 367 | } 368 | } 369 | 370 | // Replace and remove the saved node 371 | if (f) { 372 | void *tmp = f->value; 373 | f->value = q->value; 374 | q->value = tmp; 375 | 376 | p->link[p->link[1] == q] = q->link[q->link[0] == NULL]; 377 | 378 | if (node_cb) { 379 | node_cb(self, q); 380 | } 381 | q = NULL; 382 | } 383 | 384 | // Update the root (it may be different) 385 | self->root = head.link[1]; 386 | 387 | // Make the root black for simplified logic 388 | if (self->root != NULL) { 389 | self->root->red = 0; 390 | } 391 | 392 | --self->size; 393 | } 394 | return 1; 395 | } 396 | 397 | int 398 | rb_tree_remove (struct rb_tree *self, void *value) { 399 | int result = 0; 400 | if (self) { 401 | result = rb_tree_remove_with_cb(self, value, rb_tree_node_dealloc_cb); 402 | } 403 | return result; 404 | } 405 | 406 | size_t 407 | rb_tree_size (struct rb_tree *self) { 408 | size_t result = 0; 409 | if (self) { 410 | result = self->size; 411 | } 412 | return result; 413 | } 414 | 415 | // rb_iter 416 | 417 | struct rb_iter * 418 | rb_iter_alloc () { 419 | return malloc(sizeof(struct rb_iter)); 420 | } 421 | 422 | struct rb_iter * 423 | rb_iter_init (struct rb_iter *self) { 424 | if (self) { 425 | self->tree = NULL; 426 | self->node = NULL; 427 | self->top = 0; 428 | } 429 | return self; 430 | } 431 | 432 | struct rb_iter * 433 | rb_iter_create () { 434 | return rb_iter_init(rb_iter_alloc()); 435 | } 436 | 437 | void 438 | rb_iter_dealloc (struct rb_iter *self) { 439 | if (self) { 440 | free(self); 441 | } 442 | } 443 | 444 | // Internal function, init traversal object, dir determines whether 445 | // to begin traversal at the smallest or largest valued node. 446 | static void * 447 | rb_iter_start (struct rb_iter *self, struct rb_tree *tree, int dir) { 448 | void *result = NULL; 449 | if (self) { 450 | self->tree = tree; 451 | self->node = tree->root; 452 | self->top = 0; 453 | 454 | // Save the path for later selfersal 455 | if (self->node != NULL) { 456 | while (self->node->link[dir] != NULL) { 457 | self->path[self->top++] = self->node; 458 | self->node = self->node->link[dir]; 459 | } 460 | } 461 | 462 | result = self->node == NULL ? NULL : self->node->value; 463 | } 464 | return result; 465 | } 466 | 467 | // Traverse a red black tree in the user-specified direction (0 asc, 1 desc) 468 | static void * 469 | rb_iter_move (struct rb_iter *self, int dir) { 470 | if (self->node->link[dir] != NULL) { 471 | 472 | // Continue down this branch 473 | self->path[self->top++] = self->node; 474 | self->node = self->node->link[dir]; 475 | while ( self->node->link[!dir] != NULL ) { 476 | self->path[self->top++] = self->node; 477 | self->node = self->node->link[!dir]; 478 | } 479 | } else { 480 | 481 | // Move to the next branch 482 | struct rb_node *last = NULL; 483 | do { 484 | if (self->top == 0) { 485 | self->node = NULL; 486 | break; 487 | } 488 | last = self->node; 489 | self->node = self->path[--self->top]; 490 | } while (last == self->node->link[dir]); 491 | } 492 | return self->node == NULL ? NULL : self->node->value; 493 | } 494 | 495 | void * 496 | rb_iter_first (struct rb_iter *self, struct rb_tree *tree) { 497 | return rb_iter_start(self, tree, 0); 498 | } 499 | 500 | void * 501 | rb_iter_last (struct rb_iter *self, struct rb_tree *tree) { 502 | return rb_iter_start(self, tree, 1); 503 | } 504 | 505 | void * 506 | rb_iter_next (struct rb_iter *self) { 507 | return rb_iter_move(self, 1); 508 | } 509 | 510 | void * 511 | rb_iter_prev (struct rb_iter *self) { 512 | return rb_iter_move(self, 0); 513 | } 514 | --------------------------------------------------------------------------------