1cimport numpy as np
2from libcpp.pair cimport pair
3from libcpp.set cimport set as cset
4from libcpp.vector cimport vector
5
6from yt.utilities.lib.ewah_bool_array cimport (
7    ewah_bool_array,
8    ewah_bool_iterator,
9    ewah_map,
10    sstream,
11)
12
13ctypedef bint bitarrtype
14ctypedef pair[np.uint64_t, np.uint64_t] ind_pair
15
16cdef class FileBitmasks:
17    cdef np.uint32_t nfiles
18    cdef ewah_map** ewah_coll
19    cdef ewah_bool_array** ewah_keys
20    cdef ewah_bool_array** ewah_refn
21
22    cdef void _reset(self)
23    cdef bint _iseq(self, FileBitmasks solf)
24    cdef BoolArrayCollection _get_bitmask(self, np.uint32_t ifile)
25    cdef tuple _find_collisions(self, BoolArrayCollection coll, bint verbose=*)
26    cdef tuple _find_collisions_coarse(self, BoolArrayCollection coll, bint
27                verbose=*, file_list=*)
28    cdef tuple _find_collisions_refined(self, BoolArrayCollection coll, bint verbose=*)
29    cdef void _set(self, np.uint32_t ifile, np.uint64_t i1, np.uint64_t i2=*)
30    cdef void _set_coarse(self, np.uint32_t ifile, np.uint64_t i1)
31    cdef void _set_refined(self, np.uint32_t ifile, np.uint64_t i1, np.uint64_t i2)
32    cdef void _set_coarse_array(self, np.uint32_t ifile, np.uint8_t[:] arr)
33    cdef void _set_refined_array(self, np.uint32_t ifile, np.uint64_t mi1, np.uint8_t[:] arr)
34    cdef void _set_refined_index_array(self, np.uint32_t ifile, np.int64_t nsub_mi,
35                                       np.ndarray[np.uint64_t, ndim=1] sub_mi1,
36                                       np.ndarray[np.uint64_t, ndim=1] sub_mi2)
37    cdef void _set_map(self, np.uint32_t ifile, np.uint64_t i1, np.uint64_t i2)
38    cdef void _set_refn(self, np.uint32_t ifile, np.uint64_t i1)
39    cdef bint _get(self, np.uint32_t ifile, np.uint64_t i1, np.uint64_t i2=*)
40    cdef bint _get_coarse(self, np.uint32_t ifile, np.uint64_t i1)
41    cdef void _get_coarse_array(self, np.uint32_t ifile, np.uint64_t imax, np.uint8_t[:] arr) except *
42    cdef bint _isref(self, np.uint32_t ifile, np.uint64_t i)
43    cdef np.uint64_t _count_total(self, np.uint32_t ifile)
44    cdef np.uint64_t _count_refined(self, np.uint32_t ifile)
45    cdef np.uint64_t _count_coarse(self, np.uint32_t ifile)
46    cdef void _append(self, np.uint32_t ifile, BoolArrayCollection solf)
47    cdef bint _intersects(self, np.uint32_t ifile, BoolArrayCollection solf)
48    cdef void _logicalxor(self, np.uint32_t ifile, BoolArrayCollection solf, BoolArrayCollection out)
49    cdef void _logicaland(self, np.uint32_t ifile, BoolArrayCollection solf, BoolArrayCollection out)
50    cdef void _select_contaminated(self, np.uint32_t ifile, BoolArrayCollection mask, np.uint8_t[:] out,
51               np.uint8_t[:] secondary_files, BoolArrayCollection mask2=*)
52    cdef void _select_uncontaminated(self, np.uint32_t ifile, BoolArrayCollection mask, np.uint8_t[:] out,
53               BoolArrayCollection mask2=*)
54    cdef bytes _dumps(self, np.uint32_t ifile)
55    cdef bint _loads(self, np.uint32_t ifile, bytes s)
56    cdef bint _check(self)
57
58cdef class BoolArrayCollection:
59    cdef ewah_map* ewah_coll
60    cdef ewah_bool_array* ewah_keys
61    cdef ewah_bool_array* ewah_refn
62    cdef ewah_bool_array* ewah_coar
63
64    cdef void _reset(self)
65    cdef int _richcmp(self, BoolArrayCollection solf, int op) except -1
66    cdef void _set(self, np.uint64_t i1, np.uint64_t i2=*)
67    cdef void _set_coarse(self, np.uint64_t i1)
68    cdef void _set_refined(self, np.uint64_t i1, np.uint64_t i2)
69    cdef void _set_coarse_array(self, np.uint8_t[:] arr)
70    cdef void _set_refined_array(self, np.uint64_t mi1, np.uint8_t[:] arr)
71    cdef void _set_map(self, np.uint64_t i1, np.uint64_t i2)
72    cdef void _set_refn(self, np.uint64_t i1)
73    cdef bint _get(self, np.uint64_t i1, np.uint64_t i2=*)
74    cdef bint _get_coarse(self, np.uint64_t i1)
75    cdef void _get_coarse_array(self, np.uint64_t imax, np.uint8_t[:] arr) except *
76    cdef bint _contains(self, np.uint64_t i)
77    cdef bint _isref(self, np.uint64_t i)
78    cdef void _ewah_coarse(self)
79    cdef np.uint64_t _count_total(self)
80    cdef np.uint64_t _count_refined(self)
81    cdef np.uint64_t _count_coarse(self)
82    cdef void _append(self, BoolArrayCollection solf)
83    cdef void _logicalor(self, BoolArrayCollection solf, BoolArrayCollection out)
84    cdef bint _intersects(self, BoolArrayCollection solf)
85    cdef void _logicalxor(self, BoolArrayCollection solf, BoolArrayCollection out)
86    cdef void _logicaland(self, BoolArrayCollection solf, BoolArrayCollection out)
87    cdef void _select_contaminated(self, BoolArrayCollection mask, np.uint8_t[:] out,
88        BoolArrayCollection mask2=*)
89    cdef void _select_uncontaminated(self, BoolArrayCollection mask, np.uint8_t[:] out,
90        BoolArrayCollection mask2=*)
91    cdef void _get_ghost_zones(self, int ngz, int order1, int order2,
92                               bint periodicity[3], BoolArrayCollection out_ewah,
93                               bint coarse_ghosts=*)
94    cdef bytes _dumps(self)
95    cdef bint _loads(self, bytes s)
96    cdef bint _check(self)
97
98cdef class BoolArrayCollectionUncompressed:
99    cdef int nele1
100    cdef int nele2
101    cdef ewah_map* ewah_coll
102    cdef bitarrtype* ewah_keys
103    cdef bitarrtype* ewah_refn
104
105    cdef void _set(self, np.uint64_t i1, np.uint64_t i2=*)
106    cdef void _set_coarse(self, np.uint64_t i1)
107    cdef void _set_refined(self, np.uint64_t i1, np.uint64_t i2)
108    cdef void _set_coarse_array(self, np.uint8_t[:] arr)
109    cdef void _set_coarse_array_ptr(self, np.uint8_t *arr)
110    cdef void _set_refined_array(self, np.uint64_t mi1, np.uint8_t[:] arr)
111    cdef void _set_refined_array_ptr(self, np.uint64_t mi1, np.uint8_t *arr)
112    cdef void _set_map(self, np.uint64_t i1, np.uint64_t i2)
113    cdef void _set_refn(self, np.uint64_t i1)
114    cdef bint _get(self, np.uint64_t i1, np.uint64_t i2=*)
115    cdef bint _get_coarse(self, np.uint64_t i1)
116    cdef bint _isref(self, np.uint64_t i)
117    cdef np.uint64_t _count_total(self)
118    cdef np.uint64_t _count_refined(self)
119    cdef void _append(self, BoolArrayCollectionUncompressed solf)
120    cdef bint _intersects(self, BoolArrayCollectionUncompressed solf)
121    cdef void _compress(self, BoolArrayCollection solf)
122
123cdef class SparseUnorderedBitmaskSet:
124    cdef cset[np.uint64_t] entries
125    cdef void _set(self, np.uint64_t ind)
126    cdef void _fill(self, np.uint8_t[:] mask)
127    cdef void _fill_ewah(self, BoolArrayCollection mm)
128    cdef void _fill_bool(self, BoolArrayCollectionUncompressed mm)
129    cdef void _reset(self)
130    cdef to_array(self)
131
132cdef class SparseUnorderedBitmaskVector:
133    cdef int total
134    cdef vector[np.uint64_t] entries
135    cdef void _set(self, np.uint64_t ind)
136    cdef void _fill(self, np.uint8_t[:] mask)
137    cdef void _fill_ewah(self, BoolArrayCollection mm)
138    cdef void _fill_bool(self, BoolArrayCollectionUncompressed mm)
139    cdef void _reset(self)
140    cdef to_array(self)
141    cdef void _remove_duplicates(self)
142    cdef void _prune(self)
143
144cdef class SparseUnorderedRefinedBitmaskSet:
145    cdef cset[ind_pair] entries
146    cdef void _set(self, np.uint64_t ind1, np.uint64_t ind2)
147    cdef void _fill(self, np.uint8_t[:] mask1, np.uint8_t[:])
148    cdef void _fill_ewah(self, BoolArrayCollection mm)
149    cdef void _fill_bool(self, BoolArrayCollectionUncompressed mm)
150    cdef void _reset(self)
151    cdef to_array(self)
152
153cdef class SparseUnorderedRefinedBitmaskVector:
154    cdef int total
155    cdef vector[ind_pair] entries
156    cdef void _set(self, np.uint64_t ind1, np.uint64_t ind2)
157    cdef void _fill(self, np.uint8_t[:] mask1, np.uint8_t[:])
158    cdef void _fill_ewah(self, BoolArrayCollection mm)
159    cdef void _fill_bool(self, BoolArrayCollectionUncompressed mm)
160    cdef void _reset(self)
161    cdef to_array(self)
162    cdef void _remove_duplicates(self)
163    cdef void _prune(self)
164