├── LICENSE ├── README.md ├── benchmarks.mojo ├── flx ├── __init__.mojo ├── cache.mojo ├── data_types.mojo ├── flx_buffer.mojo ├── flx_builder.mojo └── flx_value.mojo ├── flx2 ├── __init__.mojo ├── cache.mojo ├── data_types.mojo ├── flx_buffer.mojo ├── flx_builder.mojo └── flx_value.mojo ├── samples.mojo ├── test_builder.mojo ├── test_builder2.mojo ├── test_roundtrip.mojo ├── test_roundtrip2.mojo ├── test_value_to_json.mojo └── type_conversion_sample.mojo /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /benchmarks.mojo: -------------------------------------------------------------------------------- 1 | from time import now 2 | from math.limit import max_finite, min_finite 3 | from math import min 4 | from random import random_ui64, random_float64, random_si64 5 | from flx import * 6 | 7 | fn construct_and_read_vec[count: Int, DW: DType, DR: DType = DW]() raises: 8 | var nums = DTypePointer[DW].alloc(count) 9 | for i in range(count): 10 | @parameter 11 | if DW == DType.uint8 or DW == DType.uint16 or DW == DType.uint32 or DW == DType.uint64: 12 | nums[i] = (random_ui64(0, 1 << 63).cast[DW]()) 13 | elif DW == DType.float16 or DW == DType.float32 or DW == DType.float64: 14 | nums[i] = (random_float64().cast[DW]()) 15 | else: 16 | nums[i] = (random_si64(min_finite[DType.int64](), max_finite[DType.int64]()).cast[DW]()) 17 | var size = 0 18 | var min_duration_create = max_finite[DType.int64]().to_int() 19 | var min_duration_read = max_finite[DType.int64]().to_int() 20 | for _ in range(20): 21 | var tik = now() 22 | var r = flx[DW](nums, count) 23 | var tok = now() 24 | min_duration_create = min(min_duration_create, tok - tik) 25 | size = r.get[1, Int]() 26 | var bytes = r.get[0, DTypePointer[DType.uint8]]() 27 | tik = now() 28 | var vec = FlxValue(bytes, size) 29 | tok = now() 30 | var read_duration = tok - tik 31 | for i in range(count): 32 | tik = now() 33 | var v = vec[i].get[DR]() 34 | tok = now() 35 | read_duration += tok - tik 36 | if v != nums[i].cast[DR](): 37 | print("Error at index:", i) 38 | bytes.free() 39 | min_duration_read = min(min_duration_read, read_duration) 40 | 41 | print("Constructed a vector of", count, DW, ",", size, "bytes in", min_duration_create / 1_000_000, "ms") 42 | print("Read a vector of", count, DW, ",", size, "bytes in", min_duration_read / 1_000_000, "ms") 43 | 44 | fn construct_and_read_table[columns: Int, rows: Int, DW: DType, DR: DType = DW]() raises: 45 | var counts = columns * rows 46 | var nums = List[SIMD[DW, 1]](capacity=counts) 47 | for _ in range(counts): 48 | @parameter 49 | if DW == DType.uint8 or DW == DType.uint16 or DW == DType.uint32 or DW == DType.uint64: 50 | nums.append(random_ui64(0, 1 << 63).cast[DW]()) 51 | elif DW == DType.float16 or DW == DType.float32 or DW == DType.float64: 52 | nums.append(random_float64().cast[DW]()) 53 | else: 54 | nums.append(random_si64(min_finite[DType.int64](), max_finite[DType.int64]()).cast[DW]()) 55 | var size = 0 56 | var min_duration_create = max_finite[DType.int64]().to_int() 57 | var min_duration_read = max_finite[DType.int64]().to_int() 58 | for _ in range(20): 59 | var tik = now() 60 | var df = FlxVec() 61 | for i in range(rows): 62 | var row = df^.vec() 63 | for j in range(columns): 64 | row = row^.add[DW](nums[j + i * columns]) 65 | df = row^.up_to_vec() 66 | var r = df^.finish() 67 | var tok = now() 68 | min_duration_create = min(min_duration_create, tok - tik) 69 | size = r.get[1, Int]() 70 | var bytes = r.get[0, DTypePointer[DType.uint8]]() 71 | tik = now() 72 | var vec = FlxValue(bytes, size) 73 | tok = now() 74 | var read_duration = tok - tik 75 | for i in range(rows): 76 | for j in range(columns): 77 | tik = now() 78 | var v = vec[i][j].get[DR]() 79 | tok = now() 80 | read_duration += tok - tik 81 | if v != nums[j + i * columns].cast[DR](): 82 | print("Error at row:", i, "column:", j) 83 | bytes.free() 84 | min_duration_read = min(min_duration_read, read_duration) 85 | 86 | print("Constructed a table", "rows:", rows, "columns:", columns, DW, ",", size, "bytes in", min_duration_create / 1_000_000, "ms") 87 | print("Read a table", "rows:", rows, "columns:", columns, DW, ",", size, "bytes in", min_duration_read / 1_000_000, "ms") 88 | 89 | 90 | fn main(): 91 | try: 92 | construct_and_read_vec[1_000_000, DType.uint64]() 93 | construct_and_read_vec[1_000_000, DType.uint32]() 94 | construct_and_read_vec[2_000_000, DType.uint32]() 95 | construct_and_read_vec[1_000_000, DType.uint16, DType.uint32]() 96 | construct_and_read_vec[2_000_000, DType.uint16, DType.uint32]() 97 | construct_and_read_vec[1_000_000, DType.uint8, DType.uint32]() 98 | construct_and_read_vec[2_000_000, DType.uint8, DType.uint32]() 99 | construct_and_read_vec[4_000_000, DType.uint8, DType.uint32]() 100 | construct_and_read_vec[1_000_000, DType.int64]() 101 | construct_and_read_vec[1_000_000, DType.int32]() 102 | construct_and_read_vec[1_000_000, DType.int16, DType.int32]() 103 | construct_and_read_vec[1_000_000, DType.int8, DType.int32]() 104 | construct_and_read_vec[1_000_000, DType.float64]() 105 | construct_and_read_vec[1_000_000, DType.float32]() 106 | construct_and_read_vec[1_000_000, DType.float16, DType.float32]() 107 | construct_and_read_vec[32_000, DType.float16]() 108 | construct_and_read_vec[16_000, DType.float16]() 109 | construct_and_read_table[10, 100_000, DType.uint64]() 110 | construct_and_read_table[100_000, 10, DType.uint64]() 111 | construct_and_read_table[10, 100_000, DType.float16]() 112 | construct_and_read_table[100_000, 10, DType.float16, DType.float32]() 113 | except e: 114 | print("Unexpected error:", e) 115 | -------------------------------------------------------------------------------- /flx/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .flx_builder import FlxMap, FlxVec 2 | from .flx_value import FlxValue 3 | from .flx_buffer import flx, flx_blob, flx_null -------------------------------------------------------------------------------- /flx/cache.mojo: -------------------------------------------------------------------------------- 1 | from memory import memset_zero, memcpy 2 | from .data_types import StackValue 3 | 4 | @value 5 | struct Key(CollectionElement): 6 | var pointer: DTypePointer[DType.uint8] 7 | var size: Int 8 | 9 | fn __init__(inout self, pointer: DTypePointer[DType.uint8], size: Int): 10 | var cp = DTypePointer[DType.uint8].alloc(size) 11 | memcpy(cp, pointer, size) 12 | self.pointer = cp 13 | self.size = size 14 | 15 | # alias Key = (DTypePointer[DType.uint8], Int) 16 | alias Keys = List[Key] 17 | alias Values = List[StackValue] 18 | 19 | struct _CacheStackValue(Movable, Copyable): 20 | var keys: Keys 21 | var values: Values 22 | var key_map: DTypePointer[DType.uint32] 23 | var count: Int 24 | var capacity: Int 25 | 26 | fn __init__(inout self): 27 | self.count = 0 28 | self.capacity = 16 29 | self.keys = Keys(capacity=self.capacity) 30 | self.values = Values(capacity=self.capacity) 31 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 32 | memset_zero(self.key_map, self.capacity) 33 | 34 | fn __moveinit__(inout self, owned other: Self): 35 | self.count = other.count 36 | self.capacity = other.capacity 37 | self.values = other.values^ 38 | self.key_map = other.key_map 39 | self.keys = other.keys^ 40 | 41 | fn __copyinit__(inout self, other: Self): 42 | self.count = other.count 43 | self.capacity = other.capacity 44 | var keys_count = len(other.keys) 45 | 46 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 47 | memcpy(self.key_map, other.key_map, self.capacity) 48 | # self.keys = other.keys 49 | # self.values = other.values 50 | self.keys = Keys(capacity=keys_count) 51 | self.values = Values(capacity=keys_count) 52 | for i in range(keys_count): 53 | var key = other.keys[i] 54 | var p = key.pointer 55 | var size = key.size 56 | var cp = DTypePointer[DType.uint8].alloc(size) 57 | memcpy(cp, p, size) 58 | var new_key = Key(cp, size) 59 | self.keys.append(new_key) 60 | self.values.append(other.values[i]) 61 | 62 | fn __del__(owned self): 63 | self.key_map.free() 64 | for i in range(len(self.keys)): 65 | var key = self.keys[i] 66 | key.pointer.free() 67 | 68 | fn put(inout self, key: Key, value: StackValue): 69 | if self.count / self.capacity >= 0.8: 70 | self._rehash() 71 | self._put(key, value, -1) 72 | 73 | fn _rehash(inout self): 74 | var old_mask_capacity = self.capacity >> 3 75 | self.key_map.free() 76 | self.capacity <<= 1 77 | var mask_capacity = self.capacity >> 3 78 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 79 | memset_zero(self.key_map, self.capacity) 80 | 81 | for i in range(len(self.keys)): 82 | self._put(self.keys[i], self.values[i], i + 1) 83 | 84 | fn _put(inout self, key: Key, value: StackValue, rehash_index: Int): 85 | var key_hash = self._hash(key) 86 | var modulo_mask = self.capacity - 1 87 | var key_map_index = (key_hash & modulo_mask).to_int() 88 | while True: 89 | var key_index = self.key_map.offset(key_map_index).load().to_int() 90 | if key_index == 0: 91 | var new_key_index: Int 92 | if rehash_index == -1: 93 | self.keys.append(key) 94 | self.values.append(value) 95 | self.count += 1 96 | new_key_index = len(self.keys) 97 | else: 98 | new_key_index = rehash_index 99 | self.key_map.offset(key_map_index).store(UInt32(new_key_index)) 100 | return 101 | 102 | var other_key = self.keys[key_index - 1] 103 | if self._eq(other_key, key): 104 | self.values[key_index - 1] = value 105 | return 106 | 107 | key_map_index = (key_map_index + 1) & modulo_mask 108 | 109 | fn _hash(self, key: Key) -> UInt32: 110 | var hash: UInt32 = 0 111 | var bytes = key.pointer 112 | var count = key.size 113 | while count >= 4: 114 | var c = bytes.bitcast[DType.uint32]().load() 115 | hash = _hash_word32(hash, c) 116 | bytes = bytes.offset(4) 117 | count -= 4 118 | if count >= 2: 119 | var c = bytes.bitcast[DType.uint16]().load().cast[DType.uint32]() 120 | hash = _hash_word32(hash, c) 121 | bytes = bytes.offset(2) 122 | count -= 2 123 | if count > 0: 124 | var c = bytes.load().cast[DType.uint32]() 125 | hash = _hash_word32(hash, c) 126 | return hash 127 | 128 | fn _eq(self, a: Key, b: Key) -> Bool: 129 | var bytes_a = a.pointer 130 | var bytes_b = b.pointer 131 | var count_a = a.size 132 | var count_b = b.size 133 | if count_a != count_b: 134 | return False 135 | var count = count_a 136 | while count >= 4: 137 | if bytes_a.bitcast[DType.uint32]().load() != bytes_b.bitcast[DType.uint32]().load(): 138 | return False 139 | bytes_a = bytes_a.offset(4) 140 | bytes_b = bytes_b.offset(4) 141 | count -= 4 142 | if count >= 2: 143 | if bytes_a.bitcast[DType.uint16]().load() != bytes_b.bitcast[DType.uint16]().load(): 144 | return False 145 | bytes_a = bytes_a.offset(2) 146 | bytes_b = bytes_b.offset(2) 147 | count -= 2 148 | if count > 0: 149 | return bytes_a.load() == bytes_b.load() 150 | return True 151 | 152 | fn get(self, key: Key, default: StackValue) -> StackValue: 153 | var key_hash = self._hash(key) 154 | var modulo_mask = self.capacity - 1 155 | var key_map_index = (key_hash & modulo_mask).to_int() 156 | while True: 157 | var key_index = self.key_map.offset(key_map_index).load().to_int() 158 | if key_index == 0: 159 | return default 160 | var other_key = self.keys[key_index - 1] 161 | if self._eq(other_key, key): 162 | return self.values[key_index - 1] 163 | key_map_index = (key_map_index + 1) & modulo_mask 164 | 165 | from math.math import rotate_bits_left 166 | 167 | alias ROTATE = 5 168 | alias SEED32 = 0x9e_37_79_b9 169 | 170 | @always_inline 171 | fn _hash_word32(value: UInt32, word: UInt32) -> UInt32: 172 | return (rotate_bits_left[ROTATE](value) ^ word) * SEED32 173 | 174 | fn _key_string(key: Key) -> String: 175 | var bytes = key.pointer 176 | var count = key.size 177 | var result: String = "" 178 | for i in range(count): 179 | result += chr(bytes.load(i).to_int()) 180 | return result 181 | 182 | fn _key_int_string(key: Key) -> String: 183 | var bytes = key.pointer 184 | var count = key.size 185 | var result: String = "" 186 | for i in range(count): 187 | result += String(bytes.load(i).to_int()) 188 | return result 189 | 190 | 191 | struct _CacheStringOrKey(Movable, Copyable): 192 | # offsets and counts 193 | var ocs: List[(Int, Int)] 194 | var key_map: DTypePointer[DType.uint32] 195 | var count: Int 196 | var capacity: Int 197 | 198 | fn __init__(inout self): 199 | self.count = 0 200 | self.capacity = 16 201 | self.ocs = List[(Int, Int)](capacity=self.capacity) 202 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 203 | memset_zero(self.key_map, self.capacity) 204 | 205 | fn __moveinit__(inout self, owned other: Self): 206 | self.count = other.count 207 | self.capacity = other.capacity 208 | self.ocs = other.ocs^ 209 | self.key_map = other.key_map 210 | 211 | fn __copyinit__(inout self, other: Self): 212 | self.count = other.count 213 | self.capacity = other.capacity 214 | # TODO: copies elements one by one because otherwise it throws a core dump 215 | # self.ocs = other.ocs 216 | self.ocs = List[(Int, Int)](capacity=self.capacity) 217 | for i in range(self.capacity): 218 | self.ocs[i] = other.ocs[i] 219 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 220 | memcpy(self.key_map, other.key_map, self.capacity) 221 | 222 | fn __del__(owned self): 223 | self.key_map.free() 224 | 225 | fn put(inout self, oc: (Int, Int), pointer: DTypePointer[DType.uint8]): 226 | if self.count / self.capacity >= 0.8: 227 | self._rehash(pointer) 228 | self._put(oc, pointer, -1) 229 | 230 | fn get(self, bc: (DTypePointer[DType.uint8], Int), pointer: DTypePointer[DType.uint8]) -> Int: 231 | var bytes = bc.get[0, DTypePointer[DType.uint8]]() 232 | var count = bc.get[1, Int]() 233 | var key_hash = self._hash(bytes, count) 234 | var modulo_mask = self.capacity - 1 235 | var key_map_index = (key_hash & modulo_mask).to_int() 236 | while True: 237 | var key_index = self.key_map.offset(key_map_index).load().to_int() 238 | if key_index == 0: 239 | return -1 240 | var other_oc = self.ocs[key_index - 1] 241 | if self._eq(count, other_oc.get[1, Int](), bytes, pointer.offset(other_oc.get[0, Int]())): 242 | return other_oc.get[0, Int]() 243 | key_map_index = (key_map_index + 1) & modulo_mask 244 | 245 | fn _rehash(inout self, pointer: DTypePointer[DType.uint8]): 246 | self.key_map.free() 247 | self.capacity <<= 1 248 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 249 | memset_zero(self.key_map, self.capacity) 250 | for i in range(len(self.ocs)): 251 | self._put(self.ocs[i], pointer, i + 1) 252 | 253 | fn _put(inout self, oc: (Int, Int), pointer: DTypePointer[DType.uint8], rehash_index: Int): 254 | var bytes = pointer.offset(oc.get[0, Int]()) 255 | var count = oc.get[1, Int]() 256 | var key_hash = self._hash(bytes, count) 257 | var modulo_mask = self.capacity - 1 258 | var key_map_index = (key_hash & modulo_mask).to_int() 259 | while True: 260 | var key_index = self.key_map.offset(key_map_index).load().to_int() 261 | if key_index == 0: 262 | var new_key_index: Int 263 | if rehash_index == -1: 264 | self.ocs.append(oc) 265 | self.count += 1 266 | new_key_index = len(self.ocs) 267 | else: 268 | new_key_index = rehash_index 269 | self.key_map.offset(key_map_index).store(UInt32(new_key_index)) 270 | return 271 | 272 | var other_ol = self.ocs[key_index - 1] 273 | if self._eq(count, other_ol.get[1, Int](), bytes, pointer.offset(other_ol.get[0, Int]())): 274 | return 275 | 276 | key_map_index = (key_map_index + 1) & modulo_mask 277 | 278 | fn _hash(self, _bytes: DTypePointer[DType.uint8], _count: Int) -> UInt32: 279 | var bytes = _bytes 280 | var count = _count 281 | var hash: UInt32 = 0 282 | while count >= 4: 283 | var c = bytes.bitcast[DType.uint32]().load() 284 | hash = _hash_word32(hash, c) 285 | bytes = bytes.offset(4) 286 | count -= 4 287 | if count >= 2: 288 | var c = bytes.bitcast[DType.uint16]().load().cast[DType.uint32]() 289 | hash = _hash_word32(hash, c) 290 | bytes = bytes.offset(2) 291 | count -= 2 292 | if count > 0: 293 | var c = bytes.load().cast[DType.uint32]() 294 | hash = _hash_word32(hash, c) 295 | return hash 296 | 297 | fn _eq(self, _count_a: Int, _count_b: Int, _bytes_a: DTypePointer[DType.uint8], _bytes_b: DTypePointer[DType.uint8]) -> Bool: 298 | var bytes_a = _bytes_a 299 | var bytes_b = _bytes_b 300 | var count_a = _count_a 301 | var count_b = _count_b 302 | if count_a != count_b: 303 | return False 304 | var count = count_a 305 | while count >= 8: 306 | if bytes_a.bitcast[DType.uint64]().load() != bytes_b.bitcast[DType.uint64]().load(): 307 | return False 308 | bytes_a = bytes_a.offset(8) 309 | bytes_b = bytes_b.offset(8) 310 | count -= 8 311 | if count >= 4: 312 | if bytes_a.bitcast[DType.uint32]().load() != bytes_b.bitcast[DType.uint32]().load(): 313 | return False 314 | bytes_a = bytes_a.offset(4) 315 | bytes_b = bytes_b.offset(4) 316 | count -= 4 317 | if count >= 2: 318 | if bytes_a.bitcast[DType.uint16]().load() != bytes_b.bitcast[DType.uint16]().load(): 319 | return False 320 | bytes_a = bytes_a.offset(2) 321 | bytes_b = bytes_b.offset(2) 322 | count -= 2 323 | if count > 0: 324 | return bytes_a.load() == bytes_b.load() 325 | return True 326 | -------------------------------------------------------------------------------- /flx/data_types.mojo: -------------------------------------------------------------------------------- 1 | from math.limit import max_finite, min_finite 2 | from math.bit import bswap, bit_length 3 | from memory.unsafe import bitcast 4 | from sys.info import is_big_endian 5 | 6 | alias is_be = is_big_endian() 7 | 8 | @value 9 | @register_passable("trivial") 10 | struct ValueBitWidth: 11 | alias width8 = ValueBitWidth(0) 12 | alias width16 = ValueBitWidth(1) 13 | alias width32 = ValueBitWidth(2) 14 | alias width64 = ValueBitWidth(3) 15 | 16 | alias min8 = min_finite[DType.int8]().to_int() 17 | alias max8 = max_finite[DType.int8]().to_int() 18 | alias umax8 = max_finite[DType.uint8]().cast[DType.uint64]() 19 | alias min16 = min_finite[DType.int16]().to_int() 20 | alias max16 = max_finite[DType.int16]().to_int() 21 | alias umax16 = max_finite[DType.uint16]().cast[DType.uint64]() 22 | alias min32 = min_finite[DType.int32]().to_int() 23 | alias max32 = max_finite[DType.int32]().to_int() 24 | alias umax32 = max_finite[DType.uint32]().cast[DType.uint64]() 25 | 26 | var value: UInt8 27 | 28 | @always_inline 29 | fn __lt__(self, other: ValueBitWidth) -> Bool: 30 | return self.value < other.value 31 | 32 | @always_inline 33 | fn __le__(self, other: ValueBitWidth) -> Bool: 34 | return self.value <= other.value 35 | 36 | @always_inline 37 | fn __eq__(self, other: ValueBitWidth) -> Bool: 38 | return self.value == other.value 39 | 40 | @always_inline 41 | @staticmethod 42 | fn of[D: DType](n: SIMD[D, 1]) -> ValueBitWidth: 43 | @parameter 44 | if D == DType.uint8 or D == DType.int8 or D == DType.bool: 45 | return ValueBitWidth.width8 46 | elif D == DType.uint16 or D == DType.int16 or D == DType.float16: 47 | return ValueBitWidth.width16 48 | elif D == DType.uint32 or D == DType.int32 or D == DType.float32: 49 | return ValueBitWidth.width32 50 | else: 51 | return ValueBitWidth.width64 52 | 53 | @always_inline 54 | @staticmethod 55 | fn of(n: Int) -> ValueBitWidth: 56 | return ValueBitWidth((bit_length(Int64(n)) >> 3).to_int()) 57 | 58 | @always_inline 59 | @staticmethod 60 | fn of(n: UInt64) -> ValueBitWidth: 61 | return ValueBitWidth((bit_length(n) >> 3).to_int()) 62 | 63 | @always_inline 64 | fn padding_size(buffer_size: UInt64, scalar_size: UInt64) -> UInt64: 65 | return (~buffer_size + 1) & (scalar_size - 1) 66 | 67 | @value 68 | @register_passable("trivial") 69 | struct ValueType: 70 | alias Null = ValueType(0) 71 | alias Int = ValueType(1) 72 | alias UInt = ValueType(2) 73 | alias Float = ValueType(3) 74 | alias Key = ValueType(4) 75 | alias String = ValueType(5) 76 | alias IndirectInt = ValueType(6) 77 | alias IndirectUInt = ValueType(7) 78 | alias IndirectFloat = ValueType(8) 79 | alias Map = ValueType(9) 80 | alias Vector = ValueType(10) 81 | alias VectorInt = ValueType(11) 82 | alias VectorUInt = ValueType(12) 83 | alias VectorFloat = ValueType(13) 84 | alias VectorKey = ValueType(14) 85 | alias VectorString = ValueType(15) 86 | alias VectorInt2 = ValueType(16) 87 | alias VectorUInt2 = ValueType(17) 88 | alias VectorFloat2 = ValueType(18) 89 | alias VectorInt3 = ValueType(19) 90 | alias VectorUInt3 = ValueType(20) 91 | alias VectorFloat3 = ValueType(21) 92 | alias VectorInt4 = ValueType(22) 93 | alias VectorUInt4 = ValueType(23) 94 | alias VectorFloat4 = ValueType(24) 95 | alias Blob = ValueType(25) 96 | alias Bool = ValueType(26) 97 | alias VectorBool = ValueType(36) 98 | 99 | alias NullPackedType = 0 100 | 101 | var value: UInt8 102 | 103 | @always_inline 104 | fn __eq__(self, other: Self) -> Bool: 105 | return self.value == other.value 106 | 107 | @always_inline 108 | fn __ne__(self, other: Self) -> Bool: 109 | return self.value != other.value 110 | 111 | @always_inline 112 | fn __lt__(self, other: Self) -> Bool: 113 | return self.value < other.value 114 | 115 | @always_inline 116 | fn __le__(self, other: Self) -> Bool: 117 | return self.value <= other.value 118 | 119 | @always_inline 120 | fn __sub__(self, other: Self) -> Self: 121 | return Self(self.value - other.value) 122 | 123 | @always_inline 124 | fn __add__(self, other: Self) -> Self: 125 | return Self(self.value + other.value) 126 | 127 | @always_inline 128 | fn __add__(self, other: UInt8) -> Self: 129 | return Self(self.value + other.value) 130 | 131 | @always_inline 132 | fn __mod__(self, other: UInt8) -> Self: 133 | return Self(self.value % other) 134 | 135 | @always_inline 136 | fn __floordiv__(self, other: UInt8) -> Self: 137 | return Self(self.value // other) 138 | 139 | @always_inline 140 | fn __lshift__(self, other: UInt8) -> Self: 141 | return Self(self.value << other) 142 | 143 | @always_inline 144 | fn is_inline(self) -> Bool: 145 | return self == ValueType.Bool or self <= ValueType.Float 146 | 147 | @always_inline 148 | fn is_typed_vector_element(self) -> Bool: 149 | return self == ValueType.Bool or ValueType.Int <= self <= ValueType.String 150 | 151 | @always_inline 152 | fn is_typed_vector(self) -> Bool: 153 | return self == ValueType.VectorBool or ValueType.VectorInt <= self <= ValueType.VectorString 154 | 155 | @always_inline 156 | fn is_fixed_typed_vector(self) -> Bool: 157 | return ValueType.VectorInt2 <= self <= ValueType.VectorFloat4 158 | 159 | @always_inline 160 | fn is_a_vector(self) -> Bool: 161 | return self == ValueType.VectorBool or ValueType.Vector <= self <= ValueType.VectorFloat4 162 | 163 | @always_inline 164 | fn to_typed_vector(self, length: UInt8) raises -> ValueType: 165 | if length == 0: 166 | return self - ValueType.Int + ValueType.VectorInt 167 | if length == 2: 168 | return self - ValueType.Int + ValueType.VectorInt2 169 | if length == 3: 170 | return self - ValueType.Int + ValueType.VectorInt3 171 | if length == 4: 172 | return self - ValueType.Int + ValueType.VectorInt4 173 | raise "Unexpected length " + String(length) 174 | 175 | @always_inline 176 | fn typed_vector_element_type(self) -> ValueType: 177 | return self - ValueType.VectorInt + ValueType.Int 178 | 179 | @always_inline 180 | fn fixed_typed_vector_element_type(self) -> ValueType: 181 | return ((self - ValueType.VectorInt2) % 3) + ValueType.Int 182 | 183 | @always_inline 184 | fn fixed_typed_vector_element_size(self) -> Int: 185 | return ((self - ValueType.VectorInt2) // 3).value.to_int() + 2 186 | 187 | @always_inline 188 | fn packed_type(self, bit_width: ValueBitWidth) -> UInt8: 189 | return (self << 2).value | bit_width.value 190 | 191 | @staticmethod 192 | @always_inline 193 | fn of[D: DType]() -> Self: 194 | @parameter 195 | if D == DType.uint8 or D == DType.uint16 or D == DType.uint32 or D == DType.uint64: 196 | return ValueType.UInt 197 | elif D == DType.float16 or D == DType.float32 or D == DType.float64: 198 | return ValueType.Float 199 | elif D == DType.bool: 200 | return ValueType.Bool 201 | else: 202 | return ValueType.Int 203 | 204 | 205 | @value 206 | @register_passable("trivial") 207 | struct StackValue(CollectionElement): 208 | var value: SIMD[DType.uint8, 8] 209 | var width: ValueBitWidth 210 | var type: ValueType 211 | 212 | alias Null = StackValue(0, ValueBitWidth.width8, ValueType.Null) 213 | 214 | @staticmethod 215 | @always_inline 216 | fn of[D: DType](v: SIMD[D, 1]) -> Self: 217 | # TODO: make it better 218 | var value = SIMD[DType.uint8, 8](0) 219 | @parameter 220 | if D == DType.bool: 221 | value[0] = v.cast[DType.uint8]() 222 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 223 | elif D == DType.uint8 or D == DType.int8: 224 | var v1 = bitcast[DType.uint8, 1](v) 225 | value[0] = v1[0] 226 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 227 | elif D == DType.uint16 or D == DType.int16 or D == DType.float16: 228 | var v1 = bitcast[DType.uint8, 2](v) 229 | @parameter 230 | if is_be: 231 | v1 = bswap(v1) 232 | value[0] = v1[0] 233 | value[1] = v1[1] 234 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 235 | elif D == DType.uint32 or D == DType.int32 or D == DType.float32: 236 | var v1 = bitcast[DType.uint8, 4](v) 237 | @parameter 238 | if is_be: 239 | v1 = bswap(v1) 240 | value[0] = v1[0] 241 | value[1] = v1[1] 242 | value[2] = v1[2] 243 | value[3] = v1[3] 244 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 245 | else: 246 | var v1 = bitcast[DType.uint8, 8](v) 247 | @parameter 248 | if is_be: 249 | v1 = bswap(v1) 250 | return StackValue(bitcast[DType.uint8, 8](v), ValueBitWidth.of(v), ValueType.of[D]()) 251 | 252 | @staticmethod 253 | @always_inline 254 | fn of(v: Int) -> Self: 255 | var value = bitcast[DType.uint8, 8](Int64(v)) 256 | return StackValue(value, ValueBitWidth.of(v), ValueType.Int) 257 | 258 | @always_inline 259 | fn stored_width(self, bit_width: ValueBitWidth = ValueBitWidth.width8) -> ValueBitWidth: 260 | if self.type.is_inline(): 261 | return bit_width if self.width < bit_width else self.width 262 | return self.width 263 | 264 | @always_inline 265 | fn stored_packed_type(self, bit_width: ValueBitWidth = ValueBitWidth.width8) -> UInt8: 266 | return self.type.packed_type(self.stored_width(bit_width)) 267 | 268 | @always_inline 269 | fn element_width(self, size: UInt64, index: Int) raises -> ValueBitWidth: 270 | if self.type.is_inline(): 271 | return self.width 272 | for i in range(4): 273 | var width = UInt64(1 << i) 274 | var offset_loc = size + padding_size(size, width) + UInt64(index * width) 275 | var offset = offset_loc - self.as_uint() 276 | var bit_width = ValueBitWidth.of(offset.to_int()) 277 | if (1 << bit_width.value).cast[DType.uint64]() == width: 278 | return bit_width 279 | 280 | raise "Element with size " + String(size) + " and index " + String(index) + " is of unknown width" 281 | 282 | @always_inline 283 | fn as_float(self) -> Float64: 284 | return bitcast[DType.float64, 1](self.value) 285 | 286 | @always_inline 287 | fn as_int(self) -> Int64: 288 | return bitcast[DType.int64, 1](self.value) 289 | 290 | @always_inline 291 | fn as_uint(self) -> UInt64: 292 | return bitcast[DType.uint64, 1](self.value) 293 | 294 | fn to_value(self, byte_width: UInt64) -> SIMD[DType.uint8, 8]: 295 | var self_byte_width = (1 << self.width.value).cast[DType.uint64]() 296 | if self_byte_width == byte_width or self.type == ValueType.UInt: 297 | return self.value 298 | else: 299 | if self_byte_width > 2: 300 | if self_byte_width == 4: 301 | var v = self.value.slice[4]() 302 | if byte_width == 2: 303 | if self.type == ValueType.Float: 304 | return StackValue.of(bitcast[DType.float32, 1](v).cast[DType.float16]()).value 305 | else: 306 | return StackValue.of(bitcast[DType.int32, 1](v).cast[DType.int16]()).value 307 | else: 308 | # byte_width == 8 309 | if self.type == ValueType.Float: 310 | return StackValue.of(bitcast[DType.float32, 1](v).cast[DType.float64]()).value 311 | else: 312 | return StackValue.of(bitcast[DType.int32, 1](v).cast[DType.int64]()).value 313 | else: 314 | # self_byte_width == 8 315 | if byte_width == 2: 316 | if self.type == ValueType.Float: 317 | return StackValue.of(bitcast[DType.float64, 1](self.value).cast[DType.float16]()).value 318 | else: 319 | return StackValue.of(bitcast[DType.int64, 1](self.value).cast[DType.int16]()).value 320 | else: 321 | if self.type == ValueType.Float: 322 | return StackValue.of(bitcast[DType.float64, 1](self.value).cast[DType.float32]()).value 323 | else: 324 | return StackValue.of(bitcast[DType.int64, 1](self.value).cast[DType.int32]()).value 325 | else: 326 | if self_byte_width == 2: 327 | var v = self.value.slice[2]() 328 | if byte_width == 4: 329 | if self.type == ValueType.Float: 330 | return StackValue.of(bitcast[DType.float16, 1](v).cast[DType.float32]()).value 331 | else: 332 | return StackValue.of(bitcast[DType.int16, 1](v).cast[DType.int32]()).value 333 | else: 334 | if self.type == ValueType.Float: 335 | return StackValue.of(bitcast[DType.float16, 1](v).cast[DType.float64]()).value 336 | else: 337 | return StackValue.of(bitcast[DType.int16, 1](v).cast[DType.int64]()).value 338 | else: 339 | if byte_width > 2: 340 | if byte_width == 4: 341 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int32]()).value 342 | else: 343 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int64]()).value 344 | else: 345 | if byte_width == 2: 346 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int16]()).value 347 | else: 348 | return self.value 349 | 350 | @always_inline 351 | fn is_float32(self) -> Bool: 352 | return self.type == ValueType.Float and self.width == ValueBitWidth.width32 353 | 354 | @always_inline 355 | fn is_offset(self) -> Bool: 356 | return not self.type.is_inline() 357 | 358 | fn __ne__(self, other: Self) -> Bool: 359 | return self.value != other.value 360 | or self.width.value != other.width.value 361 | or self.type != other.type 362 | -------------------------------------------------------------------------------- /flx/flx_buffer.mojo: -------------------------------------------------------------------------------- 1 | from .data_types import StackValue, ValueBitWidth, padding_size, ValueType 2 | from .cache import _CacheStackValue, Key, _CacheStringOrKey 3 | from memory import memcpy, memset_zero 4 | from memory.unsafe import bitcast 5 | from math import max 6 | 7 | fn flx_null() -> (DTypePointer[DType.uint8], Int): 8 | var buffer = FlxBuffer(16) 9 | buffer.add_null() 10 | return finish_ignoring_excetion(buffer^) 11 | 12 | fn flx(v: Int) -> (DTypePointer[DType.uint8], Int): 13 | var buffer = FlxBuffer(16) 14 | buffer.add(v) 15 | return finish_ignoring_excetion(buffer^) 16 | 17 | fn flx[D: DType](v: SIMD[D, 1]) -> (DTypePointer[DType.uint8], Int): 18 | var buffer = FlxBuffer(16) 19 | buffer.add(v) 20 | return finish_ignoring_excetion(buffer^) 21 | 22 | fn flx(v: String) -> (DTypePointer[DType.uint8], Int): 23 | var buffer = FlxBuffer(len(v) + 16) 24 | buffer.add(v) 25 | return finish_ignoring_excetion(buffer^) 26 | 27 | fn flx_blob(v: DTypePointer[DType.uint8], length: Int) -> (DTypePointer[DType.uint8], Int): 28 | var buffer = FlxBuffer(length + 32) 29 | buffer.blob(v, length) 30 | return finish_ignoring_excetion(buffer^) 31 | 32 | fn flx[D: DType](v: DTypePointer[D], length: Int) -> (DTypePointer[DType.uint8], Int): 33 | var buffer = FlxBuffer(length * sizeof[D]() + 1024) 34 | buffer.add(v, length) 35 | return finish_ignoring_excetion(buffer^) 36 | 37 | struct FlxBuffer[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Copyable, Movable): 38 | var _stack: List[StackValue] 39 | var _stack_positions: List[Int] 40 | var _stack_is_vector: List[SIMD[DType.bool, 1]] 41 | var _bytes: DTypePointer[DType.uint8] 42 | var _size: UInt64 43 | var _offset: UInt64 44 | var _finished: Bool 45 | var _string_cache: _CacheStringOrKey 46 | var _key_cache: _CacheStringOrKey 47 | var _keys_vec_cache: _CacheStackValue 48 | var _reference_cache: _CacheStackValue 49 | 50 | fn __init__(inout self, size: UInt64 = 1 << 11): 51 | self._size = size 52 | self._stack = List[StackValue]() 53 | self._stack_positions = List[Int]() 54 | self._stack_is_vector = List[SIMD[DType.bool, 1]]() 55 | self._bytes = DTypePointer[DType.uint8].alloc(size.to_int()) 56 | self._offset = 0 57 | self._finished = False 58 | self._string_cache = _CacheStringOrKey() 59 | self._key_cache = _CacheStringOrKey() 60 | self._keys_vec_cache = _CacheStackValue() 61 | self._reference_cache = _CacheStackValue() 62 | 63 | fn __moveinit__(inout self, owned other: Self): 64 | self._size = other._size 65 | self._stack = other._stack^ 66 | self._stack_positions = other._stack_positions^ 67 | self._stack_is_vector = other._stack_is_vector^ 68 | self._bytes = other._bytes 69 | self._offset = other._offset 70 | self._finished = other._finished 71 | self._string_cache = other._string_cache^ 72 | self._key_cache = other._key_cache^ 73 | self._keys_vec_cache = other._keys_vec_cache^ 74 | self._reference_cache = other._reference_cache^ 75 | 76 | fn __copyinit__(inout self, other: Self): 77 | self._size = other._size 78 | self._stack = other._stack 79 | self._stack_positions = other._stack_positions 80 | self._stack_is_vector = other._stack_is_vector 81 | self._bytes = DTypePointer[DType.uint8].alloc(other._size.to_int()) 82 | memcpy(self._bytes, other._bytes, other._offset.to_int()) 83 | self._offset = other._offset 84 | self._finished = other._finished 85 | self._string_cache = other._string_cache 86 | self._key_cache = other._key_cache 87 | self._keys_vec_cache = other._keys_vec_cache 88 | self._reference_cache = other._reference_cache 89 | 90 | fn __del__(owned self): 91 | if not self._finished: 92 | self._bytes.free() 93 | 94 | fn add_null(inout self): 95 | self._stack.append(StackValue.Null) 96 | 97 | fn add[D: DType](inout self, value: SIMD[D, 1]): 98 | self._stack.append(StackValue.of(value)) 99 | 100 | fn add(inout self, value: Int): 101 | self._stack.append(StackValue.of(value)) 102 | 103 | fn add(inout self, value: String): 104 | self._add_string[as_key=False](value) 105 | 106 | fn key(inout self, value: String): 107 | self._add_string[as_key=True](value) 108 | 109 | fn _add_string[as_key: Bool](inout self, value: String): 110 | var byte_length = len(value) 111 | var bit_width = ValueBitWidth.of(byte_length) 112 | var bytes = value._as_ptr().bitcast[DType.uint8]() 113 | 114 | @parameter 115 | if dedup_string and not as_key: 116 | var cached_offset = self._string_cache.get((bytes, byte_length), self._bytes) 117 | if cached_offset != -1: 118 | self._stack.append(StackValue(bitcast[DType.uint8, 8](Int64(cached_offset)), bit_width, ValueType.String)) 119 | return 120 | 121 | @parameter 122 | if dedup_key and as_key: 123 | var cached_offset = self._key_cache.get((bytes, byte_length), self._bytes) 124 | if cached_offset != -1: 125 | self._stack.append(StackValue(bitcast[DType.uint8, 8](Int64(cached_offset)), bit_width, ValueType.Key)) 126 | return 127 | 128 | @parameter 129 | if not as_key: 130 | var byte_width = self._align(bit_width) 131 | self._write(byte_length, byte_width) 132 | 133 | var offset = self._offset 134 | var new_offest = self._new_offset(byte_length) 135 | memcpy(self._bytes.offset(self._offset.to_int()), bytes, byte_length) 136 | self._offset = new_offest 137 | self._write(0) 138 | 139 | @parameter 140 | if dedup_string and not as_key: 141 | self._string_cache.put((offset.to_int(), byte_length), self._bytes) 142 | @parameter 143 | if dedup_key and as_key: 144 | self._key_cache.put((offset.to_int(), byte_length), self._bytes) 145 | 146 | @parameter 147 | if as_key: 148 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Key)) 149 | else: 150 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.String)) 151 | value._strref_keepalive() 152 | 153 | fn blob(inout self, value: DTypePointer[DType.uint8], length: Int): 154 | var bit_width = ValueBitWidth.of(length) 155 | var byte_width = self._align(bit_width) 156 | self._write(length, byte_width) 157 | var offset = self._offset 158 | var new_offest = self._new_offset(length) 159 | memcpy(self._bytes.offset(self._offset.to_int()), value, length) 160 | self._offset = new_offest 161 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Blob)) 162 | 163 | fn add_indirect[D: DType](inout self, value: SIMD[D, 1]): 164 | var value_type = ValueType.of[D]() 165 | if value_type == ValueType.Int or value_type == ValueType.UInt or value_type == ValueType.Float: 166 | var bit_width = ValueBitWidth.of(value) 167 | var byte_width = self._align(bit_width) 168 | var offset = self._offset 169 | self._write(StackValue.of(value), byte_width) 170 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, value_type + 5)) 171 | else: 172 | self._stack.append(StackValue.of(value)) 173 | 174 | fn add[D: DType](inout self, value: DTypePointer[D], length: Int): 175 | var len_bit_width = ValueBitWidth.of(length) 176 | var elem_bit_width = ValueBitWidth.of(SIMD[D, 1](0)) 177 | if len_bit_width <= elem_bit_width: 178 | var bit_width = len_bit_width if elem_bit_width < len_bit_width else elem_bit_width 179 | var byte_width = self._align(bit_width) 180 | self._write(length, byte_width) 181 | var offset = self._offset 182 | var byte_length = sizeof[D]() * length 183 | var new_offest = self._new_offset(byte_length) 184 | memcpy(self._bytes.offset(self._offset.to_int()), value.bitcast[DType.uint8](), byte_length) 185 | self._offset = new_offest 186 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.of[D]() + ValueType.Vector)) 187 | else: 188 | self.start_vector() 189 | for i in range(length): 190 | self.add[D](value.load(i)) 191 | try: 192 | self.end() 193 | except: 194 | pass 195 | 196 | fn add_referenced(inout self, reference_key: String) raises: 197 | var key = Key(reference_key._as_ptr().bitcast[DType.uint8](), len(reference_key)) 198 | var stack_value = self._reference_cache.get(key, StackValue.Null) 199 | key.pointer.free() 200 | if stack_value.type == ValueType.Null: 201 | raise "No value for reference key " + reference_key 202 | self._stack.append(stack_value) 203 | 204 | fn start_vector(inout self): 205 | self._stack_positions.append(len(self._stack)) 206 | self._stack_is_vector.append(True) 207 | 208 | fn start_map(inout self): 209 | self._stack_positions.append(len(self._stack)) 210 | self._stack_is_vector.append(False) 211 | 212 | fn end(inout self, reference_key: String = "") raises: 213 | var position = self._stack_positions.pop_back() 214 | var is_vector = self._stack_is_vector.pop_back() 215 | if is_vector: 216 | self._end_vector(position) 217 | else: 218 | self._sort_keys_and_end_map(position) 219 | if len(reference_key) > 0: 220 | var key = Key(reference_key._as_ptr().bitcast[DType.uint8](), len(reference_key)) 221 | self._reference_cache.put(key, self._stack[len(self._stack) - 1]) 222 | 223 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 224 | return self._finish() 225 | 226 | fn _finish(inout self) raises -> (DTypePointer[DType.uint8], Int): 227 | self._finished = True 228 | 229 | while len(self._stack_positions) > 0: 230 | self.end() 231 | 232 | if len(self._stack) != 1: 233 | raise "Stack needs to have only one element. Instead of: " + String(len(self._stack)) 234 | 235 | var value = self._stack.pop_back() 236 | var byte_width = self._align(value.element_width(self._offset, 0)) 237 | self._write(value, byte_width) 238 | self._write(value.stored_packed_type()) 239 | self._write(byte_width.cast[DType.uint8]()) 240 | return self._bytes, self._offset.to_int() 241 | 242 | fn _align(inout self, bit_width: ValueBitWidth) -> UInt64: 243 | var byte_width = 1 << bit_width.value.to_int() 244 | self._offset += padding_size(self._offset, byte_width) 245 | return byte_width 246 | 247 | fn _write(inout self, value: StackValue, byte_width: UInt64): 248 | self._grow_bytes_if_needed(self._offset + byte_width) 249 | if value.is_offset(): 250 | var rel_offset = self._offset - value.as_uint() 251 | # Safety check not implemented for now as it is internal call and should be safe 252 | # if byte_width == 8 or rel_offset < (1 << (byte_width * 8)): 253 | self._write(rel_offset, byte_width) 254 | else: 255 | var new_offset = self._new_offset(byte_width) 256 | self._bytes.store(self._offset.to_int(), value.to_value(byte_width)) 257 | self._offset = new_offset 258 | 259 | fn _write(inout self, value: UInt64, byte_width: UInt64): 260 | self._grow_bytes_if_needed(self._offset + byte_width) 261 | var new_offset = self._new_offset(byte_width) 262 | self._bytes.store(self._offset.to_int(), bitcast[DType.uint8, 8](value)) 263 | # We write 8 bytes but the offset is still set to byte_width 264 | self._offset = new_offset 265 | 266 | fn _write(inout self, value: UInt8): 267 | self._grow_bytes_if_needed(self._offset + 1) 268 | var new_offset = self._new_offset(1) 269 | self._bytes.offset(self._offset.to_int()).store(value) 270 | self._offset = new_offset 271 | 272 | fn _new_offset(inout self, byte_width: UInt64) -> UInt64: 273 | var new_offset = self._offset + byte_width 274 | var min_size = self._offset + max(byte_width, 8) 275 | self._grow_bytes_if_needed(min_size) 276 | return new_offset 277 | 278 | fn _grow_bytes_if_needed(inout self, min_size: UInt64): 279 | var prev_size = self._size 280 | while self._size < min_size: 281 | self._size <<= 1 282 | if prev_size < self._size: 283 | var prev_bytes = self._bytes 284 | self._bytes = DTypePointer[DType.uint8].alloc(self._size.to_int()) 285 | memcpy(self._bytes, prev_bytes, self._offset.to_int()) 286 | prev_bytes.free() 287 | 288 | fn _end_vector(inout self, position: Int) raises: 289 | var length = len(self._stack) - position 290 | var vec = self._create_vector(position, length, 1) 291 | self._stack.resize(position, StackValue.Null) 292 | self._stack.append(vec) 293 | 294 | fn _sort_keys_and_end_map(inout self, position: Int) raises: 295 | if (len(self._stack) - position) & 1 == 1: 296 | raise "The stack needs to hold key value pairs (even number of elements). Check if you combined [key] with [add] method calls properly." 297 | for i in range(position + 2, len(self._stack), 2): 298 | var key = self._stack[i] 299 | var value = self._stack[i + 1] 300 | var j = i - 2 301 | while j >= position and self._should_flip(self._stack[j], key): 302 | self._stack[j + 2] = self._stack[j] 303 | self._stack[j + 3] = self._stack[j + 1] 304 | j -= 2 305 | self._stack[j + 2] = key 306 | self._stack[j + 3] = value 307 | self._end_map(position) 308 | 309 | fn _should_flip(self, a: StackValue, b: StackValue) raises -> Bool: 310 | if a.type != ValueType.Key or b.type != ValueType.Key: 311 | raise "Stack values are not keys " + String(a.type.value) + " " + String(a.type.value) 312 | var index = 0 313 | while True: 314 | var c1 = self._bytes.load(a.as_uint().to_int() + index) 315 | var c2 = self._bytes.load(b.as_uint().to_int() + index) 316 | if c1 < c2: 317 | return False 318 | if c1 > c2: 319 | return True 320 | if c1 == 0 and c2 == 0: 321 | return False 322 | index += 1 323 | 324 | fn _end_map(inout self, start: Int) raises: 325 | var length = (len(self._stack) - start) >> 1 326 | var keys = StackValue.Null 327 | @parameter 328 | if dedup_key and dedup_keys_vec: 329 | var keys_vec = self._create_keys_vec_value(start, length) 330 | var cached = self._keys_vec_cache.get(keys_vec, StackValue.Null) 331 | if cached != StackValue.Null: 332 | keys = cached 333 | keys_vec.pointer.free() 334 | else: 335 | keys = self._create_vector(start, length, 2) 336 | self._keys_vec_cache.put(keys_vec, keys) 337 | else: 338 | keys = self._create_vector(start, length, 2) 339 | var map = self._create_vector(start + 1, length, 2, keys) 340 | self._stack.resize(start, StackValue.Null) 341 | self._stack.append(map) 342 | 343 | fn _create_keys_vec_value(self, start: Int, length: Int) -> Key: 344 | var size = length * 8 345 | var result = DTypePointer[DType.uint8].alloc(size) 346 | var offset = 0 347 | memset_zero(result, size) 348 | for i in range(start, len(self._stack), 2): 349 | result.store(offset, self._stack[i].value) 350 | offset += 8 351 | var key = Key(result, size) 352 | result.free() 353 | return key 354 | 355 | fn _create_vector(inout self, start: Int, length: Int, step: Int, keys: StackValue = StackValue.Null) raises -> StackValue: 356 | var bit_width = ValueBitWidth.of(UInt64(length)) 357 | var prefix_elements = 1 358 | if keys != StackValue.Null: 359 | prefix_elements += 2 360 | var keys_bit_width = keys.element_width(self._offset, 0) 361 | if bit_width < keys_bit_width: 362 | bit_width = keys_bit_width 363 | 364 | var typed = False 365 | var vec_elem_type = ValueType.Null 366 | if length > 0: 367 | vec_elem_type = self._stack[start].type 368 | typed = vec_elem_type.is_typed_vector_element() 369 | if keys != StackValue.Null: 370 | typed = False 371 | for i in range(start, len(self._stack), step): 372 | var elem_bit_width = self._stack[i].element_width(self._offset, i + prefix_elements) 373 | if bit_width < elem_bit_width: 374 | bit_width = elem_bit_width 375 | if vec_elem_type != self._stack[i].type: 376 | typed = False 377 | if bit_width == ValueBitWidth.width64 and typed == False: 378 | break 379 | var byte_width = self._align(bit_width) 380 | if keys != StackValue.Null: 381 | self._write(keys, byte_width) 382 | self._write((1 << keys.width.value).to_int(), byte_width) 383 | self._write(UInt64(length), byte_width) 384 | var offset = self._offset 385 | for i in range(start, len(self._stack), step): 386 | self._write(self._stack[i], byte_width) 387 | if not typed: 388 | for i in range(start, len(self._stack), step): 389 | self._write(self._stack[i].stored_packed_type()) 390 | if keys != StackValue.Null: 391 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Map) 392 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Vector) 393 | 394 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Vector + vec_elem_type) 395 | 396 | fn finish_ignoring_excetion(owned flx: FlxBuffer) -> (DTypePointer[DType.uint8], Int): 397 | try: 398 | return flx^.finish() 399 | except e: 400 | # should never happen 401 | print("Unexpected error:", e) 402 | return DTypePointer[DType.uint8](), -1 403 | -------------------------------------------------------------------------------- /flx/flx_builder.mojo: -------------------------------------------------------------------------------- 1 | from .flx_buffer import FlxBuffer 2 | 3 | struct FlxMap[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Movable): 4 | var buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec] 5 | 6 | fn __init__(inout self, owned buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]): 7 | self.buffer = buffer^ 8 | 9 | fn __init__(inout self): 10 | self.buffer = FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]() 11 | self.buffer.start_map() 12 | 13 | fn __moveinit__(inout self, owned other: Self): 14 | self.buffer = other.buffer^ 15 | 16 | fn add(owned self, key: String, value: Int) -> Self: 17 | self.buffer.key(key) 18 | self.buffer.add(value) 19 | return self^ 20 | 21 | fn add[D: DType](owned self, key: String, value: SIMD[D, 1]) -> Self: 22 | self.buffer.key(key) 23 | self.buffer.add(value) 24 | return self^ 25 | 26 | fn add_indirect[D: DType](owned self, key: String, value: SIMD[D, 1]) -> Self: 27 | self.buffer.key(key) 28 | self.buffer.add_indirect(value) 29 | return self^ 30 | 31 | fn add_referenced(owned self, key: String, ref_key: String) raises -> Self: 32 | self.buffer.key(key) 33 | self.buffer.add_referenced(ref_key) 34 | return self^ 35 | 36 | fn add(owned self, key: String, value: String) -> Self: 37 | self.buffer.key(key) 38 | self.buffer.add(value) 39 | return self^ 40 | 41 | fn add(owned self, key: String, value: DTypePointer[DType.uint8], length: Int) -> Self: 42 | self.buffer.key(key) 43 | self.buffer.blob(value, length) 44 | return self^ 45 | 46 | fn map(owned self, key: String) -> Self: 47 | self.buffer.key(key) 48 | self.buffer.start_map() 49 | return self^ 50 | 51 | fn vec(owned self, key: String) -> FlxVec[dedup_string, dedup_key, dedup_keys_vec]: 52 | # TODO: investigate ownership transfer instead of copy 53 | var buffer = self.buffer 54 | buffer.key(key) 55 | buffer.start_vector() 56 | return FlxVec[dedup_string, dedup_key, dedup_keys_vec](buffer^) 57 | 58 | fn up_to_map(owned self, ref_key: String = "") raises -> Self: 59 | var depth = len(self.buffer._stack_is_vector) 60 | if depth < 2: 61 | raise "This map is not nested, please call finish instead" 62 | self.buffer.end(ref_key) 63 | if self.buffer._stack_is_vector[depth - 2]: 64 | raise "This map is nested in a vector, please call up_to_vec instead" 65 | return self^ 66 | 67 | fn up_to_vec(owned self, ref_key: String = "") raises -> FlxVec[dedup_string, dedup_key, dedup_keys_vec]: 68 | # TODO: investigate ownership transfer instead of copy 69 | var buffer = self.buffer 70 | var depth = len(buffer._stack_is_vector) 71 | if depth < 2: 72 | raise "This map is not nested, please call finish instead" 73 | buffer.end(ref_key) 74 | if not buffer._stack_is_vector[depth - 2]: 75 | raise "This map is nested in a map, please call up_to_map instead" 76 | return FlxVec[dedup_string, dedup_key, dedup_keys_vec](buffer^) 77 | 78 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 79 | return self.buffer._finish() 80 | 81 | struct FlxVec[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Movable): 82 | var buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec] 83 | 84 | fn __init__(inout self, owned buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]): 85 | self.buffer = buffer^ 86 | 87 | fn __init__(inout self): 88 | self.buffer = FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]() 89 | self.buffer.start_vector() 90 | 91 | fn __moveinit__(inout self, owned other: Self): 92 | self.buffer = other.buffer^ 93 | 94 | fn add(owned self, value: Int) -> Self: 95 | self.buffer.add(value) 96 | return self^ 97 | 98 | fn add[D: DType](owned self, value: SIMD[D, 1]) -> Self: 99 | self.buffer.add(value) 100 | return self^ 101 | 102 | fn add_indirect[D: DType](owned self, value: SIMD[D, 1]) -> Self: 103 | self.buffer.add_indirect(value) 104 | return self^ 105 | 106 | fn add_referenced(owned self, ref_key: String) raises -> Self: 107 | self.buffer.add_referenced(ref_key) 108 | return self^ 109 | 110 | fn add(owned self, value: String) -> Self: 111 | self.buffer.add(value) 112 | return self^ 113 | 114 | fn add(owned self, value: DTypePointer[DType.uint8], length: Int) -> Self: 115 | self.buffer.blob(value, length) 116 | return self^ 117 | 118 | fn map(owned self) -> FlxMap[dedup_string, dedup_key, dedup_keys_vec]: 119 | # TODO: investigate ownership transfer instead of copy 120 | var buffer = self.buffer 121 | buffer.start_map() 122 | return FlxMap[dedup_string, dedup_key, dedup_keys_vec](buffer^) 123 | 124 | fn vec(owned self) -> Self: 125 | self.buffer.start_vector() 126 | return self^ 127 | 128 | fn null(owned self) -> Self: 129 | self.buffer.add_null() 130 | return self^ 131 | 132 | fn up_to_map(owned self, ref_key: String = "") raises -> FlxMap[dedup_string, dedup_key, dedup_keys_vec]: 133 | # TODO: investigate ownership transfer instead of copy 134 | var buffer = self.buffer 135 | var depth = len(buffer._stack_is_vector) 136 | if depth < 2: 137 | raise "This vec is not nested, please call finish instead" 138 | buffer.end(ref_key) 139 | if self.buffer._stack_is_vector[depth - 2]: 140 | raise "This vec is nested in a vec, please call up_to_vec instead" 141 | return FlxMap[dedup_string, dedup_key, dedup_keys_vec](buffer^) 142 | 143 | fn up_to_vec(owned self, ref_key: String = "") raises -> Self: 144 | var depth = len(self.buffer._stack_is_vector) 145 | if depth < 2: 146 | raise "This vec is not nested, please call finish instead" 147 | self.buffer.end(ref_key) 148 | if not self.buffer._stack_is_vector[depth - 2]: 149 | raise "This vec is nested in a map, please call up_to_map instead" 150 | return self^ 151 | 152 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 153 | return self.buffer._finish() 154 | -------------------------------------------------------------------------------- /flx/flx_value.mojo: -------------------------------------------------------------------------------- 1 | from .data_types import ValueType, is_be 2 | from math.bit import bswap 3 | 4 | struct FlxValue(Sized): 5 | var _bytes: DTypePointer[DType.uint8] 6 | var _byte_width: UInt8 7 | var _parent_byte_width: UInt8 8 | var _type: ValueType 9 | 10 | @always_inline 11 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], parent_byte_width: UInt8, packed_type: UInt8): 12 | self._bytes = bytes 13 | self._parent_byte_width = parent_byte_width 14 | self._byte_width = 1 << (packed_type & 3) 15 | self._type = ValueType(packed_type >> 2) 16 | 17 | @always_inline 18 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], parent_byte_width: UInt8, byte_width: UInt8, type: ValueType): 19 | self._bytes = bytes 20 | self._parent_byte_width = parent_byte_width 21 | self._byte_width = byte_width 22 | self._type = type 23 | 24 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], length: Int) raises: 25 | if length < 3: 26 | raise "Length should be at least 3, was: " + String(length) 27 | var parent_byte_width = bytes.load(length - 1) 28 | var packed_type = bytes.load(length - 2) 29 | var offset = length - parent_byte_width.to_int() - 2 30 | self._bytes = bytes.offset(offset) 31 | self._parent_byte_width = parent_byte_width 32 | self._byte_width = 1 << (packed_type & 3) 33 | self._type = ValueType(packed_type >> 2) 34 | 35 | fn __init__(inout self, bytes_and_length: (DTypePointer[DType.uint8], Int)) raises: 36 | var bytes = bytes_and_length.get[0, DTypePointer[DType.uint8]]() 37 | var length = bytes_and_length.get[1, Int]() 38 | if length < 3: 39 | raise "Length should be at least 3, was: " + String(length) 40 | var parent_byte_width = bytes.load(length - 1) 41 | var packed_type = bytes.load(length - 2) 42 | var offset = length - parent_byte_width.to_int() - 2 43 | self._bytes = bytes.offset(offset) 44 | self._parent_byte_width = parent_byte_width 45 | self._byte_width = 1 << (packed_type & 3) 46 | self._type = ValueType(packed_type >> 2) 47 | 48 | fn __moveinit__(inout self, owned other: Self): 49 | self._bytes = other._bytes 50 | self._parent_byte_width = other._parent_byte_width 51 | self._byte_width = other._byte_width 52 | self._type = other._type 53 | 54 | @always_inline 55 | fn __len__(self) -> Int: 56 | if self.is_null(): 57 | return 0 58 | if self.is_vec(): 59 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 60 | return read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 61 | if not self._type.is_fixed_typed_vector() 62 | else self._type.fixed_typed_vector_element_size() 63 | if self._type == ValueType.String or self.is_blob() or self.is_map(): 64 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 65 | return read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 66 | if self._type == ValueType.Key: 67 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 68 | var size = 0 69 | while p.offset(size).load() != 0: 70 | size += 1 71 | return 1 72 | 73 | @always_inline 74 | fn is_null(self) -> Bool: 75 | return self._type == ValueType.Null 76 | 77 | @always_inline 78 | fn is_a[D: DType](self) -> Bool: 79 | return self._type == ValueType.of[D]() 80 | 81 | @always_inline 82 | fn is_map(self) -> Bool: 83 | return self._type == ValueType.Map 84 | 85 | @always_inline 86 | fn is_vec(self) -> Bool: 87 | return self._type.is_a_vector() 88 | 89 | @always_inline 90 | fn is_string(self) -> Bool: 91 | return self._type == ValueType.String or self._type == ValueType.Key 92 | 93 | @always_inline 94 | fn is_blob(self) -> Bool: 95 | return self._type == ValueType.Blob 96 | 97 | @always_inline 98 | fn is_int(self) -> Bool: 99 | return self._type == ValueType.Int 100 | or self._type == ValueType.UInt 101 | or self._type == ValueType.IndirectInt 102 | or self._type == ValueType.IndirectUInt 103 | 104 | @always_inline 105 | fn is_float(self) -> Bool: 106 | return self._type == ValueType.Float 107 | or self._type == ValueType.IndirectFloat 108 | 109 | @always_inline 110 | fn is_bool(self) -> Bool: 111 | return self._type == ValueType.Bool 112 | 113 | @always_inline 114 | fn get[D: DType](self) raises -> SIMD[D, 1]: 115 | if self._type != ValueType.of[D](): 116 | raise "Value is not of type " + D.__str__() + " type id: " + String(self._type.value) 117 | if sizeof[D]() != self._parent_byte_width.to_int(): 118 | raise "Value byte width is " + String(self._parent_byte_width.to_int()) + " which does not conform with " + D.__str__() 119 | @parameter 120 | if is_be: 121 | return bswap(self._bytes.bitcast[D]().load()) 122 | else: 123 | return self._bytes.bitcast[D]().load() 124 | 125 | @always_inline 126 | fn int(self) raises -> Int: 127 | if self._type == ValueType.Int: 128 | return read_int(self._bytes, self._parent_byte_width) 129 | if self._type == ValueType.UInt: 130 | return read_uint(self._bytes, self._parent_byte_width) 131 | if self._type == ValueType.IndirectInt: 132 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 133 | return read_int(p, self._byte_width) 134 | if self._type == ValueType.IndirectUInt: 135 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 136 | return read_uint(p, self._byte_width) 137 | raise "Type is not an int or uint, type id: " + String(self._type.value) 138 | 139 | @always_inline 140 | fn float(self) raises -> Float64: 141 | if self._type == ValueType.Float: 142 | return read_float(self._bytes, self._parent_byte_width) 143 | if self._type == ValueType.IndirectFloat: 144 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 145 | return read_float(p, self._byte_width) 146 | raise "Type is not a float, type id: " + String(self._type.value) 147 | 148 | @always_inline 149 | fn bool(self) raises -> Bool: 150 | if self._type == ValueType.Bool: 151 | return read_uint(self._bytes, self._parent_byte_width) == 1 152 | raise "Type is not a bool, type id: " + String(self._type.value) 153 | 154 | @always_inline 155 | fn string(self) raises -> String: 156 | if self._type == ValueType.String: 157 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 158 | var size = read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 159 | var size_width = self._byte_width 160 | while p.offset(size).load() != 0: 161 | size_width <<= 1 162 | size = read_uint(p.offset(-size_width.to_int()), size_width) 163 | var p1 = Pointer[Int8].alloc(size + 1) 164 | memcpy(p1, p.bitcast[DType.int8](), size + 1) 165 | return String(p1, size + 1) 166 | if self._type == ValueType.Key: 167 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 168 | var size = 0 169 | while p.offset(size).load() != 0: 170 | size += 1 171 | var p1 = Pointer[Int8].alloc(size + 1) 172 | memcpy(p1, p.bitcast[DType.int8](), size + 1) 173 | return String(p1, size + 1) 174 | raise "Type is not convertable to string, type id: " + String(self._type.value) 175 | 176 | @always_inline 177 | fn blob(self) raises -> (DTypePointer[DType.uint8], Int): 178 | if not self.is_blob(): 179 | raise "Type is not blob, type id: " + String(self._type.value) 180 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 181 | var size = read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 182 | return (p, size) 183 | 184 | @always_inline 185 | fn vec(self) raises -> FlxVecValue: 186 | if not self._type.is_a_vector(): 187 | raise "Value is not a vector. Type id: " + String(self._type.value) 188 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 189 | var size = read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 190 | if not self._type.is_fixed_typed_vector() 191 | else self._type.fixed_typed_vector_element_size() 192 | return FlxVecValue(p, self._byte_width, self._type, size) 193 | 194 | @always_inline 195 | fn map(self) raises -> FlxMapValue: 196 | if self._type != ValueType.Map: 197 | raise "Value is not a map. Type id: " + String(self._type.value) 198 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 199 | var size = read_uint(p.offset(-self._byte_width.to_int()), self._byte_width) 200 | return FlxMapValue(p, self._byte_width, size) 201 | 202 | @always_inline 203 | fn has_key(self, key: String) raises -> Bool: 204 | if not self.is_map(): 205 | return False 206 | return self.map().key_index(key) >= 0 207 | 208 | @always_inline 209 | fn __getitem__(self, index: Int) raises -> FlxValue: 210 | return self.vec()[index] 211 | 212 | @always_inline 213 | fn __getitem__(self, key: String) raises -> FlxValue: 214 | return self.map()[key] 215 | 216 | fn json(self) raises -> String: 217 | if self.is_null(): 218 | return "null" 219 | if self.is_bool(): 220 | return "true" if self.bool() else "false" 221 | if self.is_int(): 222 | return self.int() 223 | if self.is_float(): 224 | return self.float() 225 | if self.is_string(): 226 | return '"' + self.string() + '"' 227 | if self.is_vec(): 228 | var result: String = "[" 229 | for i in range(self.__len__()): 230 | result += self[i].json() 231 | if i < self.__len__() - 1: 232 | result += "," 233 | result += "]" 234 | return result 235 | if self.is_map(): 236 | var result: String = "{" 237 | var map = self.map() 238 | var keys = map.keys() 239 | var values = map.values() 240 | for i in range(self.__len__()): 241 | result += '"' + keys[i].string() + '":' + values[i].json() 242 | if i < self.__len__() - 1: 243 | result += "," 244 | result += "}" 245 | return result 246 | raise "Unexpected type id: " + String(self._type.value) 247 | 248 | 249 | struct FlxVecValue(Sized): 250 | var _bytes: DTypePointer[DType.uint8] 251 | var _byte_width: UInt8 252 | var _type: ValueType 253 | var _length: Int 254 | 255 | @always_inline 256 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], byte_width: UInt8, type: ValueType, length: Int): 257 | self._bytes = bytes 258 | self._byte_width = byte_width 259 | self._type = type 260 | self._length = length 261 | 262 | @always_inline 263 | fn __len__(self) -> Int: 264 | return self._length 265 | 266 | @always_inline 267 | fn __getitem__(self, index: Int) raises -> FlxValue: 268 | if index < 0 or index >= self._length: 269 | raise "Bad index: " + String(index) + ". Lenght: " + String(self._length) 270 | if self._type.is_typed_vector(): 271 | return FlxValue( 272 | self._bytes.offset(index * self._byte_width.to_int()), 273 | self._byte_width, 274 | 1, 275 | self._type.typed_vector_element_type() 276 | ) 277 | if self._type.is_fixed_typed_vector(): 278 | return FlxValue( 279 | self._bytes.offset(index * self._byte_width.to_int()), 280 | self._byte_width, 281 | 1, 282 | self._type.fixed_typed_vector_element_type() 283 | ) 284 | if self._type == ValueType.Vector: 285 | var packed_type = self._bytes.offset(self._length * self._byte_width.to_int() + index).load() 286 | return FlxValue( 287 | self._bytes.offset(index * self._byte_width.to_int()), 288 | self._byte_width, 289 | packed_type 290 | ) 291 | raise "Is not an expected vector type. Type id: " + String(self._type.value) 292 | 293 | struct FlxMapValue(Sized): 294 | var _bytes: DTypePointer[DType.uint8] 295 | var _byte_width: UInt8 296 | var _length: Int 297 | 298 | @always_inline 299 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], byte_width: UInt8, length: Int): 300 | self._bytes = bytes 301 | self._byte_width = byte_width 302 | self._length = length 303 | 304 | @always_inline 305 | fn __len__(self) -> Int: 306 | return self._length 307 | 308 | @always_inline 309 | fn __getitem__(self, key: String) raises -> FlxValue: 310 | var index = self.key_index(key) 311 | if index < 0: 312 | raise "Key " + key + " could not be found" 313 | return self.values()[index] 314 | 315 | @always_inline 316 | fn keys(self) -> FlxVecValue: 317 | var p1 = self._bytes.offset(self._byte_width.to_int() * -3) 318 | var p2 = jump_to_indirect(p1, self._byte_width) 319 | var byte_width = read_uint(p1.offset(self._byte_width.to_int()), self._byte_width) 320 | return FlxVecValue(p2, byte_width, ValueType.VectorKey, self._length) 321 | 322 | @always_inline 323 | fn values(self) -> FlxVecValue: 324 | return FlxVecValue(self._bytes, self._byte_width, ValueType.Vector, self._length) 325 | 326 | @always_inline 327 | fn key_index(self, key: String) raises -> Int: 328 | var a = key._as_ptr().bitcast[DType.uint8]() 329 | var keys = self.keys() 330 | var low = 0 331 | var high = self._length - 1 332 | while low <= high: 333 | var mid = (low + high) >> 1 334 | var mid_key = keys[mid] 335 | var b = jump_to_indirect(mid_key._bytes, mid_key._parent_byte_width) 336 | var diff = cmp(a, b, len(key) + 1) 337 | if diff == 0: 338 | return mid 339 | if diff < 0: 340 | high = mid - 1 341 | else: 342 | low = mid + 1 343 | return -1 344 | 345 | 346 | @always_inline 347 | fn jump_to_indirect(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> DTypePointer[DType.uint8]: 348 | return bytes.offset(-read_uint(bytes, byte_width)) 349 | 350 | @always_inline 351 | fn read_uint(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> Int: 352 | if byte_width < 4: 353 | if byte_width == 1: 354 | return bytes.load().to_int() 355 | else: 356 | @parameter 357 | if is_be: 358 | return bswap(bytes.bitcast[DType.uint16]().load()).to_int() 359 | else: 360 | return bytes.bitcast[DType.uint16]().load().to_int() 361 | else: 362 | if byte_width == 4: 363 | @parameter 364 | if is_be: 365 | return bswap(bytes.bitcast[DType.uint32]().load()).to_int() 366 | else: 367 | return bytes.bitcast[DType.uint32]().load().to_int() 368 | else: 369 | @parameter 370 | if is_be: 371 | return bswap(bytes.bitcast[DType.uint64]().load()).to_int() 372 | else: 373 | return bytes.bitcast[DType.uint64]().load().to_int() 374 | 375 | @always_inline 376 | fn read_int(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> Int: 377 | if byte_width < 4: 378 | if byte_width == 1: 379 | return bytes.bitcast[DType.int8]().load().to_int() 380 | else: 381 | @parameter 382 | if is_be: 383 | return bswap(bytes.bitcast[DType.int16]().load()).to_int() 384 | else: 385 | return bytes.bitcast[DType.int16]().load().to_int() 386 | else: 387 | if byte_width == 4: 388 | @parameter 389 | if is_be: 390 | return bswap(bytes.bitcast[DType.int32]().load()).to_int() 391 | else: 392 | return bytes.bitcast[DType.int32]().load().to_int() 393 | else: 394 | @parameter 395 | if is_be: 396 | return bswap(bytes.bitcast[DType.int64]().load()).to_int() 397 | else: 398 | return bytes.bitcast[DType.int64]().load().to_int() 399 | 400 | @always_inline 401 | fn read_float(bytes: DTypePointer[DType.uint8], byte_width: UInt8) raises -> Float64: 402 | if byte_width == 8: 403 | @parameter 404 | if is_be: 405 | return bswap(bytes.bitcast[DType.float64]().load()) 406 | else: 407 | return bytes.bitcast[DType.float64]().load() 408 | if byte_width == 4: 409 | @parameter 410 | if is_be: 411 | return bswap(bytes.bitcast[DType.float32]().load()).cast[DType.float64]() 412 | else: 413 | return bytes.bitcast[DType.float32]().load().cast[DType.float64]() 414 | if byte_width == 2: 415 | @parameter 416 | if is_be: 417 | return bswap(bytes.bitcast[DType.float16]().load()).cast[DType.float64]() 418 | else: 419 | return bytes.bitcast[DType.float16]().load().cast[DType.float64]() 420 | raise "Unexpected byte width: " + String(byte_width) 421 | 422 | @always_inline 423 | fn cmp(a: DTypePointer[DType.uint8], b: DTypePointer[DType.uint8], length: Int) -> Int: 424 | for i in range(length): 425 | var diff = a.load(i).to_int() - b.load(i).to_int() 426 | if diff != 0: 427 | return diff 428 | return 0 429 | -------------------------------------------------------------------------------- /flx2/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .flx_builder import FlxMap, FlxVec 2 | from .flx_value import FlxValue 3 | from .flx_buffer import flx, flx_blob, flx_null -------------------------------------------------------------------------------- /flx2/cache.mojo: -------------------------------------------------------------------------------- 1 | from memory import memset_zero, memcpy 2 | from math.bit import bit_length 3 | from .data_types import StackValue 4 | 5 | @value 6 | struct Key(CollectionElement): 7 | var pointer: DTypePointer[DType.uint8] 8 | var size: Int 9 | 10 | fn __init__(inout self, pointer: DTypePointer[DType.uint8], size: Int): 11 | var cp = DTypePointer[DType.uint8].alloc(size) 12 | memcpy(cp, pointer, size) 13 | self.pointer = cp 14 | self.size = size 15 | 16 | # alias Key = (DTypePointer[DType.uint8], Int) 17 | alias Keys = List[Key] 18 | alias Values = List[StackValue] 19 | 20 | struct _CacheStackValue(Movable, Copyable): 21 | var keys: Keys 22 | var values: Values 23 | var key_map: DTypePointer[DType.uint32] 24 | var count: Int 25 | var capacity: Int 26 | 27 | fn __init__(inout self): 28 | self.count = 0 29 | self.capacity = 16 30 | self.keys = Keys(capacity=self.capacity) 31 | self.values = Values(capacity=self.capacity) 32 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 33 | memset_zero(self.key_map, self.capacity) 34 | 35 | fn __moveinit__(inout self, owned other: Self): 36 | self.count = other.count 37 | self.capacity = other.capacity 38 | self.values = other.values^ 39 | self.key_map = other.key_map 40 | self.keys = other.keys^ 41 | 42 | fn __copyinit__(inout self, other: Self): 43 | self.count = other.count 44 | self.capacity = other.capacity 45 | var keys_count = len(other.keys) 46 | 47 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 48 | memcpy(self.key_map, other.key_map, self.capacity) 49 | # self.keys = other.keys 50 | # self.values = other.values 51 | self.keys = Keys(capacity=keys_count) 52 | self.values = Values(capacity=keys_count) 53 | for i in range(keys_count): 54 | var key = other.keys[i] 55 | var p = key.pointer 56 | var size = key.size 57 | var cp = DTypePointer[DType.uint8].alloc(size) 58 | memcpy(cp, p, size) 59 | var new_key = Key(cp, size) 60 | self.keys.append(new_key) 61 | self.values.append(other.values[i]) 62 | 63 | fn __del__(owned self): 64 | self.key_map.free() 65 | for i in range(len(self.keys)): 66 | var key = self.keys[i] 67 | key.pointer.free() 68 | 69 | fn put(inout self, key: Key, value: StackValue): 70 | if self.count / self.capacity >= 0.8: 71 | self._rehash() 72 | self._put(key, value, -1) 73 | 74 | fn _rehash(inout self): 75 | var old_mask_capacity = self.capacity >> 3 76 | self.key_map.free() 77 | self.capacity <<= 1 78 | var mask_capacity = self.capacity >> 3 79 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 80 | memset_zero(self.key_map, self.capacity) 81 | 82 | for i in range(len(self.keys)): 83 | self._put(self.keys[i], self.values[i], i + 1) 84 | 85 | fn _put(inout self, key: Key, value: StackValue, rehash_index: Int): 86 | var key_hash = self._hash(key) 87 | var modulo_mask = self.capacity - 1 88 | var key_map_index = (key_hash & modulo_mask).to_int() 89 | while True: 90 | var key_index = self.key_map.offset(key_map_index).load().to_int() 91 | if key_index == 0: 92 | var new_key_index: Int 93 | if rehash_index == -1: 94 | self.keys.append(key) 95 | self.values.append(value) 96 | self.count += 1 97 | new_key_index = len(self.keys) 98 | else: 99 | new_key_index = rehash_index 100 | self.key_map.offset(key_map_index).store(UInt32(new_key_index)) 101 | return 102 | 103 | var other_key = self.keys[key_index - 1] 104 | if self._eq(other_key, key): 105 | self.values[key_index - 1] = value 106 | return 107 | 108 | key_map_index = (key_map_index + 1) & modulo_mask 109 | 110 | fn _hash(self, key: Key) -> UInt32: 111 | var hash: UInt32 = 0 112 | var bytes = key.pointer 113 | var count = key.size 114 | while count >= 4: 115 | var c = bytes.bitcast[DType.uint32]().load() 116 | hash = _hash_word32(hash, c) 117 | bytes = bytes.offset(4) 118 | count -= 4 119 | if count >= 2: 120 | var c = bytes.bitcast[DType.uint16]().load().cast[DType.uint32]() 121 | hash = _hash_word32(hash, c) 122 | bytes = bytes.offset(2) 123 | count -= 2 124 | if count > 0: 125 | var c = bytes.load().cast[DType.uint32]() 126 | hash = _hash_word32(hash, c) 127 | return hash 128 | 129 | fn _eq(self, a: Key, b: Key) -> Bool: 130 | var bytes_a = a.pointer 131 | var bytes_b = b.pointer 132 | var count_a = a.size 133 | var count_b = b.size 134 | if count_a != count_b: 135 | return False 136 | var count = count_a 137 | while count >= 4: 138 | if bytes_a.bitcast[DType.uint32]().load() != bytes_b.bitcast[DType.uint32]().load(): 139 | return False 140 | bytes_a = bytes_a.offset(4) 141 | bytes_b = bytes_b.offset(4) 142 | count -= 4 143 | if count >= 2: 144 | if bytes_a.bitcast[DType.uint16]().load() != bytes_b.bitcast[DType.uint16]().load(): 145 | return False 146 | bytes_a = bytes_a.offset(2) 147 | bytes_b = bytes_b.offset(2) 148 | count -= 2 149 | if count > 0: 150 | return bytes_a.load() == bytes_b.load() 151 | return True 152 | 153 | fn get(self, key: Key, default: StackValue) -> StackValue: 154 | var key_hash = self._hash(key) 155 | var modulo_mask = self.capacity - 1 156 | var key_map_index = (key_hash & modulo_mask).to_int() 157 | while True: 158 | var key_index = self.key_map.offset(key_map_index).load().to_int() 159 | if key_index == 0: 160 | return default 161 | var other_key = self.keys[key_index - 1] 162 | if self._eq(other_key, key): 163 | return self.values[key_index - 1] 164 | key_map_index = (key_map_index + 1) & modulo_mask 165 | 166 | from math.math import rotate_bits_left 167 | 168 | alias ROTATE = 5 169 | alias SEED32 = 0x9e_37_79_b9 170 | 171 | @always_inline 172 | fn _hash_word32(value: UInt32, word: UInt32) -> UInt32: 173 | return (rotate_bits_left[ROTATE](value) ^ word) * SEED32 174 | 175 | fn _key_string(key: Key) -> String: 176 | var bytes = key.pointer 177 | var count = key.size 178 | var result: String = "" 179 | for i in range(count): 180 | result += chr(bytes.load(i).to_int()) 181 | return result 182 | 183 | fn _key_int_string(key: Key) -> String: 184 | var bytes = key.pointer 185 | var count = key.size 186 | var result: String = "" 187 | for i in range(count): 188 | result += String(bytes.load(i).to_int()) 189 | return result 190 | 191 | 192 | struct _CacheStringOrKey[is_string: Bool = True](Movable, Copyable): 193 | # offsets and counts 194 | var ocs: List[(Int, Int)] 195 | var key_map: DTypePointer[DType.uint32] 196 | var count: Int 197 | var capacity: Int 198 | 199 | fn __init__(inout self): 200 | self.count = 0 201 | self.capacity = 16 202 | self.ocs = List[(Int, Int)](capacity=self.capacity) 203 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 204 | memset_zero(self.key_map, self.capacity) 205 | 206 | fn __moveinit__(inout self, owned other: Self): 207 | self.count = other.count 208 | self.capacity = other.capacity 209 | self.ocs = other.ocs^ 210 | self.key_map = other.key_map 211 | 212 | fn __copyinit__(inout self, other: Self): 213 | self.count = other.count 214 | self.capacity = other.capacity 215 | # TODO: copies elements one by one because otherwise it throws a core dump 216 | # self.ocs = other.ocs 217 | self.ocs = List[(Int, Int)](capacity=self.capacity) 218 | for i in range(self.capacity): 219 | self.ocs[i] = other.ocs[i] 220 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 221 | memcpy(self.key_map, other.key_map, self.capacity) 222 | 223 | fn __del__(owned self): 224 | self.key_map.free() 225 | 226 | fn put(inout self, oc: (Int, Int), pointer: DTypePointer[DType.uint8]): 227 | if self.count / self.capacity >= 0.8: 228 | self._rehash(pointer) 229 | self._put(oc, pointer, -1) 230 | 231 | fn get(self, bc: (DTypePointer[DType.uint8], Int), pointer: DTypePointer[DType.uint8]) -> Int: 232 | var bytes = bc.get[0, DTypePointer[DType.uint8]]() 233 | var count = bc.get[1, Int]() 234 | var key_hash = self._hash(bytes, count) 235 | var modulo_mask = self.capacity - 1 236 | var key_map_index = (key_hash & modulo_mask).to_int() 237 | while True: 238 | var key_index = self.key_map.offset(key_map_index).load().to_int() 239 | if key_index == 0: 240 | return -1 241 | var other_oc = self.ocs[key_index - 1] 242 | @parameter 243 | if is_string: 244 | var count = other_oc.get[1, Int]() 245 | var size_byte_length = 1 << (bit_length(UInt64(count << 2)) >> 3) 246 | if self._eq(count, other_oc.get[1, Int](), bytes, pointer.offset(other_oc.get[0, Int]() + size_byte_length)): 247 | return other_oc.get[0, Int]() 248 | else: 249 | if self._eq(count, other_oc.get[1, Int](), bytes, pointer.offset(other_oc.get[0, Int]())): 250 | return other_oc.get[0, Int]() 251 | key_map_index = (key_map_index + 1) & modulo_mask 252 | 253 | fn _rehash(inout self, pointer: DTypePointer[DType.uint8]): 254 | self.key_map.free() 255 | self.capacity <<= 1 256 | self.key_map = DTypePointer[DType.uint32].alloc(self.capacity) 257 | memset_zero(self.key_map, self.capacity) 258 | for i in range(len(self.ocs)): 259 | self._put(self.ocs[i], pointer, i + 1) 260 | 261 | fn _put(inout self, oc: (Int, Int), pointer: DTypePointer[DType.uint8], rehash_index: Int): 262 | var count = oc.get[1, Int]() 263 | var bytes: DTypePointer[DType.uint8] 264 | 265 | @parameter 266 | if is_string: 267 | var size_byte_length = 1 << (bit_length(UInt64(count << 2)) >> 3) 268 | bytes = pointer.offset(oc.get[0, Int]() + size_byte_length) 269 | else: 270 | bytes = pointer.offset(oc.get[0, Int]()) 271 | 272 | var key_hash = self._hash(bytes, count) 273 | var modulo_mask = self.capacity - 1 274 | var key_map_index = (key_hash & modulo_mask).to_int() 275 | while True: 276 | var key_index = self.key_map.offset(key_map_index).load().to_int() 277 | if key_index == 0: 278 | var new_key_index: Int 279 | if rehash_index == -1: 280 | self.ocs.append(oc) 281 | self.count += 1 282 | new_key_index = len(self.ocs) 283 | else: 284 | new_key_index = rehash_index 285 | self.key_map.offset(key_map_index).store(UInt32(new_key_index)) 286 | return 287 | 288 | var other_ol = self.ocs[key_index - 1] 289 | if self._eq(count, other_ol.get[1, Int](), bytes, pointer.offset(other_ol.get[0, Int]())): 290 | return 291 | 292 | key_map_index = (key_map_index + 1) & modulo_mask 293 | 294 | fn _hash(self, _bytes: DTypePointer[DType.uint8], _count: Int) -> UInt32: 295 | var bytes = _bytes 296 | var count = _count 297 | var hash: UInt32 = 0 298 | while count >= 4: 299 | var c = bytes.bitcast[DType.uint32]().load() 300 | hash = _hash_word32(hash, c) 301 | bytes = bytes.offset(4) 302 | count -= 4 303 | if count >= 2: 304 | var c = bytes.bitcast[DType.uint16]().load().cast[DType.uint32]() 305 | hash = _hash_word32(hash, c) 306 | bytes = bytes.offset(2) 307 | count -= 2 308 | if count > 0: 309 | var c = bytes.load().cast[DType.uint32]() 310 | hash = _hash_word32(hash, c) 311 | return hash 312 | 313 | fn _eq(self, _count_a: Int, _count_b: Int, _bytes_a: DTypePointer[DType.uint8], _bytes_b: DTypePointer[DType.uint8]) -> Bool: 314 | var bytes_a = _bytes_a 315 | var bytes_b = _bytes_b 316 | var count_a = _count_a 317 | var count_b = _count_b 318 | if count_a != count_b: 319 | return False 320 | var count = count_a 321 | while count >= 8: 322 | if bytes_a.bitcast[DType.uint64]().load() != bytes_b.bitcast[DType.uint64]().load(): 323 | return False 324 | bytes_a = bytes_a.offset(8) 325 | bytes_b = bytes_b.offset(8) 326 | count -= 8 327 | if count >= 4: 328 | if bytes_a.bitcast[DType.uint32]().load() != bytes_b.bitcast[DType.uint32]().load(): 329 | return False 330 | bytes_a = bytes_a.offset(4) 331 | bytes_b = bytes_b.offset(4) 332 | count -= 4 333 | if count >= 2: 334 | if bytes_a.bitcast[DType.uint16]().load() != bytes_b.bitcast[DType.uint16]().load(): 335 | return False 336 | bytes_a = bytes_a.offset(2) 337 | bytes_b = bytes_b.offset(2) 338 | count -= 2 339 | if count > 0: 340 | return bytes_a.load() == bytes_b.load() 341 | return True 342 | -------------------------------------------------------------------------------- /flx2/data_types.mojo: -------------------------------------------------------------------------------- 1 | from math.limit import max_finite, min_finite 2 | from math.bit import bswap, bit_length 3 | from memory.unsafe import bitcast 4 | from sys.info import is_big_endian 5 | 6 | alias is_be = is_big_endian() 7 | 8 | @value 9 | @register_passable("trivial") 10 | struct ValueBitWidth: 11 | alias width8 = ValueBitWidth(0) 12 | alias width16 = ValueBitWidth(1) 13 | alias width32 = ValueBitWidth(2) 14 | alias width64 = ValueBitWidth(3) 15 | 16 | alias min8 = min_finite[DType.int8]().to_int() 17 | alias max8 = max_finite[DType.int8]().to_int() 18 | alias umax8 = max_finite[DType.uint8]().cast[DType.uint64]() 19 | alias min16 = min_finite[DType.int16]().to_int() 20 | alias max16 = max_finite[DType.int16]().to_int() 21 | alias umax16 = max_finite[DType.uint16]().cast[DType.uint64]() 22 | alias min32 = min_finite[DType.int32]().to_int() 23 | alias max32 = max_finite[DType.int32]().to_int() 24 | alias umax32 = max_finite[DType.uint32]().cast[DType.uint64]() 25 | 26 | var value: UInt8 27 | 28 | @always_inline 29 | fn __lt__(self, other: ValueBitWidth) -> Bool: 30 | return self.value < other.value 31 | 32 | @always_inline 33 | fn __le__(self, other: ValueBitWidth) -> Bool: 34 | return self.value <= other.value 35 | 36 | @always_inline 37 | fn __eq__(self, other: ValueBitWidth) -> Bool: 38 | return self.value == other.value 39 | 40 | @always_inline 41 | @staticmethod 42 | fn of[D: DType](n: SIMD[D, 1]) -> ValueBitWidth: 43 | @parameter 44 | if D == DType.uint8 or D == DType.int8 or D == DType.bool: 45 | return ValueBitWidth.width8 46 | elif D == DType.uint16 or D == DType.int16 or D == DType.float16: 47 | return ValueBitWidth.width16 48 | elif D == DType.uint32 or D == DType.int32 or D == DType.float32: 49 | return ValueBitWidth.width32 50 | else: 51 | return ValueBitWidth.width64 52 | 53 | @always_inline 54 | @staticmethod 55 | fn of(n: Int) -> ValueBitWidth: 56 | return ValueBitWidth((bit_length(Int64(n)) >> 3).to_int()) 57 | 58 | @always_inline 59 | @staticmethod 60 | fn of(n: UInt64) -> ValueBitWidth: 61 | return ValueBitWidth((bit_length(n) >> 3).to_int()) 62 | 63 | @always_inline 64 | fn padding_size(buffer_size: UInt64, scalar_size: UInt64) -> UInt64: 65 | return (~buffer_size + 1) & (scalar_size - 1) 66 | 67 | @value 68 | @register_passable("trivial") 69 | struct ValueType: 70 | alias Null = ValueType(0) 71 | alias Int = ValueType(1) 72 | alias UInt = ValueType(2) 73 | alias Float = ValueType(3) 74 | alias Key = ValueType(4) 75 | alias String = ValueType(5) 76 | alias IndirectInt = ValueType(6) 77 | alias IndirectUInt = ValueType(7) 78 | alias IndirectFloat = ValueType(8) 79 | alias Map = ValueType(9) 80 | alias Vector = ValueType(10) 81 | alias VectorInt = ValueType(11) 82 | alias VectorUInt = ValueType(12) 83 | alias VectorFloat = ValueType(13) 84 | alias VectorKey = ValueType(14) 85 | alias VectorString = ValueType(15) 86 | alias VectorInt2 = ValueType(16) 87 | alias VectorUInt2 = ValueType(17) 88 | alias VectorFloat2 = ValueType(18) 89 | alias VectorInt3 = ValueType(19) 90 | alias VectorUInt3 = ValueType(20) 91 | alias VectorFloat3 = ValueType(21) 92 | alias VectorInt4 = ValueType(22) 93 | alias VectorUInt4 = ValueType(23) 94 | alias VectorFloat4 = ValueType(24) 95 | alias Blob = ValueType(25) 96 | alias Bool = ValueType(26) 97 | alias VectorBool = ValueType(36) 98 | 99 | alias NullPackedType = 0 100 | 101 | var value: UInt8 102 | 103 | @always_inline 104 | fn __eq__(self, other: Self) -> Bool: 105 | return self.value == other.value 106 | 107 | @always_inline 108 | fn __ne__(self, other: Self) -> Bool: 109 | return self.value != other.value 110 | 111 | @always_inline 112 | fn __lt__(self, other: Self) -> Bool: 113 | return self.value < other.value 114 | 115 | @always_inline 116 | fn __le__(self, other: Self) -> Bool: 117 | return self.value <= other.value 118 | 119 | @always_inline 120 | fn __sub__(self, other: Self) -> Self: 121 | return Self(self.value - other.value) 122 | 123 | @always_inline 124 | fn __add__(self, other: Self) -> Self: 125 | return Self(self.value + other.value) 126 | 127 | @always_inline 128 | fn __add__(self, other: UInt8) -> Self: 129 | return Self(self.value + other.value) 130 | 131 | @always_inline 132 | fn __mod__(self, other: UInt8) -> Self: 133 | return Self(self.value % other) 134 | 135 | @always_inline 136 | fn __floordiv__(self, other: UInt8) -> Self: 137 | return Self(self.value // other) 138 | 139 | @always_inline 140 | fn __lshift__(self, other: UInt8) -> Self: 141 | return Self(self.value << other) 142 | 143 | @always_inline 144 | fn is_inline(self) -> Bool: 145 | return self == ValueType.Bool or self <= ValueType.Float 146 | 147 | @always_inline 148 | fn is_typed_vector_element(self) -> Bool: 149 | return self == ValueType.Bool or ValueType.Int <= self <= ValueType.String 150 | 151 | @always_inline 152 | fn is_typed_vector(self) -> Bool: 153 | return self == ValueType.VectorBool or ValueType.VectorInt <= self <= ValueType.VectorString 154 | 155 | @always_inline 156 | fn is_fixed_typed_vector(self) -> Bool: 157 | return ValueType.VectorInt2 <= self <= ValueType.VectorFloat4 158 | 159 | @always_inline 160 | fn is_a_vector(self) -> Bool: 161 | return self == ValueType.VectorBool or ValueType.Vector <= self <= ValueType.VectorFloat4 162 | 163 | @always_inline 164 | fn to_typed_vector(self, length: UInt8) raises -> ValueType: 165 | if length == 0: 166 | return self - ValueType.Int + ValueType.VectorInt 167 | if length == 2: 168 | return self - ValueType.Int + ValueType.VectorInt2 169 | if length == 3: 170 | return self - ValueType.Int + ValueType.VectorInt3 171 | if length == 4: 172 | return self - ValueType.Int + ValueType.VectorInt4 173 | raise "Unexpected length " + String(length) 174 | 175 | @always_inline 176 | fn typed_vector_element_type(self) -> ValueType: 177 | return self - ValueType.VectorInt + ValueType.Int 178 | 179 | @always_inline 180 | fn fixed_typed_vector_element_type(self) -> ValueType: 181 | return ((self - ValueType.VectorInt2) % 3) + ValueType.Int 182 | 183 | @always_inline 184 | fn fixed_typed_vector_element_size(self) -> Int: 185 | return ((self - ValueType.VectorInt2) // 3).value.to_int() + 2 186 | 187 | @always_inline 188 | fn packed_type(self, bit_width: ValueBitWidth) -> UInt8: 189 | return (self << 2).value | bit_width.value 190 | 191 | @staticmethod 192 | @always_inline 193 | fn of[D: DType]() -> Self: 194 | @parameter 195 | if D == DType.uint8 or D == DType.uint16 or D == DType.uint32 or D == DType.uint64: 196 | return ValueType.UInt 197 | elif D == DType.float16 or D == DType.float32 or D == DType.float64: 198 | return ValueType.Float 199 | elif D == DType.bool: 200 | return ValueType.Bool 201 | else: 202 | return ValueType.Int 203 | 204 | 205 | @value 206 | @register_passable("trivial") 207 | struct StackValue(CollectionElement): 208 | var value: SIMD[DType.uint8, 8] 209 | var width: ValueBitWidth 210 | var type: ValueType 211 | 212 | alias Null = StackValue(0, ValueBitWidth.width8, ValueType.Null) 213 | 214 | @staticmethod 215 | @always_inline 216 | fn of[D: DType](v: SIMD[D, 1]) -> Self: 217 | # TODO: make it better 218 | var value = SIMD[DType.uint8, 8](0) 219 | @parameter 220 | if D == DType.bool: 221 | value[0] = v.cast[DType.uint8]() 222 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 223 | elif D == DType.uint8 or D == DType.int8: 224 | var v1 = bitcast[DType.uint8, 1](v) 225 | value[0] = v1[0] 226 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 227 | elif D == DType.uint16 or D == DType.int16 or D == DType.float16: 228 | var v1 = bitcast[DType.uint8, 2](v) 229 | @parameter 230 | if is_be: 231 | v1 = bswap(v1) 232 | value[0] = v1[0] 233 | value[1] = v1[1] 234 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 235 | elif D == DType.uint32 or D == DType.int32 or D == DType.float32: 236 | var v1 = bitcast[DType.uint8, 4](v) 237 | @parameter 238 | if is_be: 239 | v1 = bswap(v1) 240 | value[0] = v1[0] 241 | value[1] = v1[1] 242 | value[2] = v1[2] 243 | value[3] = v1[3] 244 | return StackValue(value, ValueBitWidth.of(v), ValueType.of[D]()) 245 | else: 246 | var v1 = bitcast[DType.uint8, 8](v) 247 | @parameter 248 | if is_be: 249 | v1 = bswap(v1) 250 | return StackValue(bitcast[DType.uint8, 8](v), ValueBitWidth.of(v), ValueType.of[D]()) 251 | 252 | @staticmethod 253 | @always_inline 254 | fn of(v: Int) -> Self: 255 | var value = bitcast[DType.uint8, 8](Int64(v)) 256 | return StackValue(value, ValueBitWidth.of(v), ValueType.Int) 257 | 258 | @always_inline 259 | fn stored_width(self, bit_width: ValueBitWidth = ValueBitWidth.width8) -> ValueBitWidth: 260 | if self.type.is_inline(): 261 | return bit_width if self.width < bit_width else self.width 262 | return self.width 263 | 264 | @always_inline 265 | fn stored_packed_type(self, bit_width: ValueBitWidth = ValueBitWidth.width8) -> UInt8: 266 | return self.type.packed_type(self.stored_width(bit_width)) 267 | 268 | @always_inline 269 | fn element_width(self, size: UInt64, index: Int) raises -> ValueBitWidth: 270 | if self.type.is_inline(): 271 | return self.width 272 | for i in range(4): 273 | var width = UInt64(1 << i) 274 | var offset_loc = size + padding_size(size, width) + UInt64(index * width) 275 | var offset = offset_loc - self.as_uint() 276 | var bit_width = ValueBitWidth.of(offset.to_int()) 277 | if (1 << bit_width.value).cast[DType.uint64]() == width: 278 | return bit_width 279 | 280 | raise "Element with size " + String(size) + " and index " + String(index) + " is of unknown width" 281 | 282 | @always_inline 283 | fn as_float(self) -> Float64: 284 | return bitcast[DType.float64, 1](self.value) 285 | 286 | @always_inline 287 | fn as_int(self) -> Int64: 288 | return bitcast[DType.int64, 1](self.value) 289 | 290 | @always_inline 291 | fn as_uint(self) -> UInt64: 292 | return bitcast[DType.uint64, 1](self.value) 293 | 294 | fn to_value(self, byte_width: UInt64) -> SIMD[DType.uint8, 8]: 295 | var self_byte_width = (1 << self.width.value).cast[DType.uint64]() 296 | if self_byte_width == byte_width or self.type == ValueType.UInt: 297 | return self.value 298 | else: 299 | if self_byte_width > 2: 300 | if self_byte_width == 4: 301 | var v = self.value.slice[4]() 302 | if byte_width == 2: 303 | if self.type == ValueType.Float: 304 | return StackValue.of(bitcast[DType.float32, 1](v).cast[DType.float16]()).value 305 | else: 306 | return StackValue.of(bitcast[DType.int32, 1](v).cast[DType.int16]()).value 307 | else: 308 | # byte_width == 8 309 | if self.type == ValueType.Float: 310 | return StackValue.of(bitcast[DType.float32, 1](v).cast[DType.float64]()).value 311 | else: 312 | return StackValue.of(bitcast[DType.int32, 1](v).cast[DType.int64]()).value 313 | else: 314 | # self_byte_width == 8 315 | if byte_width == 2: 316 | if self.type == ValueType.Float: 317 | return StackValue.of(bitcast[DType.float64, 1](self.value).cast[DType.float16]()).value 318 | else: 319 | return StackValue.of(bitcast[DType.int64, 1](self.value).cast[DType.int16]()).value 320 | else: 321 | if self.type == ValueType.Float: 322 | return StackValue.of(bitcast[DType.float64, 1](self.value).cast[DType.float32]()).value 323 | else: 324 | return StackValue.of(bitcast[DType.int64, 1](self.value).cast[DType.int32]()).value 325 | else: 326 | if self_byte_width == 2: 327 | var v = self.value.slice[2]() 328 | if byte_width == 4: 329 | if self.type == ValueType.Float: 330 | return StackValue.of(bitcast[DType.float16, 1](v).cast[DType.float32]()).value 331 | else: 332 | return StackValue.of(bitcast[DType.int16, 1](v).cast[DType.int32]()).value 333 | else: 334 | if self.type == ValueType.Float: 335 | return StackValue.of(bitcast[DType.float16, 1](v).cast[DType.float64]()).value 336 | else: 337 | return StackValue.of(bitcast[DType.int16, 1](v).cast[DType.int64]()).value 338 | else: 339 | if byte_width > 2: 340 | if byte_width == 4: 341 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int32]()).value 342 | else: 343 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int64]()).value 344 | else: 345 | if byte_width == 2: 346 | return StackValue.of(bitcast[DType.int8, 1](self.value[0]).cast[DType.int16]()).value 347 | else: 348 | return self.value 349 | 350 | @always_inline 351 | fn is_float32(self) -> Bool: 352 | return self.type == ValueType.Float and self.width == ValueBitWidth.width32 353 | 354 | @always_inline 355 | fn is_offset(self) -> Bool: 356 | return not self.type.is_inline() 357 | 358 | fn __ne__(self, other: Self) -> Bool: 359 | return self.value != other.value 360 | or self.width.value != other.width.value 361 | or self.type != other.type 362 | -------------------------------------------------------------------------------- /flx2/flx_buffer.mojo: -------------------------------------------------------------------------------- 1 | from .data_types import StackValue, ValueBitWidth, padding_size, ValueType 2 | from .cache import _CacheStackValue, Key, _CacheStringOrKey 3 | from memory import memcpy, memset_zero 4 | from memory.unsafe import bitcast 5 | from math import max 6 | 7 | fn flx_null() -> (DTypePointer[DType.uint8], Int): 8 | var buffer = FlxBuffer(16) 9 | buffer.add_null() 10 | return finish_ignoring_excetion(buffer^) 11 | 12 | fn flx(v: Int) -> (DTypePointer[DType.uint8], Int): 13 | var buffer = FlxBuffer(16) 14 | buffer.add(v) 15 | return finish_ignoring_excetion(buffer^) 16 | 17 | fn flx[D: DType](v: SIMD[D, 1]) -> (DTypePointer[DType.uint8], Int): 18 | var buffer = FlxBuffer(16) 19 | buffer.add(v) 20 | return finish_ignoring_excetion(buffer^) 21 | 22 | fn flx(v: String) -> (DTypePointer[DType.uint8], Int): 23 | var buffer = FlxBuffer(len(v) + 16) 24 | buffer.add(v) 25 | return finish_ignoring_excetion(buffer^) 26 | 27 | fn flx_blob(v: DTypePointer[DType.uint8], length: Int) -> (DTypePointer[DType.uint8], Int): 28 | var buffer = FlxBuffer(length + 32) 29 | buffer.blob(v, length) 30 | return finish_ignoring_excetion(buffer^) 31 | 32 | fn flx[D: DType](v: DTypePointer[D], length: Int) -> (DTypePointer[DType.uint8], Int): 33 | var buffer = FlxBuffer(length * sizeof[D]() + 1024) 34 | buffer.add(v, length) 35 | return finish_ignoring_excetion(buffer^) 36 | 37 | struct FlxBuffer[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Copyable, Movable): 38 | var _stack: List[StackValue] 39 | var _stack_positions: List[Int] 40 | var _stack_is_vector: List[SIMD[DType.bool, 1]] 41 | var _bytes: DTypePointer[DType.uint8] 42 | var _size: UInt64 43 | var _offset: UInt64 44 | var _finished: Bool 45 | var _string_cache: _CacheStringOrKey 46 | var _key_cache: _CacheStringOrKey[False] 47 | var _keys_vec_cache: _CacheStackValue 48 | var _reference_cache: _CacheStackValue 49 | 50 | fn __init__(inout self, size: UInt64 = 1 << 11): 51 | self._size = size 52 | self._stack = List[StackValue]() 53 | self._stack_positions = List[Int]() 54 | self._stack_is_vector = List[SIMD[DType.bool, 1]]() 55 | self._bytes = DTypePointer[DType.uint8].alloc(size.to_int()) 56 | self._offset = 0 57 | self._finished = False 58 | self._string_cache = _CacheStringOrKey() 59 | self._key_cache = _CacheStringOrKey[is_string=False]() 60 | self._keys_vec_cache = _CacheStackValue() 61 | self._reference_cache = _CacheStackValue() 62 | 63 | fn __moveinit__(inout self, owned other: Self): 64 | self._size = other._size 65 | self._stack = other._stack^ 66 | self._stack_positions = other._stack_positions^ 67 | self._stack_is_vector = other._stack_is_vector^ 68 | self._bytes = other._bytes 69 | self._offset = other._offset 70 | self._finished = other._finished 71 | self._string_cache = other._string_cache^ 72 | self._key_cache = other._key_cache^ 73 | self._keys_vec_cache = other._keys_vec_cache^ 74 | self._reference_cache = other._reference_cache^ 75 | 76 | fn __copyinit__(inout self, other: Self): 77 | self._size = other._size 78 | self._stack = other._stack 79 | self._stack_positions = other._stack_positions 80 | self._stack_is_vector = other._stack_is_vector 81 | self._bytes = DTypePointer[DType.uint8].alloc(other._size.to_int()) 82 | memcpy(self._bytes, other._bytes, other._offset.to_int()) 83 | self._offset = other._offset 84 | self._finished = other._finished 85 | self._string_cache = other._string_cache 86 | self._key_cache = other._key_cache 87 | self._keys_vec_cache = other._keys_vec_cache 88 | self._reference_cache = other._reference_cache 89 | 90 | fn __del__(owned self): 91 | if not self._finished: 92 | self._bytes.free() 93 | 94 | fn add_null(inout self): 95 | self._stack.append(StackValue.Null) 96 | 97 | fn add[D: DType](inout self, value: SIMD[D, 1]): 98 | self._stack.append(StackValue.of(value)) 99 | 100 | fn add(inout self, value: Int): 101 | self._stack.append(StackValue.of(value)) 102 | 103 | fn add(inout self, value: String): 104 | self._add_string[as_key=False](value) 105 | 106 | fn key(inout self, value: String): 107 | self._add_string[as_key=True](value) 108 | 109 | fn _add_string[as_key: Bool](inout self, value: String): 110 | var byte_length = len(value) 111 | var bit_width = ValueBitWidth.of(byte_length << 2) 112 | var bytes = value._as_ptr().bitcast[DType.uint8]() 113 | var offset = self._offset 114 | 115 | @parameter 116 | if dedup_string and not as_key: 117 | var cached_offset = self._string_cache.get((bytes, byte_length), self._bytes) 118 | if cached_offset != -1: 119 | self._stack.append(StackValue(bitcast[DType.uint8, 8](Int64(cached_offset)), ValueBitWidth.width8, ValueType.String)) 120 | return 121 | 122 | @parameter 123 | if dedup_key and as_key: 124 | var cached_offset = self._key_cache.get((bytes, byte_length), self._bytes) 125 | if cached_offset != -1: 126 | self._stack.append(StackValue(bitcast[DType.uint8, 8](Int64(cached_offset)), bit_width, ValueType.Key)) 127 | return 128 | 129 | @parameter 130 | if not as_key: 131 | var byte_width = self._align(bit_width) 132 | self._write(byte_length << 2 | bit_width.value.to_int(), byte_width) 133 | 134 | var new_offest = self._new_offset(byte_length) 135 | memcpy(self._bytes.offset(self._offset.to_int()), bytes, byte_length) 136 | self._offset = new_offest 137 | self._write(0) 138 | 139 | @parameter 140 | if dedup_string and not as_key: 141 | self._string_cache.put((offset.to_int(), byte_length), self._bytes) 142 | @parameter 143 | if dedup_key and as_key: 144 | self._key_cache.put((offset.to_int(), byte_length), self._bytes) 145 | 146 | @parameter 147 | if as_key: 148 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Key)) 149 | else: 150 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.String)) 151 | value._strref_keepalive() 152 | 153 | fn blob(inout self, value: DTypePointer[DType.uint8], length: Int): 154 | var offset = self._offset 155 | var bit_width = ValueBitWidth.of(length << 2) 156 | var byte_width = self._align(bit_width) 157 | self._write(length << 2 | bit_width.value.to_int(), byte_width) 158 | var new_offest = self._new_offset(length) 159 | memcpy(self._bytes.offset(self._offset.to_int()), value, length) 160 | self._offset = new_offest 161 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Blob)) 162 | 163 | fn add_indirect[D: DType](inout self, value: SIMD[D, 1]): 164 | var value_type = ValueType.of[D]() 165 | if value_type == ValueType.Int or value_type == ValueType.UInt or value_type == ValueType.Float: 166 | var bit_width = ValueBitWidth.of(value) 167 | var byte_width = self._align(bit_width) 168 | var offset = self._offset 169 | self._write(StackValue.of(value), byte_width) 170 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), bit_width, value_type + 5)) 171 | else: 172 | self._stack.append(StackValue.of(value)) 173 | 174 | fn add[D: DType](inout self, value: DTypePointer[D], length: Int): 175 | var len_bit_width = ValueBitWidth.of(length << 2) 176 | var elem_bit_width = ValueBitWidth.of(SIMD[D, 1](0)) 177 | var offset = self._offset 178 | var len_byte_width = self._align(len_bit_width) 179 | self._write((length << 2) | len_bit_width.value.to_int(), len_byte_width) 180 | _ = self._align(elem_bit_width) 181 | var byte_length = sizeof[D]() * length 182 | var new_offest = self._new_offset(byte_length) 183 | memcpy(self._bytes.offset(self._offset.to_int()), value.bitcast[DType.uint8](), byte_length) 184 | self._offset = new_offest 185 | self._stack.append(StackValue(bitcast[DType.uint8, 8](offset), elem_bit_width, ValueType.of[D]() + ValueType.Vector)) 186 | 187 | fn add_referenced(inout self, reference_key: String) raises: 188 | var key = Key(reference_key._as_ptr().bitcast[DType.uint8](), len(reference_key)) 189 | var stack_value = self._reference_cache.get(key, StackValue.Null) 190 | key.pointer.free() 191 | if stack_value.type == ValueType.Null: 192 | raise "No value for reference key " + reference_key 193 | self._stack.append(stack_value) 194 | 195 | fn start_vector(inout self): 196 | self._stack_positions.append(len(self._stack)) 197 | self._stack_is_vector.append(True) 198 | 199 | fn start_map(inout self): 200 | self._stack_positions.append(len(self._stack)) 201 | self._stack_is_vector.append(False) 202 | 203 | fn end(inout self, reference_key: String = "") raises: 204 | var position = self._stack_positions.pop_back() 205 | var is_vector = self._stack_is_vector.pop_back() 206 | if is_vector: 207 | self._end_vector(position) 208 | else: 209 | self._sort_keys_and_end_map(position) 210 | if len(reference_key) > 0: 211 | var key = Key(reference_key._as_ptr().bitcast[DType.uint8](), len(reference_key)) 212 | self._reference_cache.put(key, self._stack[len(self._stack) - 1]) 213 | 214 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 215 | return self._finish() 216 | 217 | fn _finish(inout self) raises -> (DTypePointer[DType.uint8], Int): 218 | self._finished = True 219 | 220 | while len(self._stack_positions) > 0: 221 | self.end() 222 | 223 | if len(self._stack) != 1: 224 | raise "Stack needs to have only one element. Instead of: " + String(len(self._stack)) 225 | 226 | var value = self._stack.pop_back() 227 | var byte_width = self._align(value.element_width(self._offset, 0)) 228 | self._write(value, byte_width) 229 | self._write(value.stored_packed_type()) 230 | self._write(byte_width.cast[DType.uint8]()) 231 | return self._bytes, self._offset.to_int() 232 | 233 | fn _align(inout self, bit_width: ValueBitWidth) -> UInt64: 234 | var byte_width = 1 << bit_width.value.to_int() 235 | self._offset += padding_size(self._offset, byte_width) 236 | return byte_width 237 | 238 | fn _write(inout self, value: StackValue, byte_width: UInt64): 239 | self._grow_bytes_if_needed(self._offset + byte_width) 240 | if value.is_offset(): 241 | var rel_offset = self._offset - value.as_uint() 242 | # Safety check not implemented for now as it is internal call and should be safe 243 | # if byte_width == 8 or rel_offset < (1 << (byte_width * 8)): 244 | self._write(rel_offset, byte_width) 245 | else: 246 | var new_offset = self._new_offset(byte_width) 247 | self._bytes.store(self._offset.to_int(), value.to_value(byte_width)) 248 | self._offset = new_offset 249 | 250 | fn _write(inout self, value: UInt64, byte_width: UInt64): 251 | self._grow_bytes_if_needed(self._offset + byte_width) 252 | var new_offset = self._new_offset(byte_width) 253 | self._bytes.store(self._offset.to_int(), bitcast[DType.uint8, 8](value)) 254 | # We write 8 bytes but the offset is still set to byte_width 255 | self._offset = new_offset 256 | 257 | fn _write(inout self, value: UInt8): 258 | self._grow_bytes_if_needed(self._offset + 1) 259 | var new_offset = self._new_offset(1) 260 | self._bytes.offset(self._offset.to_int()).store(value) 261 | self._offset = new_offset 262 | 263 | fn _new_offset(inout self, byte_width: UInt64) -> UInt64: 264 | var new_offset = self._offset + byte_width 265 | var min_size = self._offset + max(byte_width, 8) 266 | self._grow_bytes_if_needed(min_size) 267 | return new_offset 268 | 269 | fn _grow_bytes_if_needed(inout self, min_size: UInt64): 270 | var prev_size = self._size 271 | while self._size < min_size: 272 | self._size <<= 1 273 | if prev_size < self._size: 274 | var prev_bytes = self._bytes 275 | self._bytes = DTypePointer[DType.uint8].alloc(self._size.to_int()) 276 | memcpy(self._bytes, prev_bytes, self._offset.to_int()) 277 | prev_bytes.free() 278 | 279 | fn _end_vector(inout self, position: Int) raises: 280 | var length = len(self._stack) - position 281 | var vec = self._create_vector(position, length, 1) 282 | self._stack.resize(position, StackValue.Null) 283 | self._stack.append(vec) 284 | 285 | fn _sort_keys_and_end_map(inout self, position: Int) raises: 286 | if (len(self._stack) - position) & 1 == 1: 287 | raise "The stack needs to hold key value pairs (even number of elements). Check if you combined [key] with [add] method calls properly." 288 | for i in range(position + 2, len(self._stack), 2): 289 | var key = self._stack[i] 290 | var value = self._stack[i + 1] 291 | var j = i - 2 292 | while j >= position and self._should_flip(self._stack[j], key): 293 | self._stack[j + 2] = self._stack[j] 294 | self._stack[j + 3] = self._stack[j + 1] 295 | j -= 2 296 | self._stack[j + 2] = key 297 | self._stack[j + 3] = value 298 | self._end_map(position) 299 | 300 | fn _should_flip(self, a: StackValue, b: StackValue) raises -> Bool: 301 | if a.type != ValueType.Key or b.type != ValueType.Key: 302 | raise "Stack values are not keys " + String(a.type.value) + " " + String(a.type.value) 303 | var index = 0 304 | while True: 305 | var c1 = self._bytes.load(a.as_uint().to_int() + index) 306 | var c2 = self._bytes.load(b.as_uint().to_int() + index) 307 | if c1 < c2: 308 | return False 309 | if c1 > c2: 310 | return True 311 | if c1 == 0 and c2 == 0: 312 | return False 313 | index += 1 314 | 315 | fn _end_map(inout self, start: Int) raises: 316 | var length = (len(self._stack) - start) >> 1 317 | var keys = StackValue.Null 318 | @parameter 319 | if dedup_key and dedup_keys_vec: 320 | var keys_vec = self._create_keys_vec_value(start, length) 321 | var cached = self._keys_vec_cache.get(keys_vec, StackValue.Null) 322 | if cached != StackValue.Null: 323 | keys = cached 324 | keys_vec.pointer.free() 325 | else: 326 | keys = self._create_vector(start, length, 2) 327 | self._keys_vec_cache.put(keys_vec, keys) 328 | else: 329 | keys = self._create_vector(start, length, 2) 330 | var map = self._create_vector(start + 1, length, 2, keys) 331 | self._stack.resize(start, StackValue.Null) 332 | self._stack.append(map) 333 | 334 | fn _create_keys_vec_value(self, start: Int, length: Int) -> Key: 335 | var size = length * 8 336 | var result = DTypePointer[DType.uint8].alloc(size) 337 | var offset = 0 338 | memset_zero(result, size) 339 | for i in range(start, len(self._stack), 2): 340 | result.store(offset, self._stack[i].value) 341 | offset += 8 342 | var key = Key(result, size) 343 | result.free() 344 | return key 345 | 346 | fn _create_vector(inout self, start: Int, length: Int, step: Int, keys: StackValue = StackValue.Null) raises -> StackValue: 347 | # var bit_width = ValueBitWidth.of(UInt64(length)) 348 | var bit_width = ValueBitWidth.width8 349 | var prefix_elements = 1 350 | if keys != StackValue.Null: 351 | prefix_elements += 2 352 | var keys_bit_width = keys.element_width(self._offset, 0) 353 | if bit_width < keys_bit_width: 354 | bit_width = keys_bit_width 355 | 356 | var typed = False 357 | var vec_elem_type = ValueType.Null 358 | if length > 0: 359 | vec_elem_type = self._stack[start].type 360 | typed = vec_elem_type.is_typed_vector_element() 361 | if keys != StackValue.Null: 362 | typed = False 363 | for i in range(start, len(self._stack), step): 364 | var elem_bit_width = self._stack[i].element_width(self._offset, i + prefix_elements) 365 | if bit_width < elem_bit_width: 366 | bit_width = elem_bit_width 367 | if vec_elem_type != self._stack[i].type: 368 | typed = False 369 | if bit_width == ValueBitWidth.width64 and typed == False: 370 | break 371 | var byte_width = self._align(bit_width) 372 | if keys != StackValue.Null: 373 | self._write(keys, byte_width) 374 | self._write((1 << keys.width.value).to_int(), byte_width) 375 | var offset = self._offset 376 | var length_bit_width = ValueBitWidth.of(UInt64(length << 2)) 377 | var length_byte_width = self._align(length_bit_width) 378 | self._write(UInt64(length << 2) | length_bit_width.value.cast[DType.uint64](), length_byte_width) 379 | _ = self._align(bit_width) 380 | for i in range(start, len(self._stack), step): 381 | self._write(self._stack[i], byte_width) 382 | if not typed: 383 | for i in range(start, len(self._stack), step): 384 | self._write(self._stack[i].stored_packed_type()) 385 | if keys != StackValue.Null: 386 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Map) 387 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Vector) 388 | 389 | return StackValue(bitcast[DType.uint8, 8](offset), bit_width, ValueType.Vector + vec_elem_type) 390 | 391 | fn finish_ignoring_excetion(owned flx: FlxBuffer) -> (DTypePointer[DType.uint8], Int): 392 | try: 393 | return flx^.finish() 394 | except e: 395 | # should never happen 396 | print("Unexpected error:", e) 397 | return DTypePointer[DType.uint8](), -1 398 | -------------------------------------------------------------------------------- /flx2/flx_builder.mojo: -------------------------------------------------------------------------------- 1 | from .flx_buffer import FlxBuffer 2 | 3 | struct FlxMap[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Movable): 4 | var buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec] 5 | 6 | fn __init__(inout self, owned buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]): 7 | self.buffer = buffer^ 8 | 9 | fn __init__(inout self): 10 | self.buffer = FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]() 11 | self.buffer.start_map() 12 | 13 | fn __moveinit__(inout self, owned other: Self): 14 | self.buffer = other.buffer^ 15 | 16 | fn add(owned self, key: String, value: Int) -> Self: 17 | self.buffer.key(key) 18 | self.buffer.add(value) 19 | return self^ 20 | 21 | fn add[D: DType](owned self, key: String, value: SIMD[D, 1]) -> Self: 22 | self.buffer.key(key) 23 | self.buffer.add(value) 24 | return self^ 25 | 26 | fn add_indirect[D: DType](owned self, key: String, value: SIMD[D, 1]) -> Self: 27 | self.buffer.key(key) 28 | self.buffer.add_indirect(value) 29 | return self^ 30 | 31 | fn add_referenced(owned self, key: String, ref_key: String) raises -> Self: 32 | self.buffer.key(key) 33 | self.buffer.add_referenced(ref_key) 34 | return self^ 35 | 36 | fn add(owned self, key: String, value: String) -> Self: 37 | self.buffer.key(key) 38 | self.buffer.add(value) 39 | return self^ 40 | 41 | fn add(owned self, key: String, value: DTypePointer[DType.uint8], length: Int) -> Self: 42 | self.buffer.key(key) 43 | self.buffer.blob(value, length) 44 | return self^ 45 | 46 | fn map(owned self, key: String) -> Self: 47 | self.buffer.key(key) 48 | self.buffer.start_map() 49 | return self^ 50 | 51 | fn vec(owned self, key: String) -> FlxVec[dedup_string, dedup_key, dedup_keys_vec]: 52 | # TODO: investigate ownership transfer instead of copy 53 | var buffer = self.buffer 54 | buffer.key(key) 55 | buffer.start_vector() 56 | return FlxVec[dedup_string, dedup_key, dedup_keys_vec](buffer^) 57 | 58 | fn up_to_map(owned self, ref_key: String = "") raises -> Self: 59 | var depth = len(self.buffer._stack_is_vector) 60 | if depth < 2: 61 | raise "This map is not nested, please call finish instead" 62 | self.buffer.end(ref_key) 63 | if self.buffer._stack_is_vector[depth - 2]: 64 | raise "This map is nested in a vector, please call up_to_vec instead" 65 | return self^ 66 | 67 | fn up_to_vec(owned self, ref_key: String = "") raises -> FlxVec[dedup_string, dedup_key, dedup_keys_vec]: 68 | # TODO: investigate ownership transfer instead of copy 69 | var buffer = self.buffer 70 | var depth = len(buffer._stack_is_vector) 71 | if depth < 2: 72 | raise "This map is not nested, please call finish instead" 73 | buffer.end(ref_key) 74 | if not buffer._stack_is_vector[depth - 2]: 75 | raise "This map is nested in a map, please call up_to_map instead" 76 | return FlxVec[dedup_string, dedup_key, dedup_keys_vec](buffer^) 77 | 78 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 79 | return self.buffer._finish() 80 | 81 | struct FlxVec[dedup_string: Bool = True, dedup_key: Bool = True, dedup_keys_vec: Bool = True](Movable): 82 | var buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec] 83 | 84 | fn __init__(inout self, owned buffer: FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]): 85 | self.buffer = buffer^ 86 | 87 | fn __init__(inout self): 88 | self.buffer = FlxBuffer[dedup_string, dedup_key, dedup_keys_vec]() 89 | self.buffer.start_vector() 90 | 91 | fn __moveinit__(inout self, owned other: Self): 92 | self.buffer = other.buffer^ 93 | 94 | fn add(owned self, value: Int) -> Self: 95 | self.buffer.add(value) 96 | return self^ 97 | 98 | fn add[D: DType](owned self, value: SIMD[D, 1]) -> Self: 99 | self.buffer.add(value) 100 | return self^ 101 | 102 | fn add_indirect[D: DType](owned self, value: SIMD[D, 1]) -> Self: 103 | self.buffer.add_indirect(value) 104 | return self^ 105 | 106 | fn add_referenced(owned self, ref_key: String) raises -> Self: 107 | self.buffer.add_referenced(ref_key) 108 | return self^ 109 | 110 | fn add(owned self, value: String) -> Self: 111 | self.buffer.add(value) 112 | return self^ 113 | 114 | fn add(owned self, value: DTypePointer[DType.uint8], length: Int) -> Self: 115 | self.buffer.blob(value, length) 116 | return self^ 117 | 118 | fn map(owned self) -> FlxMap[dedup_string, dedup_key, dedup_keys_vec]: 119 | # TODO: investigate ownership transfer instead of copy 120 | var buffer = self.buffer 121 | buffer.start_map() 122 | return FlxMap[dedup_string, dedup_key, dedup_keys_vec](buffer^) 123 | 124 | fn vec(owned self) -> Self: 125 | self.buffer.start_vector() 126 | return self^ 127 | 128 | fn null(owned self) -> Self: 129 | self.buffer.add_null() 130 | return self^ 131 | 132 | fn up_to_map(owned self, ref_key: String = "") raises -> FlxMap[dedup_string, dedup_key, dedup_keys_vec]: 133 | # TODO: investigate ownership transfer instead of copy 134 | var buffer = self.buffer 135 | var depth = len(buffer._stack_is_vector) 136 | if depth < 2: 137 | raise "This vec is not nested, please call finish instead" 138 | buffer.end(ref_key) 139 | if self.buffer._stack_is_vector[depth - 2]: 140 | raise "This vec is nested in a vec, please call up_to_vec instead" 141 | return FlxMap[dedup_string, dedup_key, dedup_keys_vec](buffer^) 142 | 143 | fn up_to_vec(owned self, ref_key: String = "") raises -> Self: 144 | var depth = len(self.buffer._stack_is_vector) 145 | if depth < 2: 146 | raise "This vec is not nested, please call finish instead" 147 | self.buffer.end(ref_key) 148 | if not self.buffer._stack_is_vector[depth - 2]: 149 | raise "This vec is nested in a map, please call up_to_map instead" 150 | return self^ 151 | 152 | fn finish(owned self) raises -> (DTypePointer[DType.uint8], Int): 153 | return self.buffer._finish() 154 | -------------------------------------------------------------------------------- /flx2/flx_value.mojo: -------------------------------------------------------------------------------- 1 | from .data_types import ValueType, ValueBitWidth, is_be 2 | from math.bit import bswap 3 | from math import max 4 | 5 | struct FlxValue(Sized): 6 | var _bytes: DTypePointer[DType.uint8] 7 | var _byte_width: UInt8 8 | var _parent_byte_width: UInt8 9 | var _type: ValueType 10 | 11 | @always_inline 12 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], parent_byte_width: UInt8, packed_type: UInt8): 13 | self._bytes = bytes 14 | self._parent_byte_width = parent_byte_width 15 | self._byte_width = 1 << (packed_type & 3) 16 | self._type = ValueType(packed_type >> 2) 17 | 18 | @always_inline 19 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], parent_byte_width: UInt8, byte_width: UInt8, type: ValueType): 20 | self._bytes = bytes 21 | self._parent_byte_width = parent_byte_width 22 | self._byte_width = byte_width 23 | self._type = type 24 | 25 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], length: Int) raises: 26 | if length < 3: 27 | raise "Length should be at least 3, was: " + String(length) 28 | var parent_byte_width = bytes.load(length - 1) 29 | var packed_type = bytes.load(length - 2) 30 | var offset = length - parent_byte_width.to_int() - 2 31 | self._bytes = bytes.offset(offset) 32 | self._parent_byte_width = parent_byte_width 33 | self._byte_width = 1 << (packed_type & 3) 34 | self._type = ValueType(packed_type >> 2) 35 | 36 | fn __init__(inout self, bytes_and_length: (DTypePointer[DType.uint8], Int)) raises: 37 | var bytes = bytes_and_length.get[0, DTypePointer[DType.uint8]]() 38 | var length = bytes_and_length.get[1, Int]() 39 | if length < 3: 40 | raise "Length should be at least 3, was: " + String(length) 41 | var parent_byte_width = bytes.load(length - 1) 42 | var packed_type = bytes.load(length - 2) 43 | var offset = length - parent_byte_width.to_int() - 2 44 | self._bytes = bytes.offset(offset) 45 | self._parent_byte_width = parent_byte_width 46 | self._byte_width = 1 << (packed_type & 3) 47 | self._type = ValueType(packed_type >> 2) 48 | 49 | fn __moveinit__(inout self, owned other: Self): 50 | self._bytes = other._bytes 51 | self._parent_byte_width = other._parent_byte_width 52 | self._byte_width = other._byte_width 53 | self._type = other._type 54 | 55 | @always_inline 56 | fn __len__(self) -> Int: 57 | if self.is_null(): 58 | return 0 59 | if self.is_vec() or self._type == ValueType.String or self.is_blob() or self.is_map(): 60 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 61 | return read_size(p).get[0, Int]() 62 | if self._type == ValueType.Key: 63 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 64 | var size = 0 65 | while p.offset(size).load() != 0: 66 | size += 1 67 | return 1 68 | 69 | @always_inline 70 | fn is_null(self) -> Bool: 71 | return self._type == ValueType.Null 72 | 73 | @always_inline 74 | fn is_a[D: DType](self) -> Bool: 75 | return self._type == ValueType.of[D]() 76 | 77 | @always_inline 78 | fn is_map(self) -> Bool: 79 | return self._type == ValueType.Map 80 | 81 | @always_inline 82 | fn is_vec(self) -> Bool: 83 | return self._type.is_a_vector() 84 | 85 | @always_inline 86 | fn is_string(self) -> Bool: 87 | return self._type == ValueType.String or self._type == ValueType.Key 88 | 89 | @always_inline 90 | fn is_blob(self) -> Bool: 91 | return self._type == ValueType.Blob 92 | 93 | @always_inline 94 | fn is_int(self) -> Bool: 95 | return self._type == ValueType.Int 96 | or self._type == ValueType.UInt 97 | or self._type == ValueType.IndirectInt 98 | or self._type == ValueType.IndirectUInt 99 | 100 | @always_inline 101 | fn is_float(self) -> Bool: 102 | return self._type == ValueType.Float 103 | or self._type == ValueType.IndirectFloat 104 | 105 | @always_inline 106 | fn is_bool(self) -> Bool: 107 | return self._type == ValueType.Bool 108 | 109 | @always_inline 110 | fn get[D: DType](self) raises -> SIMD[D, 1]: 111 | if self._type != ValueType.of[D](): 112 | raise "Value is not of type " + D.__str__() + " type id: " + String(self._type.value) 113 | if sizeof[D]() != self._parent_byte_width.to_int(): 114 | raise "Value byte width is " + String(self._parent_byte_width.to_int()) + " which does not conform with " + D.__str__() 115 | @parameter 116 | if is_be: 117 | return bswap(self._bytes.bitcast[D]().load()) 118 | else: 119 | return self._bytes.bitcast[D]().load() 120 | 121 | @always_inline 122 | fn int(self) raises -> Int: 123 | if self._type == ValueType.Int: 124 | return read_int(self._bytes, self._parent_byte_width) 125 | if self._type == ValueType.UInt: 126 | return read_uint(self._bytes, self._parent_byte_width) 127 | if self._type == ValueType.IndirectInt: 128 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 129 | return read_int(p, self._byte_width) 130 | if self._type == ValueType.IndirectUInt: 131 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 132 | return read_uint(p, self._byte_width) 133 | raise "Type is not an int or uint, type id: " + String(self._type.value) 134 | 135 | @always_inline 136 | fn float(self) raises -> Float64: 137 | if self._type == ValueType.Float: 138 | return read_float(self._bytes, self._parent_byte_width) 139 | if self._type == ValueType.IndirectFloat: 140 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 141 | return read_float(p, self._byte_width) 142 | raise "Type is not a float, type id: " + String(self._type.value) 143 | 144 | @always_inline 145 | fn bool(self) raises -> Bool: 146 | if self._type == ValueType.Bool: 147 | return read_uint(self._bytes, self._parent_byte_width) == 1 148 | raise "Type is not a bool, type id: " + String(self._type.value) 149 | 150 | @always_inline 151 | fn string(self) raises -> StringRef: 152 | if self._type == ValueType.String: 153 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 154 | var size: Int 155 | var size_width: Int 156 | size, size_width = read_size(p) 157 | return StringRef(p.offset(size_width).bitcast[DType.int8](), size) 158 | if self._type == ValueType.Key: 159 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 160 | var size = 0 161 | while p.offset(size).load() != 0: 162 | size += 1 163 | return StringRef(p.bitcast[DType.int8]()) 164 | raise "Type is not convertable to string, type id: " + String(self._type.value) 165 | 166 | @always_inline 167 | fn blob(self) raises -> (DTypePointer[DType.uint8], Int): 168 | if not self.is_blob(): 169 | raise "Type is not blob, type id: " + String(self._type.value) 170 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 171 | var size: Int 172 | var size_width: Int 173 | size, size_width = read_size(p) 174 | return (p.offset(size_width), size) 175 | 176 | @always_inline 177 | fn vec(self) raises -> FlxVecValue: 178 | if not self._type.is_a_vector(): 179 | raise "Value is not a vector. Type id: " + String(self._type.value) 180 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 181 | var size: Int 182 | var size_width: Int 183 | size, size_width = read_size(p) 184 | return FlxVecValue(p.offset(max(size_width, self._byte_width)), self._byte_width, self._type, size) 185 | 186 | @always_inline 187 | fn map(self) raises -> FlxMapValue: 188 | if self._type != ValueType.Map: 189 | raise "Value is not a map. Type id: " + String(self._type.value) 190 | var p = jump_to_indirect(self._bytes, self._parent_byte_width) 191 | var size: Int 192 | var size_width: Int 193 | size, size_width = read_size(p) 194 | return FlxMapValue(p.offset(max(size_width, self._byte_width)), self._byte_width, size) 195 | 196 | @always_inline 197 | fn has_key(self, key: String) raises -> Bool: 198 | if not self.is_map(): 199 | return False 200 | return self.map().key_index(key) >= 0 201 | 202 | @always_inline 203 | fn __getitem__(self, index: Int) raises -> FlxValue: 204 | return self.vec()[index] 205 | 206 | @always_inline 207 | fn __getitem__(self, key: String) raises -> FlxValue: 208 | return self.map()[key] 209 | 210 | fn json(self) raises -> String: 211 | if self.is_null(): 212 | return "null" 213 | if self.is_bool(): 214 | return "true" if self.bool() else "false" 215 | if self.is_int(): 216 | return self.int() 217 | if self.is_float(): 218 | return self.float() 219 | if self.is_string(): 220 | return '"' + str(self.string()) + '"' 221 | if self.is_vec(): 222 | return self.vec().json() 223 | if self.is_map(): 224 | return self.map().json() 225 | raise "Unexpected type id: " + String(self._type.value) 226 | 227 | 228 | struct FlxVecValue(Sized): 229 | var _bytes: DTypePointer[DType.uint8] 230 | var _byte_width: UInt8 231 | var _type: ValueType 232 | var _length: Int 233 | 234 | @always_inline 235 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], byte_width: UInt8, type: ValueType, length: Int): 236 | self._bytes = bytes 237 | self._byte_width = byte_width 238 | self._type = type 239 | self._length = length 240 | 241 | @always_inline 242 | fn __len__(self) -> Int: 243 | return self._length 244 | 245 | @always_inline 246 | fn __getitem__(self, index: Int) raises -> FlxValue: 247 | if index < 0 or index >= self._length: 248 | raise "Bad index: " + String(index) + ". Lenght: " + String(self._length) 249 | if self._type.is_typed_vector(): 250 | return FlxValue( 251 | self._bytes.offset(index * self._byte_width.to_int()), 252 | self._byte_width, 253 | 1, 254 | self._type.typed_vector_element_type() 255 | ) 256 | if self._type.is_fixed_typed_vector(): 257 | return FlxValue( 258 | self._bytes.offset(index * self._byte_width.to_int()), 259 | self._byte_width, 260 | 1, 261 | self._type.fixed_typed_vector_element_type() 262 | ) 263 | if self._type == ValueType.Vector: 264 | var packed_type = self._bytes.offset(self._length * self._byte_width.to_int() + index).load() 265 | return FlxValue( 266 | self._bytes.offset(index * self._byte_width.to_int()), 267 | self._byte_width, 268 | packed_type 269 | ) 270 | raise "Is not an expected vector type. Type id: " + String(self._type.value) 271 | 272 | fn json(self) raises -> String: 273 | var result: String = "[" 274 | for i in range(self.__len__()): 275 | result += self[i].json() 276 | if i < self.__len__() - 1: 277 | result += "," 278 | result += "]" 279 | return result 280 | 281 | struct FlxMapValue(Sized): 282 | var _bytes: DTypePointer[DType.uint8] 283 | var _byte_width: UInt8 284 | var _length: Int 285 | 286 | @always_inline 287 | fn __init__(inout self, bytes: DTypePointer[DType.uint8], byte_width: UInt8, length: Int): 288 | self._bytes = bytes 289 | self._byte_width = byte_width 290 | self._length = length 291 | 292 | @always_inline 293 | fn __len__(self) -> Int: 294 | return self._length 295 | 296 | @always_inline 297 | fn __getitem__(self, key: String) raises -> FlxValue: 298 | var index = self.key_index(key) 299 | if index < 0: 300 | raise "Key " + key + " could not be found" 301 | return self.values()[index] 302 | 303 | @always_inline 304 | fn keys(self) -> FlxVecValue: 305 | var p1 = self._bytes.offset(self._byte_width.to_int() * -3) 306 | var p2 = jump_to_indirect(p1, self._byte_width) 307 | var byte_width = read_uint(p1.offset(self._byte_width.to_int()), self._byte_width) 308 | return FlxVecValue(p2.offset(max(byte_width, 1 << ValueBitWidth.of(self._length).value.to_int())), byte_width, ValueType.VectorKey, self._length) 309 | 310 | @always_inline 311 | fn values(self) -> FlxVecValue: 312 | return FlxVecValue(self._bytes, self._byte_width, ValueType.Vector, self._length) 313 | 314 | @always_inline 315 | fn key_index(self, key: String) raises -> Int: 316 | var a = key._as_ptr().bitcast[DType.uint8]() 317 | var keys = self.keys() 318 | var low = 0 319 | var high = self._length - 1 320 | while low <= high: 321 | var mid = (low + high) >> 1 322 | var mid_key = keys[mid] 323 | var b = jump_to_indirect(mid_key._bytes, mid_key._parent_byte_width) 324 | var diff = cmp(a, b, len(key) + 1) 325 | if diff == 0: 326 | return mid 327 | if diff < 0: 328 | high = mid - 1 329 | else: 330 | low = mid + 1 331 | return -1 332 | 333 | fn json(self) raises -> String: 334 | var result: String = "{" 335 | var keys = self.keys() 336 | var values = self.values() 337 | for i in range(self.__len__()): 338 | result += '"' + str(keys[i].string()) + '":' + values[i].json() 339 | if i < self.__len__() - 1: 340 | result += "," 341 | result += "}" 342 | return result 343 | 344 | 345 | @always_inline 346 | fn jump_to_indirect(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> DTypePointer[DType.uint8]: 347 | return bytes.offset(-read_uint(bytes, byte_width)) 348 | 349 | @always_inline 350 | fn read_size(bytes: DTypePointer[DType.uint8]) -> Tuple[Int, Int]: 351 | var val = bytes.bitcast[DType.uint64]().load() 352 | var byte_width = (val & 3).to_int() + 1 353 | var mask = (1 << (byte_width << 3)) - 1 354 | return ((val & mask) >> 2).to_int(), byte_width 355 | # Below is a safer version, do not delete yet: 356 | # var val = bytes.load() 357 | # var byte_width = (val & 3).to_int() + 1 358 | # if byte_width > 2: 359 | # if byte_width == 4: 360 | # return (bytes.bitcast[DType.uint32]().load() >> 2).to_int(), byte_width 361 | # return (bytes.bitcast[DType.uint64]().load() >> 2).to_int(), byte_width 362 | # else: 363 | # if byte_width == 2: 364 | # return (bytes.bitcast[DType.uint16]().load() >> 2).to_int(), byte_width 365 | # return (val >> 2).to_int(), byte_width 366 | 367 | @always_inline 368 | fn read_uint(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> Int: 369 | if byte_width < 4: 370 | if byte_width == 1: 371 | return bytes.load().to_int() 372 | else: 373 | @parameter 374 | if is_be: 375 | return bswap(bytes.bitcast[DType.uint16]().load()).to_int() 376 | else: 377 | return bytes.bitcast[DType.uint16]().load().to_int() 378 | else: 379 | if byte_width == 4: 380 | @parameter 381 | if is_be: 382 | return bswap(bytes.bitcast[DType.uint32]().load()).to_int() 383 | else: 384 | return bytes.bitcast[DType.uint32]().load().to_int() 385 | else: 386 | @parameter 387 | if is_be: 388 | return bswap(bytes.bitcast[DType.uint64]().load()).to_int() 389 | else: 390 | return bytes.bitcast[DType.uint64]().load().to_int() 391 | 392 | @always_inline 393 | fn read_int(bytes: DTypePointer[DType.uint8], byte_width: UInt8) -> Int: 394 | if byte_width < 4: 395 | if byte_width == 1: 396 | return bytes.bitcast[DType.int8]().load().to_int() 397 | else: 398 | @parameter 399 | if is_be: 400 | return bswap(bytes.bitcast[DType.int16]().load()).to_int() 401 | else: 402 | return bytes.bitcast[DType.int16]().load().to_int() 403 | else: 404 | if byte_width == 4: 405 | @parameter 406 | if is_be: 407 | return bswap(bytes.bitcast[DType.int32]().load()).to_int() 408 | else: 409 | return bytes.bitcast[DType.int32]().load().to_int() 410 | else: 411 | @parameter 412 | if is_be: 413 | return bswap(bytes.bitcast[DType.int64]().load()).to_int() 414 | else: 415 | return bytes.bitcast[DType.int64]().load().to_int() 416 | 417 | @always_inline 418 | fn read_float(bytes: DTypePointer[DType.uint8], byte_width: UInt8) raises -> Float64: 419 | if byte_width == 8: 420 | @parameter 421 | if is_be: 422 | return bswap(bytes.bitcast[DType.float64]().load()) 423 | else: 424 | return bytes.bitcast[DType.float64]().load() 425 | if byte_width == 4: 426 | @parameter 427 | if is_be: 428 | return bswap(bytes.bitcast[DType.float32]().load()).cast[DType.float64]() 429 | else: 430 | return bytes.bitcast[DType.float32]().load().cast[DType.float64]() 431 | if byte_width == 2: 432 | @parameter 433 | if is_be: 434 | return bswap(bytes.bitcast[DType.float16]().load()).cast[DType.float64]() 435 | else: 436 | return bytes.bitcast[DType.float16]().load().cast[DType.float64]() 437 | raise "Unexpected byte width: " + String(byte_width) 438 | 439 | @always_inline 440 | fn cmp(a: DTypePointer[DType.uint8], b: DTypePointer[DType.uint8], length: Int) -> Int: 441 | for i in range(length): 442 | var diff = a.load(i).to_int() - b.load(i).to_int() 443 | if diff != 0: 444 | return diff 445 | return 0 446 | -------------------------------------------------------------------------------- /samples.mojo: -------------------------------------------------------------------------------- 1 | from test_builder import print_result 2 | from flx import * 3 | 4 | fn main(): 5 | print_result(flx_null()) 6 | print_result(flx(200)) 7 | print_result(flx[DType.float16](2.5)) 8 | print_result(flx[DType.float32](2.5)) 9 | print_result(flx[DType.float64](2.5)) 10 | print_result(flx("Hello 🔥")) 11 | try: 12 | var flxb = flx_buffer.FlxBuffer() 13 | flxb.key("Hello 🔥") 14 | print_result(flxb^.finish()) 15 | print_result(FlxVec().add(5).add(6).add(7).finish()) 16 | print_result(FlxVec().add[DType.uint8](5).add[DType.uint16](600).add[DType.uint8](7).finish()) 17 | print_result(FlxVec().add[DType.float16](1.1).add[DType.float32](1.1).add[DType.float64](1.1).finish()) 18 | print_result(FlxVec().add_indirect[DType.float16](1.1).add_indirect[DType.float32](1.1).add_indirect[DType.float64](1.1).finish()) 19 | print_result(FlxVec().add("maxim").add("alex").add("daria").finish()) 20 | print_result(FlxVec().add("maxim").add("alex").add("maxim").add("daria").finish()) 21 | print_result(FlxVec[dedup_string=False]().add("maxim").add("alex").add("maxim").add("daria").finish()) 22 | print_result(FlxVec().add[DType.int32](1234).add("maxim").add[DType.float16](1.5).add[DType.bool](True).finish()) 23 | print_result(FlxVec().add_indirect[DType.int32](1234).add("maxim").add_indirect[DType.float16](1.5).add[DType.bool](True).finish()) 24 | print_result(FlxVec().add(7).vec().add(8).add(9).finish()) 25 | print_result(FlxMap().add("a", 7).add("b", 8).finish()) 26 | print_result(FlxMap().add("b", 7).add("a", 8).finish()) 27 | print_result(FlxVec().map().add("a", 7).add("b", 8).up_to_vec().map().add("b", 42).add("a", 43).finish()) 28 | print_result(FlxVec[dedup_keys_vec=False]().map().add("a", 7).add("b", 8).up_to_vec().map().add("b", 42).add("a", 43).finish()) 29 | print_result(FlxVec[dedup_key=False]().map().add("a", 7).add("b", 8).up_to_vec().map().add("b", 42).add("a", 43).finish()) 30 | except e: 31 | print(e) 32 | 33 | 34 | -------------------------------------------------------------------------------- /test_builder.mojo: -------------------------------------------------------------------------------- 1 | 2 | from math.bit import bit_length 3 | from flx import FlxMap, FlxVec, flx, flx_blob, flx_null 4 | from flx.flx_buffer import FlxBuffer 5 | from memory.unsafe import bitcast 6 | from testing import assert_equal 7 | 8 | fn print_result(r: Tuple[DTypePointer[DType.uint8], Int]): 9 | var bytes = r.get[0, DTypePointer[DType.uint8]]() 10 | var length = r.get[1, Int]() 11 | print("(" + str(length) + ")[", end="") 12 | for i in range(length): 13 | var end = ", " if i < length - 1 else "]\n" 14 | print(String(bytes.load(i)), end=end) 15 | 16 | 17 | fn assert_result(r: Tuple[DTypePointer[DType.uint8], Int], *bytes: UInt8): 18 | if len(bytes) != r.get[1, Int](): 19 | print("Error, result contains", r.get[1, Int](), "bytes and you provided", len(bytes)) 20 | print_result(r) 21 | return 22 | var p = r.get[0, DTypePointer[DType.uint8]]() 23 | for i in range(len(bytes)): 24 | if p.load(i) != bytes[i]: 25 | print("Error at index", i, "byte", p.load(i), "!=", bytes[i]) 26 | print_result(r) 27 | return 28 | 29 | fn vec[D: DType](*values: SIMD[D, 1]) -> List[SIMD[D, 1]]: 30 | var result = List[SIMD[D, 1]](capacity=len(values)) 31 | for i in range(len(values)): 32 | result.append(values[i]) 33 | return result 34 | 35 | fn test_single_value_contructor(): 36 | assert_result(flx_null(), 0, 0, 1) 37 | assert_result(flx(25), 25, 4, 1) 38 | assert_result(flx(-25), 231, 4, 1) 39 | assert_result(flx(230), 230, 0, 5, 2) 40 | assert_result(flx(-230), 26, 255, 5, 2) 41 | assert_result(flx[DType.uint8](230), 230, 8, 1) 42 | assert_result(flx[DType.uint16](230), 230, 0, 9, 2) 43 | assert_result(flx[DType.uint32](230), 230, 0, 0, 0, 10, 4) 44 | assert_result(flx[DType.uint64](230), 230, 0, 0, 0, 0, 0, 0, 0, 11, 8) 45 | assert_result(flx[DType.bool](True), 1, 104, 1) 46 | assert_result(flx[DType.bool](False), 0, 104, 1) 47 | assert_result(flx[DType.float16](4.5), 128, 68, 13, 2) 48 | assert_result(flx[DType.float32](4.5), 0, 0, 144, 64, 14, 4) 49 | assert_result(flx[DType.float32](0.1), 205, 204, 204, 61, 14, 4) 50 | assert_result(flx("Maxim"), 5, 77, 97, 120, 105, 109, 0, 6, 20, 1) 51 | assert_result(flx("hello 😱"), 10, 104, 101, 108, 108, 111, 32, 240, 159, 152, 177, 0, 11, 20, 1) 52 | assert_result(flx("hello 🔥"), 10, 104, 101, 108, 108, 111, 32, 240, 159, 148, 165, 0, 11, 20, 1) 53 | var v1 = vec[DType.int8](1, 2, 3) 54 | assert_result(flx(DTypePointer[DType.int8](v1.data.value), len(v1)), 3, 1, 2, 3, 3, 44, 1) 55 | _= v1 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 56 | v1 = vec[DType.int8](-1, 2, 3) 57 | assert_result(flx(DTypePointer[DType.int8](v1.data.value), len(v1)), 3, 255, 2, 3, 3, 44, 1) 58 | _= v1 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 59 | var v2 = vec[DType.int16](1, 555, 3) 60 | assert_result(flx(DTypePointer[DType.int16](v2.data.value), len(v2)), 3, 0, 1, 0, 43, 2, 3, 0, 6, 45, 1) 61 | _= v2 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 62 | var v4 = vec[DType.int32](1, 55500, 3) 63 | assert_result( 64 | flx(DTypePointer[DType.int32](v4.data.value), len(v4)), 65 | 3, 0, 0, 0, 1, 0, 0, 0, 204, 216, 0, 0, 3, 0, 0, 0, 12, 46, 1 66 | ) 67 | _= v4 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 68 | var v8 = vec[DType.int64](1, 55555555500, 3) 69 | assert_result( 70 | flx(DTypePointer[DType.int64](v8.data.value), len(v8)), 71 | 3, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 172, 128, 94, 239, 12, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 24, 47, 1 72 | ) 73 | _= v8 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 74 | var vb = vec[DType.bool](True, False, True) 75 | assert_result(flx(DTypePointer[DType.bool](vb.data.value), len(vb)), 3, 1, 0, 1, 3, 144, 1) 76 | _= vb # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 77 | 78 | fn test_vec_construction(): 79 | try: 80 | var flx1 = FlxBuffer() 81 | flx1.start_vector() 82 | flx1.add(1) 83 | flx1.add(2) 84 | flx1.add(3) 85 | flx1.end() 86 | assert_result(flx1^.finish(), 3, 1, 2, 3, 3, 44, 1) 87 | 88 | var flx2 = FlxBuffer() 89 | flx2.start_vector() 90 | flx2.add(1) 91 | flx2.add(555) 92 | flx2.add(3) 93 | flx2.end() 94 | assert_result(flx2^.finish(), 3, 0, 1, 0, 43, 2, 3, 0, 6, 45, 1) 95 | 96 | var flx3 = FlxBuffer() 97 | flx3.start_vector() 98 | flx3.add("foo") 99 | flx3.add("bar") 100 | flx3.add("baz") 101 | flx3.end() 102 | assert_result(flx3^.finish(), 3, 102, 111, 111, 0, 3, 98, 97, 114, 0, 3, 98, 97, 122, 0, 3, 15, 11, 7, 3, 60, 1) 103 | 104 | var flx4 = FlxBuffer() 105 | flx4.start_vector() 106 | flx4.add("foo") 107 | flx4.add(1) 108 | flx4.add(-5) 109 | flx4.add[DType.float64](1.3) 110 | flx4.add[DType.bool](True) 111 | flx4.end() 112 | assert_result( 113 | flx4^.finish(), 114 | 3, 102, 111, 111, 0, 0, 0, 0, 115 | 5, 0, 0, 0, 0, 0, 0, 0, 116 | 15, 0, 0, 0, 0, 0, 0, 0, 117 | 1, 0, 0, 0, 0, 0, 0, 0, 118 | 251, 255, 255, 255, 255, 255, 255, 255, 119 | 205, 204, 204, 204, 204, 204, 244, 63, 120 | 1, 0, 0, 0, 0, 0, 0, 0, 121 | 20, 4, 4, 15, 104, 45, 43, 1 122 | ) 123 | 124 | var flx5 = FlxBuffer() 125 | flx5.start_vector() 126 | flx5.key("foo") 127 | flx5.key("bar") 128 | flx5.key("baz") 129 | flx5.end() 130 | assert_result(flx5^.finish(), 102, 111, 111, 0, 98, 97, 114, 0, 98, 97, 122, 0, 3, 13, 10, 7, 3, 56, 1) 131 | 132 | var flx6 = FlxBuffer() 133 | flx6.start_vector() 134 | flx6.start_vector() 135 | flx6.add(61) 136 | flx6.end() 137 | flx6.add(64) 138 | flx6.end() 139 | assert_result(flx6^.finish(), 1, 61, 2, 2, 64, 44, 4, 4, 40, 1) 140 | 141 | except: 142 | print("unexpected error") 143 | 144 | fn test_map_construction(): 145 | try: 146 | var flx1 = FlxBuffer() 147 | flx1.start_map() 148 | flx1.key("a") 149 | flx1.add(12) 150 | flx1.end() 151 | assert_result(flx1^.finish(), 97, 0, 1, 3, 1, 1, 1, 12, 4, 2, 36, 1) 152 | 153 | var flx2 = FlxBuffer() 154 | flx2.start_map() 155 | flx2.key("") 156 | flx2.add(45) 157 | flx2.key("a") 158 | flx2.add(12) 159 | flx2.end() 160 | assert_result(flx2^.finish(), 0, 97, 0, 2, 4, 4, 2, 1, 2, 45, 12, 4, 4, 4, 36, 1) 161 | 162 | var flx3 = FlxBuffer() 163 | flx3.start_map() 164 | flx3.key("b") 165 | flx3.add(45) 166 | flx3.key("a") 167 | flx3.add(12) 168 | flx3.end() 169 | assert_result( 170 | flx3^.finish(), 171 | 98, 0, 97, 0, 2, 3, 6, 2, 1, 2, 12, 45, 4, 4, 4, 36, 1 172 | ) 173 | except: 174 | print("unexpected error") 175 | 176 | fn test_add_referenced() raises: 177 | var flx = FlxBuffer() 178 | flx.start_vector() 179 | flx.start_vector() 180 | flx.add(1) 181 | flx.add(2) 182 | flx.add(3) 183 | flx.end("v1") 184 | flx.add_referenced("v1") 185 | flx.add_referenced("v1") 186 | flx.add_referenced("v1") 187 | flx.end() 188 | assert_result( 189 | flx^.finish(), 190 | 3, 1, 2, 3, 4, 4, 5, 6, 7, 44, 44, 44, 44, 8, 40, 1, 191 | ) 192 | 193 | fn test_map_builder(): 194 | try: 195 | assert_result( 196 | FlxMap().add("", 45).add("a", 12).finish(), 197 | 0, 97, 0, 2, 4, 4, 2, 1, 2, 45, 12, 4, 4, 4, 36, 1 198 | ) 199 | 200 | assert_result( 201 | FlxMap() 202 | .add("a", 12) 203 | .map("b") 204 | .add("c", 33) 205 | .up_to_map() 206 | .add("d", "max") 207 | .finish(), 208 | 0x61, 0x00, 0x62, 0x00, 0x63, 0x00, 0x01, 0x03, 0x01, 0x01, 0x01, 0x21, 0x04, 0x64, 0x00, 0x03, 209 | 0x6d, 0x61, 0x78, 0x00, 0x03, 0x15, 0x14, 0x0a, 0x03, 0x01, 0x03, 0x0c, 0x11, 0x0d, 0x04, 0x24, 210 | 0x14, 0x06, 0x24, 0x01 211 | ) 212 | 213 | assert_result( 214 | FlxMap() 215 | .add("a", 12) 216 | .vec("b") 217 | .add(33) 218 | .up_to_map() 219 | .add("d", "max") 220 | .finish(), 221 | 0x61, 0x00, 0x62, 0x00, 0x01, 0x21, 0x64, 0x00, 0x03, 0x6d, 0x61, 0x78, 0x00, 0x03, 0x0e, 0x0d, 222 | 0x0a, 0x03, 0x01, 0x03, 0x0c, 0x10, 0x0d, 0x04, 0x2c, 0x14, 0x06, 0x24, 0x01 223 | ) 224 | 225 | assert_result( 226 | FlxMap().add("a", 1).add_indirect[DType.int32]("b", 2333).add("c", 3).finish(), 227 | 97, 0, 98, 0, 29, 9, 0, 0, 99, 0, 3, 11, 10, 5, 3, 1, 3, 1, 14, 3, 4, 26, 4, 6, 36, 1 228 | ) 229 | except e: 230 | print("unexpected error", e) 231 | 232 | fn test_vec_builder(): 233 | try: 234 | assert_result( 235 | FlxVec().vec().add(61).up_to_vec().add(64).finish(), 236 | 1, 61, 2, 2, 64, 44, 4, 4, 40, 1 237 | ) 238 | assert_result( 239 | FlxVec().add(1).add_indirect[DType.int32](2333).add(3).finish(), 240 | 29, 9, 0, 0, 3, 1, 6, 3, 4, 26, 4, 6, 40, 1 241 | ) 242 | var vec = FlxVec() 243 | for i in range(256): 244 | vec = vec^.add[DType.bool](i & 1 == 1) 245 | assert_result( 246 | vec^.finish(), 247 | 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 248 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 249 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 250 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 251 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 252 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 253 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 254 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 255 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 256 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 257 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 258 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 259 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 260 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 261 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 262 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 263 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 264 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 265 | 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 266 | 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 2, 145, 2, 267 | ) 268 | except e: 269 | print("unexpected error", e) 270 | 271 | fn test_blob() raises: 272 | var data = DTypePointer[DType.uint8].alloc(1001) 273 | for i in range(1001): 274 | data[i] = 5 275 | 276 | var r = flx_blob(data, 1001) 277 | var b = r.get[0, DTypePointer[DType.uint8]]() 278 | var l = r.get[1, Int]() 279 | _ = assert_equal(l, 1008) 280 | _ = assert_equal(b.load(0), 233) 281 | _ = assert_equal(b.load(1), 3) 282 | for i in range(2, 1002): 283 | _ = assert_equal(b.load(i), 5) 284 | # Padding bytes are not zeroed out for now 285 | # _ = assert_equal(b.load(1003), 0) 286 | _ = assert_equal(b.load(1004), 234) 287 | _ = assert_equal(b.load(1005), 3) 288 | _ = assert_equal(b.load(1006), 101) 289 | _ = assert_equal(b.load(1007), 2) 290 | 291 | fn test_dedup_string(): 292 | try: 293 | assert_result( 294 | FlxVec().add("maxim").add("alex").add("maxim").finish(), 295 | 5, 109, 97, 120, 105, 109, 0, 4, 97, 108, 101, 120, 0, 3, 13, 7, 15, 3, 60, 1 296 | ) 297 | assert_result( 298 | FlxVec[dedup_string=False]().add("maxim").add("alex").add("maxim").finish(), 299 | 5, 109, 97, 120, 105, 109, 0, 4, 97, 108, 101, 120, 0, 300 | 5, 109, 97, 120, 105, 109, 0, 3, 20, 14, 9, 3, 60, 1 301 | ) 302 | assert_result( 303 | FlxMap().add("a", "maxim").add("b", "alex").add("c", "maxim").finish(), 304 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 99, 0, 305 | 3, 20, 12, 5, 3, 1, 3, 23, 15, 25, 20, 20, 20, 6, 36, 1 306 | ) 307 | assert_result( 308 | FlxMap[dedup_string=False]().add("a", "maxim").add("b", "alex").add("c", "maxim").finish(), 309 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 99, 0, 310 | 5, 109, 97, 120, 105, 109, 0, 3, 27, 19, 12, 3, 1, 3, 30, 22, 15, 20, 20, 20, 6, 36, 1 311 | ) 312 | except e: 313 | print("unexpected error", e) 314 | 315 | fn test_dedup_key(): 316 | try: 317 | assert_result( 318 | FlxVec[dedup_key=False]() 319 | .map() 320 | .add("a", "maxim") 321 | .add("b", "alex") 322 | .up_to_vec() 323 | .map() 324 | .add("a", "lena") 325 | .add("c", "daria") 326 | .finish(), 327 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 328 | 2, 18, 10, 2, 1, 2, 20, 12, 20, 20, 97, 0, 4, 108, 101, 110, 97, 0, 329 | 99, 0, 5, 100, 97, 114, 105, 97, 0, 2, 18, 11, 2, 1, 2, 20, 13, 20, 330 | 20, 2, 32, 6, 36, 36, 4, 40, 1, 331 | ) 332 | assert_result( 333 | FlxVec() 334 | .map() 335 | .add("a", "maxim") 336 | .add("b", "alex") 337 | .up_to_vec() 338 | .map() 339 | .add("a", "lena") 340 | .add("c", "daria") 341 | .finish(), 342 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 343 | 2, 18, 10, 2, 1, 2, 20, 12, 20, 20, 4, 108, 101, 110, 97, 0, 99, 0, 344 | 5, 100, 97, 114, 105, 97, 0, 2, 43, 11, 2, 1, 2, 20, 13, 20, 345 | 20, 2, 30, 6, 36, 36, 4, 40, 1 346 | ) 347 | assert_result( 348 | FlxVec() 349 | .map() 350 | .add("a", "maxim") 351 | .add("b", "alex") 352 | .up_to_vec() 353 | .map() 354 | .add("a", "lena") 355 | .add("b", "daria") 356 | .finish(), 357 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 358 | 2, 18, 10, 2, 1, 2, 20, 12, 20, 20, 4, 108, 101, 110, 97, 0, 359 | 5, 100, 97, 114, 105, 97, 0, 22, 1, 2, 15, 10, 20, 360 | 20, 2, 25, 6, 36, 36, 4, 40, 1, 361 | ) 362 | assert_result( 363 | FlxVec[dedup_keys_vec=False]() 364 | .map() 365 | .add("a", "maxim") 366 | .add("b", "alex") 367 | .up_to_vec() 368 | .map() 369 | .add("a", "lena") 370 | .add("b", "daria") 371 | .finish(), 372 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 373 | 2, 18, 10, 2, 1, 2, 20, 12, 20, 20, 4, 108, 101, 110, 97, 0, 374 | 5, 100, 97, 114, 105, 97, 0, 2, 41, 33, 2, 1, 2, 18, 13, 20, 375 | 20, 2, 28, 6, 36, 36, 4, 40, 1, 376 | ) 377 | assert_result( 378 | FlxVec[dedup_key=False]() 379 | .map() 380 | .add("a", "maxim") 381 | .add("b", "alex") 382 | .up_to_vec() 383 | .map() 384 | .add("a", "lena") 385 | .add("b", "daria") 386 | .finish(), 387 | 97, 0, 5, 109, 97, 120, 105, 109, 0, 98, 0, 4, 97, 108, 101, 120, 0, 388 | 2, 18, 10, 2, 1, 2, 20, 12, 20, 20, 97, 0, 4, 108, 101, 110, 97, 0, 98, 0, 389 | 5, 100, 97, 114, 105, 97, 0, 2, 18, 11, 2, 1, 2, 20, 13, 20, 390 | 20, 2, 32, 6, 36, 36, 4, 40, 1, 391 | ) 392 | except e: 393 | print("unexpected error", e) 394 | 395 | fn main() raises: 396 | test_single_value_contructor() 397 | test_vec_construction() 398 | test_map_construction() 399 | test_map_builder() 400 | test_vec_builder() 401 | test_blob() 402 | test_dedup_string() 403 | test_dedup_key() 404 | test_add_referenced() 405 | 406 | print("All Done!!!") -------------------------------------------------------------------------------- /test_builder2.mojo: -------------------------------------------------------------------------------- 1 | from math.bit import bit_length 2 | from flx2 import FlxMap, FlxVec, flx, flx_blob, flx_null 3 | from flx2.flx_buffer import FlxBuffer 4 | from memory.unsafe import bitcast 5 | from testing import assert_equal 6 | 7 | fn print_result(r: Tuple[DTypePointer[DType.uint8], Int]): 8 | var bytes = r.get[0, DTypePointer[DType.uint8]]() 9 | var length = r.get[1, Int]() 10 | print("(" + str(length) + ")[", end="") 11 | for i in range(length): 12 | var end = ", " if i < length - 1 else "]\n" 13 | print(String(bytes.load(i)), end=end) 14 | 15 | 16 | fn assert_result(r: Tuple[DTypePointer[DType.uint8], Int], *bytes: UInt8): 17 | if len(bytes) != r.get[1, Int](): 18 | print("Error, result contains", r.get[1, Int](), "bytes and you provided", len(bytes)) 19 | print_result(r) 20 | return 21 | var p = r.get[0, DTypePointer[DType.uint8]]() 22 | for i in range(len(bytes)): 23 | if p.load(i) != bytes[i]: 24 | print("Error at index", i, "byte", p.load(i), "!=", bytes[i]) 25 | print_result(r) 26 | return 27 | 28 | fn vec[D: DType](*values: SIMD[D, 1]) -> List[SIMD[D, 1]]: 29 | var result = List[SIMD[D, 1]](capacity=len(values)) 30 | for i in range(len(values)): 31 | result.append(values[i]) 32 | return result 33 | 34 | fn test_single_value_contructor(): 35 | assert_result(flx_null(), 0, 0, 1) 36 | assert_result(flx(25), 25, 4, 1) 37 | assert_result(flx(-25), 231, 4, 1) 38 | assert_result(flx(230), 230, 0, 5, 2) 39 | assert_result(flx(-230), 26, 255, 5, 2) 40 | assert_result(flx[DType.uint8](230), 230, 8, 1) 41 | assert_result(flx[DType.uint16](230), 230, 0, 9, 2) 42 | assert_result(flx[DType.uint32](230), 230, 0, 0, 0, 10, 4) 43 | assert_result(flx[DType.uint64](230), 230, 0, 0, 0, 0, 0, 0, 0, 11, 8) 44 | assert_result(flx[DType.bool](True), 1, 104, 1) 45 | assert_result(flx[DType.bool](False), 0, 104, 1) 46 | assert_result(flx[DType.float16](4.5), 128, 68, 13, 2) 47 | assert_result(flx[DType.float32](4.5), 0, 0, 144, 64, 14, 4) 48 | assert_result(flx[DType.float32](0.1), 205, 204, 204, 61, 14, 4) 49 | assert_result(flx("Maxim"), 20, 77, 97, 120, 105, 109, 0, 7, 20, 1) 50 | assert_result(flx("hello 😱"), 40, 104, 101, 108, 108, 111, 32, 240, 159, 152, 177, 0, 12, 20, 1) 51 | assert_result(flx("hello 🔥"), 40, 104, 101, 108, 108, 111, 32, 240, 159, 148, 165, 0, 12, 20, 1) 52 | var v1 = vec[DType.int8](1, 2, 3) 53 | assert_result(flx(DTypePointer[DType.int8](v1.data.value), len(v1)), 12, 1, 2, 3, 4, 44, 1) 54 | _= v1 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 55 | v1 = vec[DType.int8](-1, 2, 3) 56 | assert_result(flx(DTypePointer[DType.int8](v1.data.value), len(v1)), 12, 255, 2, 3, 4, 44, 1) 57 | _= v1 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 58 | var v2 = vec[DType.int16](1, 555, 3) 59 | assert_result(flx(DTypePointer[DType.int16](v2.data.value), len(v2)), 12, 0, 1, 0, 43, 2, 3, 0, 8, 45, 1) 60 | _= v2 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 61 | var v4 = vec[DType.int32](1, 55500, 3) 62 | assert_result( 63 | flx(DTypePointer[DType.int32](v4.data.value), len(v4)), 64 | 12, 0, 0, 0, 1, 0, 0, 0, 204, 216, 0, 0, 3, 0, 0, 0, 16, 46, 1 65 | ) 66 | _= v4 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 67 | var v8 = vec[DType.int64](1, 55555555500, 3) 68 | assert_result( 69 | flx(DTypePointer[DType.int64](v8.data.value), len(v8)), 70 | 12, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 172, 128, 94, 239, 12, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 32, 47, 1 71 | ) 72 | _= v8 # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 73 | var vb = vec[DType.bool](True, False, True) 74 | assert_result(flx(DTypePointer[DType.bool](vb.data.value), len(vb)), 12, 1, 0, 1, 4, 144, 1) 75 | _= vb # needed hack becasue of ASAP descruction policy, will be removed with proper LifeTime feature 76 | 77 | fn test_vec_construction(): 78 | try: 79 | var flx1 = FlxBuffer() 80 | flx1.start_vector() 81 | flx1.add(1) 82 | flx1.add(2) 83 | flx1.add(3) 84 | flx1.end() 85 | assert_result(flx1^.finish(), 12, 1, 2, 3, 4, 44, 1) 86 | 87 | var flx2 = FlxBuffer() 88 | flx2.start_vector() 89 | flx2.add(1) 90 | flx2.add(555) 91 | flx2.add(3) 92 | flx2.end() 93 | assert_result(flx2^.finish(), 12, 0, 1, 0, 43, 2, 3, 0, 8, 45, 1) 94 | 95 | var flx3 = FlxBuffer() 96 | flx3.start_vector() 97 | flx3.add("foo") 98 | flx3.add("bar") 99 | flx3.add("baz") 100 | flx3.end() 101 | assert_result(flx3^.finish(), 12, 102, 111, 111, 0, 12, 98, 97, 114, 0, 12, 98, 97, 122, 0, 12, 16, 12, 8, 4, 60, 1) 102 | 103 | var flx4 = FlxBuffer() 104 | flx4.start_vector() 105 | flx4.add("foo") 106 | flx4.add(1) 107 | flx4.add(-5) 108 | flx4.add[DType.float64](1.3) 109 | flx4.add[DType.bool](True) 110 | flx4.end() 111 | assert_result( 112 | flx4^.finish(), 113 | 12, 102, 111, 111, 0, 0, 0, 0, 114 | 20, 0, 0, 0, 0, 0, 0, 0, 115 | 16, 0, 0, 0, 0, 0, 0, 0, 116 | 1, 0, 0, 0, 0, 0, 0, 0, 117 | 251, 255, 255, 255, 255, 255, 255, 255, 118 | 205, 204, 204, 204, 204, 204, 244, 63, 119 | 1, 0, 0, 0, 0, 0, 0, 0, 120 | 20, 4, 4, 15, 104, 53, 43, 1 121 | ) 122 | 123 | var flx5 = FlxBuffer() 124 | flx5.start_vector() 125 | flx5.key("foo") 126 | flx5.key("bar") 127 | flx5.key("baz") 128 | flx5.end() 129 | assert_result(flx5^.finish(), 102, 111, 111, 0, 98, 97, 114, 0, 98, 97, 122, 0, 12, 13, 10, 7, 4, 56, 1) 130 | 131 | var flx6 = FlxBuffer() 132 | flx6.start_vector() 133 | flx6.start_vector() 134 | flx6.add(61) 135 | flx6.end() 136 | flx6.add(64) 137 | flx6.end() 138 | assert_result(flx6^.finish(), 4, 61, 8, 3, 64, 44, 4, 5, 40, 1) 139 | 140 | except: 141 | print("unexpected error") 142 | 143 | fn test_map_construction(): 144 | try: 145 | var flx1 = FlxBuffer() 146 | flx1.start_map() 147 | flx1.key("a") 148 | flx1.add(12) 149 | flx1.end() 150 | assert_result(flx1^.finish(), 97, 0, 4, 3, 2, 1, 4, 12, 4, 3, 36, 1) 151 | 152 | var flx2 = FlxBuffer() 153 | flx2.start_map() 154 | flx2.key("") 155 | flx2.add(45) 156 | flx2.key("a") 157 | flx2.add(12) 158 | flx2.end() 159 | assert_result(flx2^.finish(), 0, 97, 0, 8, 4, 4, 3, 1, 8, 45, 12, 4, 4, 5, 36, 1) 160 | 161 | var flx3 = FlxBuffer() 162 | flx3.start_map() 163 | flx3.key("b") 164 | flx3.add(45) 165 | flx3.key("a") 166 | flx3.add(12) 167 | flx3.end() 168 | assert_result( 169 | flx3^.finish(), 170 | 98, 0, 97, 0, 8, 3, 6, 3, 1, 8, 12, 45, 4, 4, 5, 36, 1 171 | ) 172 | except: 173 | print("unexpected error") 174 | 175 | fn test_add_referenced() raises: 176 | var flx = FlxBuffer() 177 | flx.start_vector() 178 | flx.start_vector() 179 | flx.add(1) 180 | flx.add(2) 181 | flx.add(3) 182 | flx.end("v1") 183 | flx.add_referenced("v1") 184 | flx.add_referenced("v1") 185 | flx.add_referenced("v1") 186 | flx.end() 187 | assert_result( 188 | flx^.finish(), 189 | 12, 1, 2, 3, 16, 5, 6, 7, 8, 44, 44, 44, 44, 9, 40, 1, 190 | ) 191 | 192 | fn test_map_builder(): 193 | try: 194 | assert_result( 195 | FlxMap().add("", 45).add("a", 12).finish(), 196 | 0, 97, 0, 8, 4, 4, 3, 1, 8, 45, 12, 4, 4, 5, 36, 1 197 | ) 198 | 199 | assert_result( 200 | FlxMap() 201 | .add("a", 12) 202 | .map("b") 203 | .add("c", 33) 204 | .up_to_map() 205 | .add("d", "max") 206 | .finish(), 207 | 97, 0, 98, 0, 99, 0, 4, 3, 2, 1, 4, 33, 4, 100, 0, 12, 109, 97, 120, 0, 208 | 12, 21, 20, 10, 4, 1, 12, 12, 18, 14, 4, 36, 20, 7, 36, 1 209 | ) 210 | 211 | assert_result( 212 | FlxMap() 213 | .add("a", 12) 214 | .vec("b") 215 | .add(33) 216 | .up_to_map() 217 | .add("d", "max") 218 | .finish(), 219 | 97, 0, 98, 0, 4, 33, 100, 0, 12, 109, 97, 120, 0, 220 | 12, 14, 13, 10, 4, 1, 12, 12, 17, 14, 4, 44, 20, 7, 36, 1 221 | ) 222 | 223 | assert_result( 224 | FlxMap().add("a", 1).add_indirect[DType.int32]("b", 2333).add("c", 3).finish(), 225 | 97, 0, 98, 0, 29, 9, 0, 0, 99, 0, 12, 11, 10, 5, 4, 1, 12, 1, 14, 3, 4, 26, 4, 7, 36, 1 226 | ) 227 | except e: 228 | print("unexpected error", e) 229 | 230 | fn test_vec_builder(): 231 | try: 232 | assert_result( 233 | FlxVec().vec().add(61).up_to_vec().add(64).finish(), 234 | 4, 61, 8, 3, 64, 44, 4, 5, 40, 1 235 | ) 236 | assert_result( 237 | FlxVec().add(1).add_indirect[DType.int32](2333).add(3).finish(), 238 | 29, 9, 0, 0, 12, 1, 6, 3, 4, 26, 4, 7, 40, 1 239 | ) 240 | var vec = FlxVec() 241 | for i in range(256): 242 | vec = vec^.add[DType.bool](i & 1 == 1) 243 | assert_result( 244 | vec^.finish(), 245 | 1, 4, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 246 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 247 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 248 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 249 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 250 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 251 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 252 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 253 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 254 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 255 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 256 | 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 1, 144, 2 257 | ) 258 | except e: 259 | print("unexpected error", e) 260 | 261 | fn test_blob() raises: 262 | var data = DTypePointer[DType.uint8].alloc(1001) 263 | for i in range(1001): 264 | data[i] = 5 265 | 266 | var r = flx_blob(data, 1001) 267 | var b = r.get[0, DTypePointer[DType.uint8]]() 268 | var l = r.get[1, Int]() 269 | _ = assert_equal(l, 1008) 270 | _ = assert_equal(b.load(0), 165) 271 | _ = assert_equal(b.load(1), 15) 272 | for i in range(2, 1002): 273 | _ = assert_equal(b.load(i), 5) 274 | # Padding bytes are not zeroed out for now 275 | # _ = assert_equal(b.load(1003), 0) 276 | _ = assert_equal(b.load(1004), 236) 277 | _ = assert_equal(b.load(1005), 3) 278 | _ = assert_equal(b.load(1006), 101) 279 | _ = assert_equal(b.load(1007), 2) 280 | 281 | fn test_dedup_string(): 282 | try: 283 | assert_result( 284 | FlxVec().add("maxim").add("alex").add("maxim").finish(), 285 | 20, 109, 97, 120, 105, 109, 0, 16, 97, 108, 101, 120, 0, 12, 14, 8, 16, 4, 60, 1 286 | ) 287 | assert_result( 288 | FlxVec[dedup_string=False]().add("maxim").add("alex").add("maxim").finish(), 289 | 20, 109, 97, 120, 105, 109, 0, 16, 97, 108, 101, 120, 0, 290 | 20, 109, 97, 120, 105, 109, 0, 12, 21, 15, 10, 4, 60, 1 291 | ) 292 | assert_result( 293 | FlxMap().add("a", "maxim").add("b", "alex").add("c", "maxim").finish(), 294 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 99, 0, 295 | 12, 20, 12, 5, 4, 1, 12, 24, 16, 26, 20, 20, 20, 7, 36, 1 296 | ) 297 | assert_result( 298 | FlxMap[dedup_string=False]().add("a", "maxim").add("b", "alex").add("c", "maxim").finish(), 299 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 99, 0, 300 | 20, 109, 97, 120, 105, 109, 0, 12, 27, 19, 12, 4, 1, 12, 31, 23, 16, 20, 20, 20, 7, 36, 1 301 | ) 302 | except e: 303 | print("unexpected error", e) 304 | 305 | fn test_dedup_key(): 306 | try: 307 | assert_result( 308 | FlxVec[dedup_key=False]() 309 | .map() 310 | .add("a", "maxim") 311 | .add("b", "alex") 312 | .up_to_vec() 313 | .map() 314 | .add("a", "lena") 315 | .add("c", "daria") 316 | .finish(), 317 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 318 | 8, 18, 10, 3, 1, 8, 21, 13, 20, 20, 97, 0, 16, 108, 101, 110, 97, 0, 319 | 99, 0, 20, 100, 97, 114, 105, 97, 0, 8, 18, 11, 3, 1, 8, 21, 14, 20, 320 | 20, 8, 33, 7, 36, 36, 5, 40, 1, 321 | ) 322 | assert_result( 323 | FlxVec() 324 | .map() 325 | .add("a", "maxim") 326 | .add("b", "alex") 327 | .up_to_vec() 328 | .map() 329 | .add("a", "lena") 330 | .add("c", "daria") 331 | .finish(), 332 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 333 | 8, 18, 10, 3, 1, 8, 21, 13, 20, 20, 16, 108, 101, 110, 97, 0, 99, 0, 334 | 20, 100, 97, 114, 105, 97, 0, 8, 43, 11, 3, 1, 8, 21, 14, 20, 335 | 20, 8, 31, 7, 36, 36, 5, 40, 1 336 | ) 337 | assert_result( 338 | FlxVec() 339 | .map() 340 | .add("a", "maxim") 341 | .add("b", "alex") 342 | .up_to_vec() 343 | .map() 344 | .add("a", "lena") 345 | .add("b", "daria") 346 | .finish(), 347 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 348 | 8, 18, 10, 3, 1, 8, 21, 13, 20, 20, 16, 108, 101, 110, 97, 0, 349 | 20, 100, 97, 114, 105, 97, 0, 23, 1, 8, 16, 11, 20, 350 | 20, 8, 26, 7, 36, 36, 5, 40, 1, 351 | ) 352 | assert_result( 353 | FlxVec[dedup_keys_vec=False]() 354 | .map() 355 | .add("a", "maxim") 356 | .add("b", "alex") 357 | .up_to_vec() 358 | .map() 359 | .add("a", "lena") 360 | .add("b", "daria") 361 | .finish(), 362 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 363 | 8, 18, 10, 3, 1, 8, 21, 13, 20, 20, 16, 108, 101, 110, 97, 0, 364 | 20, 100, 97, 114, 105, 97, 0, 8, 41, 33, 3, 1, 8, 19, 14, 20, 365 | 20, 8, 29, 7, 36, 36, 5, 40, 1, 366 | ) 367 | assert_result( 368 | FlxVec[dedup_key=False]() 369 | .map() 370 | .add("a", "maxim") 371 | .add("b", "alex") 372 | .up_to_vec() 373 | .map() 374 | .add("a", "lena") 375 | .add("b", "daria") 376 | .finish(), 377 | 97, 0, 20, 109, 97, 120, 105, 109, 0, 98, 0, 16, 97, 108, 101, 120, 0, 378 | 8, 18, 10, 3, 1, 8, 21, 13, 20, 20, 97, 0, 16, 108, 101, 110, 97, 0, 98, 0, 379 | 20, 100, 97, 114, 105, 97, 0, 8, 18, 11, 3, 1, 8, 21, 14, 20, 380 | 20, 8, 33, 7, 36, 36, 5, 40, 1, 381 | ) 382 | except e: 383 | print("unexpected error", e) 384 | 385 | fn main() raises: 386 | test_single_value_contructor() 387 | test_vec_construction() 388 | test_map_construction() 389 | test_map_builder() 390 | test_vec_builder() 391 | test_blob() 392 | test_dedup_string() 393 | test_dedup_key() 394 | test_add_referenced() 395 | 396 | print("All Done!!!") -------------------------------------------------------------------------------- /test_roundtrip.mojo: -------------------------------------------------------------------------------- 1 | from flx import FlxMap, FlxVec, FlxValue, flx, flx_blob 2 | from testing import assert_equal, assert_almost_equal 3 | 4 | fn test_string() raises: 5 | var r = flx("Hello world") 6 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 7 | _ = assert_equal("Hello world", value.string()) 8 | 9 | fn test_blob() raises: 10 | var data = DTypePointer[DType.uint8].alloc(1001) 11 | for i in range(1001): 12 | data[i] = 5 13 | 14 | var r = flx_blob(data, 1001) 15 | 16 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 17 | var blob = value.blob() 18 | _ = assert_equal(blob.get[1, Int](), 1001) 19 | for i in range(1001): 20 | _ = assert_equal(blob.get[0, DTypePointer[DType.uint8]]().load(i), 5) 21 | 22 | fn test_int() raises: 23 | var r = flx(12345) 24 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 25 | _ = assert_equal(value.int(), 12345) 26 | 27 | r = flx[DType.uint32](12345) 28 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 29 | _ = assert_equal(value.int(), 12345) 30 | 31 | r = flx[DType.uint64](12345) 32 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 33 | _ = assert_equal(value.int(), 12345) 34 | 35 | r = flx[DType.int64](12345) 36 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 37 | _ = assert_equal(value.int(), 12345) 38 | 39 | r = flx[DType.int8](12345) 40 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 41 | _ = assert_equal(value.int(), 57) 42 | 43 | fn test_float() raises: 44 | var r = flx[DType.float64](123.45) 45 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 46 | _ = assert_equal(value.float(), 123.45) 47 | 48 | r = flx[DType.float32](123.45) 49 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 50 | _ = assert_almost_equal(value.float(), 123.45) 51 | 52 | r = flx[DType.float16](123.45) 53 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 54 | _ = assert_almost_equal(value.float(), 123.4375) # Rounding error becuase of 16 bit precision 55 | 56 | fn test_vec() raises: 57 | var r1 = FlxVec().add(1).add(2).add(3).finish() 58 | var value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 59 | _ = assert_equal(value[0].get[DType.int8](), 1) 60 | _ = assert_equal(value[0].int(), 1) 61 | _ = assert_equal(value[1].get[DType.int8](), 2) 62 | _ = assert_equal(value[1].int(), 2) 63 | _ = assert_equal(value[2].get[DType.int8](), 3) 64 | _ = assert_equal(value[2].int(), 3) 65 | 66 | var r2 = FlxVec().add("a").vec().add("b").add("c").finish() 67 | value = FlxValue(r2.get[0, DTypePointer[DType.uint8]](), r2.get[1, Int]()) 68 | _ = assert_equal(value[0].string(), "a") 69 | _ = assert_equal(value[1][0].string(), "b") 70 | _ = assert_equal(value[1][1].string(), "c") 71 | 72 | var r3 = FlxVec().add[DType.float16](1.1).add[DType.float32](1.1).add[DType.float64](1.1).finish() 73 | value = FlxValue(r3.get[0, DTypePointer[DType.uint8]](), r3.get[1, Int]()) 74 | _ = assert_equal(value[0].get[DType.float64](), Float16(1.1)) 75 | _ = assert_equal(value[1].get[DType.float64](), Float32(1.1)) 76 | _ = assert_equal(value[2].get[DType.float64](), 1.1) 77 | 78 | var r4 = FlxVec().add[DType.int8](-13).add[DType.int16](-13).add[DType.int32](-13).add[DType.int64](-13).finish() 79 | value = FlxValue(r4.get[0, DTypePointer[DType.uint8]](), r4.get[1, Int]()) 80 | _ = assert_equal(value[0].get[DType.int64](), -13) 81 | _ = assert_equal(value[0].int(), -13) 82 | _ = assert_equal(value[1].get[DType.int64](), -13) 83 | _ = assert_equal(value[1].int(), -13) 84 | _ = assert_equal(value[2].get[DType.int64](), -13) 85 | _ = assert_equal(value[2].int(), -13) 86 | _ = assert_equal(value[3].get[DType.int64](), -13) 87 | _ = assert_equal(value[3].int(), -13) 88 | 89 | var vec = FlxVec() 90 | for i in range(256): 91 | vec = vec^.add[DType.bool](i & 1 == 1) 92 | var r5 = vec^.finish() 93 | value = FlxValue(r5.get[0, DTypePointer[DType.uint8]](), r5.get[1, Int]()) 94 | _ = assert_equal(256, value.__len__()) 95 | for i in range(256): 96 | _ = assert_equal(i & 1 == 1, value[i].bool()) 97 | 98 | var r6 = FlxVec().add[DType.int32](1234).add("maxim").add[DType.float16](1.5).add[DType.bool](True).finish() 99 | value = FlxValue(r6) 100 | _ = assert_equal(value[0].int(), 1234) 101 | _ = assert_equal(value[1].string(), "maxim") 102 | _ = assert_equal(value[2].get[DType.float32](), 1.5) 103 | _ = assert_equal(value[3].bool(), True) 104 | 105 | fn test_map() raises: 106 | var r1 = FlxMap().add("a", 12).add("b", 45).finish() 107 | var value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 108 | _ = assert_equal(value["a"].get[DType.int8](), 12) 109 | _ = assert_equal(value["a"].int(), 12) 110 | _ = assert_equal(value["b"].get[DType.int8](), 45) 111 | _ = assert_equal(value["b"].int(), 45) 112 | 113 | r1 = FlxMap().add("b", 12).add("a", 45).finish() 114 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 115 | _ = assert_equal(value["b"].get[DType.int8](), 12) 116 | _ = assert_equal(value["a"].get[DType.int8](), 45) 117 | 118 | r1 = FlxMap().add("name", "Maxim").add("age", 42).add[DType.float32]("weight", 72.5).add[DType.bool]("friendly", True).finish() 119 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 120 | _ = assert_equal(value["name"].string(), "Maxim") 121 | _ = assert_equal(value["age"].int(), 42) 122 | _ = assert_equal(value["weight"].get[DType.float32](), 72.5) 123 | _ = assert_equal(value["friendly"].bool(), True) 124 | 125 | r1 = FlxMap() 126 | .add("name", "Maxim") 127 | .add("age", 42) 128 | .add[DType.float32]("weight", 72.5) 129 | .map("address") 130 | .add("city", "Bla") 131 | .add("zip", "12345") 132 | .add("countryCode", "XX") 133 | .up_to_map() 134 | .vec("flags") 135 | .add[DType.bool](True) 136 | .add[DType.bool](False) 137 | .add[DType.bool](True) 138 | .add[DType.bool](True) 139 | .finish() 140 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 141 | _ = assert_equal(value["name"].string(), "Maxim") 142 | _ = assert_equal(value["age"].get[DType.int32](), 42) 143 | _ = assert_equal(value["weight"].get[DType.float32](), 72.5) 144 | _ = assert_equal(value["address"].__len__(), 3) 145 | _ = assert_equal(value["address"]["city"].string(), "Bla") 146 | _ = assert_equal(value["address"]["zip"].string(), "12345") 147 | _ = assert_equal(value["address"]["countryCode"].string(), "XX") 148 | _ = assert_equal(value["flags"].__len__(), 4) 149 | _ = assert_equal(value["flags"][0].get[DType.bool](), True) 150 | _ = assert_equal(value["flags"][1].get[DType.bool](), False) 151 | _ = assert_equal(value["flags"][2].get[DType.bool](), True) 152 | _ = assert_equal(value["flags"][3].get[DType.bool](), True) 153 | 154 | r1 = FlxVec() 155 | .map() 156 | .add("a", "maxim") 157 | .add("b", "alex") 158 | .up_to_vec() 159 | .map() 160 | .add("a", "lena") 161 | .add("c", "daria") 162 | .finish() 163 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 164 | _ = assert_equal(value.__len__(), 2) 165 | _ = assert_equal(value[0].__len__(), 2) 166 | _ = assert_equal(value[0]["a"].string(), "maxim") 167 | _ = assert_equal(value[0]["b"].string(), "alex") 168 | _ = assert_equal(value[1].__len__(), 2) 169 | _ = assert_equal(value[1]["a"].string(), "lena") 170 | _ = assert_equal(value[1]["c"].string(), "daria") 171 | 172 | r1 = FlxVec[dedup_keys_vec=False]() 173 | .map() 174 | .add("a", "maxim") 175 | .add("b", "alex") 176 | .up_to_vec() 177 | .map() 178 | .add("a", "lena") 179 | .add("b", "daria") 180 | .finish() 181 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 182 | _ = assert_equal(value.__len__(), 2) 183 | _ = assert_equal(value[0].__len__(), 2) 184 | _ = assert_equal(value[0]["a"].string(), "maxim") 185 | _ = assert_equal(value[0]["b"].string(), "alex") 186 | _ = assert_equal(value[1].__len__(), 2) 187 | _ = assert_equal(value[1]["a"].string(), "lena") 188 | _ = assert_equal(value[1]["b"].string(), "daria") 189 | 190 | 191 | fn test_indirect() raises: 192 | var r1 = FlxVec().add(1).add_indirect[DType.int32](2333).add(3).finish() 193 | var value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 194 | _ = assert_equal(value.__len__(), 3) 195 | _ = assert_equal(value[0].int(), 1) 196 | _ = assert_equal(value[1].int(), 2333) 197 | _ = assert_equal(value[2].int(), 3) 198 | 199 | r1 = FlxMap().add("a", 1).add_indirect[DType.int32]("b", 2333).add("c", 3).finish() 200 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 201 | _ = assert_equal(value.__len__(), 3) 202 | _ = assert_equal(value["a"].int(), 1) 203 | _ = assert_equal(value["b"].int(), 2333) 204 | _ = assert_equal(value["c"].int(), 3) 205 | 206 | fn test_referenced() raises: 207 | var r1 = FlxVec() 208 | .map() 209 | .vec("vector") 210 | .add("Hello") 211 | .add("World") 212 | .up_to_map(ref_key="v1") 213 | .up_to_vec() 214 | .vec() 215 | .map() 216 | .add("name", "Maxim") 217 | .add("city", "Berlin") 218 | .add_referenced("text", "v1") 219 | .up_to_vec(ref_key="p1") 220 | .add_referenced("p1") 221 | .up_to_vec(ref_key="v2") 222 | .add_referenced("v2") 223 | .add_referenced("p1") 224 | .finish() 225 | var value = FlxValue(r1) 226 | assert_equal(len(value), 4) 227 | assert_equal(len(value[0]), 1) 228 | assert_equal(len(value[0]["vector"]), 2) 229 | assert_equal(value[0]["vector"][0].string(), "Hello") 230 | assert_equal(value[0]["vector"][1].string(), "World") 231 | assert_equal(len(value[1]), 2) 232 | assert_equal(len(value[1][0]), 3) 233 | assert_equal(value[1][0]["name"].string(), "Maxim") 234 | assert_equal(value[1][0]["city"].string(), "Berlin") 235 | assert_equal(value[1][0]["text"][0].string(), "Hello") 236 | assert_equal(value[1][0]["text"][1].string(), "World") 237 | assert_equal(len(value[1][1]), 3) 238 | assert_equal(value[1][1]["name"].string(), "Maxim") 239 | assert_equal(value[1][1]["city"].string(), "Berlin") 240 | assert_equal(value[1][1]["text"][0].string(), "Hello") 241 | assert_equal(value[1][1]["text"][1].string(), "World") 242 | assert_equal(len(value[2]), 2) 243 | assert_equal(len(value[2][0]), 3) 244 | assert_equal(value[2][0]["name"].string(), "Maxim") 245 | assert_equal(value[2][0]["city"].string(), "Berlin") 246 | assert_equal(value[2][0]["text"][0].string(), "Hello") 247 | assert_equal(value[2][0]["text"][1].string(), "World") 248 | assert_equal(len(value[2][1]), 3) 249 | assert_equal(value[2][1]["name"].string(), "Maxim") 250 | assert_equal(value[2][1]["city"].string(), "Berlin") 251 | assert_equal(value[2][1]["text"][0].string(), "Hello") 252 | assert_equal(value[2][1]["text"][1].string(), "World") 253 | assert_equal(value[3]["name"].string(), "Maxim") 254 | assert_equal(value[3]["city"].string(), "Berlin") 255 | assert_equal(value[3]["text"][0].string(), "Hello") 256 | assert_equal(value[3]["text"][1].string(), "World") 257 | var json = value.json() 258 | assert_equal( 259 | json, 260 | '[{"vector":["Hello","World"]},[{"city":"Berlin","name":"Maxim","text":["Hello","World"]},{"city":"Berlin","name":"Maxim","text":["Hello","World"]}],[{"city":"Berlin","name":"Maxim","text":["Hello","World"]},{"city":"Berlin","name":"Maxim","text":["Hello","World"]}],{"city":"Berlin","name":"Maxim","text":["Hello","World"]}]' 261 | ) 262 | assert_equal(len(json), 324) 263 | assert_equal(r1.get[1, Int](), 91) 264 | 265 | fn main(): 266 | try: 267 | test_string() 268 | test_blob() 269 | test_vec() 270 | test_map() 271 | test_int() 272 | test_float() 273 | test_indirect() 274 | test_referenced() 275 | except e: 276 | print("unexpected error", e) 277 | 278 | print("All done!!!") -------------------------------------------------------------------------------- /test_roundtrip2.mojo: -------------------------------------------------------------------------------- 1 | from flx2 import FlxMap, FlxVec, FlxValue, flx, flx_blob 2 | from testing import assert_equal, assert_almost_equal 3 | from test_builder2 import print_result 4 | 5 | fn test_string() raises: 6 | var r = flx("Hello world") 7 | var value = FlxValue(r) 8 | _ = assert_equal("Hello world", value.string()) 9 | 10 | fn test_blob() raises: 11 | var data = DTypePointer[DType.uint8].alloc(1001) 12 | for i in range(1001): 13 | data[i] = 5 14 | 15 | var r = flx_blob(data, 1001) 16 | 17 | var b = r.get[0, DTypePointer[DType.uint8]]() 18 | var l = r.get[1, Int]() 19 | assert_equal(l, 1008) 20 | assert_equal(b.load(0), 165) 21 | assert_equal(b.load(1), 15) 22 | 23 | 24 | var value = FlxValue(b, l) 25 | var blob = value.blob() 26 | _ = assert_equal(blob.get[1, Int](), 1001) 27 | for i in range(1001): 28 | _ = assert_equal(blob.get[0, DTypePointer[DType.uint8]]().load(i), 5) 29 | 30 | fn test_int() raises: 31 | var r = flx(12345) 32 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 33 | _ = assert_equal(value.int(), 12345) 34 | 35 | r = flx[DType.uint32](12345) 36 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 37 | _ = assert_equal(value.int(), 12345) 38 | 39 | r = flx[DType.uint64](12345) 40 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 41 | _ = assert_equal(value.int(), 12345) 42 | 43 | r = flx[DType.int64](12345) 44 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 45 | _ = assert_equal(value.int(), 12345) 46 | 47 | r = flx[DType.int8](12345) 48 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 49 | _ = assert_equal(value.int(), 57) 50 | 51 | fn test_float() raises: 52 | var r = flx[DType.float64](123.45) 53 | var value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 54 | _ = assert_equal(value.float(), 123.45) 55 | 56 | r = flx[DType.float32](123.45) 57 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 58 | _ = assert_almost_equal(value.float(), 123.45) 59 | 60 | r = flx[DType.float16](123.45) 61 | value = FlxValue(r.get[0, DTypePointer[DType.uint8]](), r.get[1, Int]()) 62 | _ = assert_almost_equal(value.float(), 123.4375) # Rounding error becuase of 16 bit precision 63 | 64 | fn test_vec() raises: 65 | var r1 = FlxVec().add(1).add(2).add(3).finish() 66 | var value = FlxValue(r1) 67 | _ = assert_equal(value[0].get[DType.int8](), 1) 68 | _ = assert_equal(value[0].int(), 1) 69 | _ = assert_equal(value[1].get[DType.int8](), 2) 70 | _ = assert_equal(value[1].int(), 2) 71 | _ = assert_equal(value[2].get[DType.int8](), 3) 72 | _ = assert_equal(value[2].int(), 3) 73 | 74 | var r2 = FlxVec().add("a").vec().add("b").add("c").finish() 75 | value = FlxValue(r2) 76 | _ = assert_equal(value[0].string(), "a") 77 | _ = assert_equal(value[1][0].string(), "b") 78 | _ = assert_equal(value[1][1].string(), "c") 79 | 80 | var r3 = FlxVec().add[DType.float16](1.1).add[DType.float32](1.1).add[DType.float64](1.1).finish() 81 | value = FlxValue(r3) 82 | _ = assert_equal(value[0].get[DType.float64](), Float16(1.1)) 83 | _ = assert_equal(value[1].get[DType.float64](), Float32(1.1)) 84 | _ = assert_equal(value[2].get[DType.float64](), 1.1) 85 | 86 | var r4 = FlxVec().add[DType.int8](-13).add[DType.int16](-13).add[DType.int32](-13).add[DType.int64](-13).finish() 87 | value = FlxValue(r4.get[0, DTypePointer[DType.uint8]](), r4.get[1, Int]()) 88 | _ = assert_equal(value[0].get[DType.int64](), -13) 89 | _ = assert_equal(value[0].int(), -13) 90 | _ = assert_equal(value[1].get[DType.int64](), -13) 91 | _ = assert_equal(value[1].int(), -13) 92 | _ = assert_equal(value[2].get[DType.int64](), -13) 93 | _ = assert_equal(value[2].int(), -13) 94 | _ = assert_equal(value[3].get[DType.int64](), -13) 95 | _ = assert_equal(value[3].int(), -13) 96 | 97 | var vec = FlxVec() 98 | for i in range(256): 99 | vec = vec^.add[DType.bool](i & 1 == 1) 100 | var r5 = vec^.finish() 101 | value = FlxValue(r5.get[0, DTypePointer[DType.uint8]](), r5.get[1, Int]()) 102 | _ = assert_equal(256, value.__len__()) 103 | for i in range(256): 104 | _ = assert_equal(i & 1 == 1, value[i].bool()) 105 | 106 | var r6 = FlxVec().add[DType.int32](1234).add("maxim").add[DType.float16](1.5).add[DType.bool](True).finish() 107 | value = FlxValue(r6.get[0, DTypePointer[DType.uint8]](), r6.get[1, Int]()) 108 | _ = assert_equal(value[0].int(), 1234) 109 | _ = assert_equal(value[1].string(), "maxim") 110 | _ = assert_equal(value[2].get[DType.float32](), 1.5) 111 | _ = assert_equal(value[3].bool(), True) 112 | 113 | fn test_map() raises: 114 | var r1 = FlxMap().add("a", 12).add("b", 45).finish() 115 | var value = FlxValue(r1) 116 | _ = assert_equal(value["a"].get[DType.int8](), 12) 117 | _ = assert_equal(value["a"].int(), 12) 118 | _ = assert_equal(value["b"].get[DType.int8](), 45) 119 | _ = assert_equal(value["b"].int(), 45) 120 | 121 | r1 = FlxMap().add("b", 12).add("a", 45).finish() 122 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 123 | _ = assert_equal(value["b"].get[DType.int8](), 12) 124 | _ = assert_equal(value["a"].get[DType.int8](), 45) 125 | 126 | r1 = FlxMap().add("name", "Maxim").add("age", 42).add[DType.float32]("weight", 72.5).add[DType.bool]("friendly", True).finish() 127 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 128 | _ = assert_equal(value["name"].string(), "Maxim") 129 | _ = assert_equal(value["age"].int(), 42) 130 | _ = assert_equal(value["weight"].get[DType.float32](), 72.5) 131 | _ = assert_equal(value["friendly"].bool(), True) 132 | 133 | r1 = FlxMap() 134 | .add("name", "Maxim") 135 | .add("age", 42) 136 | .add[DType.float32]("weight", 72.5) 137 | .map("address") 138 | .add("city", "Bla") 139 | .add("zip", "12345") 140 | .add("countryCode", "XX") 141 | .up_to_map() 142 | .vec("flags") 143 | .add[DType.bool](True) 144 | .add[DType.bool](False) 145 | .add[DType.bool](True) 146 | .add[DType.bool](True) 147 | .finish() 148 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 149 | _ = assert_equal(value["name"].string(), "Maxim") 150 | _ = assert_equal(value["age"].get[DType.int32](), 42) 151 | _ = assert_equal(value["weight"].get[DType.float32](), 72.5) 152 | _ = assert_equal(value["address"].__len__(), 3) 153 | _ = assert_equal(value["address"]["city"].string(), "Bla") 154 | _ = assert_equal(value["address"]["zip"].string(), "12345") 155 | _ = assert_equal(value["address"]["countryCode"].string(), "XX") 156 | _ = assert_equal(value["flags"].__len__(), 4) 157 | _ = assert_equal(value["flags"][0].get[DType.bool](), True) 158 | _ = assert_equal(value["flags"][1].get[DType.bool](), False) 159 | _ = assert_equal(value["flags"][2].get[DType.bool](), True) 160 | _ = assert_equal(value["flags"][3].get[DType.bool](), True) 161 | 162 | r1 = FlxVec() 163 | .map() 164 | .add("a", "maxim") 165 | .add("b", "alex") 166 | .up_to_vec() 167 | .map() 168 | .add("a", "lena") 169 | .add("c", "daria") 170 | .finish() 171 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 172 | _ = assert_equal(value.__len__(), 2) 173 | _ = assert_equal(value[0].__len__(), 2) 174 | _ = assert_equal(value[0]["a"].string(), "maxim") 175 | _ = assert_equal(value[0]["b"].string(), "alex") 176 | _ = assert_equal(value[1].__len__(), 2) 177 | _ = assert_equal(value[1]["a"].string(), "lena") 178 | _ = assert_equal(value[1]["c"].string(), "daria") 179 | 180 | r1 = FlxVec[dedup_keys_vec=False]() 181 | .map() 182 | .add("a", "maxim") 183 | .add("b", "alex") 184 | .up_to_vec() 185 | .map() 186 | .add("a", "lena") 187 | .add("b", "daria") 188 | .finish() 189 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 190 | _ = assert_equal(value.__len__(), 2) 191 | _ = assert_equal(value[0].__len__(), 2) 192 | _ = assert_equal(value[0]["a"].string(), "maxim") 193 | _ = assert_equal(value[0]["b"].string(), "alex") 194 | _ = assert_equal(value[1].__len__(), 2) 195 | _ = assert_equal(value[1]["a"].string(), "lena") 196 | _ = assert_equal(value[1]["b"].string(), "daria") 197 | 198 | 199 | fn test_indirect() raises: 200 | var r1 = FlxVec().add(1).add_indirect[DType.int32](2333).add(3).finish() 201 | var value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 202 | _ = assert_equal(value.__len__(), 3) 203 | _ = assert_equal(value[0].int(), 1) 204 | _ = assert_equal(value[1].int(), 2333) 205 | _ = assert_equal(value[2].int(), 3) 206 | 207 | r1 = FlxMap().add("a", 1).add_indirect[DType.int32]("b", 2333).add("c", 3).finish() 208 | value = FlxValue(r1.get[0, DTypePointer[DType.uint8]](), r1.get[1, Int]()) 209 | _ = assert_equal(value.__len__(), 3) 210 | _ = assert_equal(value["a"].int(), 1) 211 | _ = assert_equal(value["b"].int(), 2333) 212 | _ = assert_equal(value["c"].int(), 3) 213 | 214 | fn test_referenced() raises: 215 | var r1 = FlxVec() 216 | .map() 217 | .vec("vector") 218 | .add("Hello") 219 | .add("World") 220 | .up_to_map(ref_key="v1") 221 | .up_to_vec() 222 | .vec() 223 | .map() 224 | .add("name", "Maxim") 225 | .add("city", "Berlin") 226 | .add_referenced("text", "v1") 227 | .up_to_vec(ref_key="p1") 228 | .add_referenced("p1") 229 | .up_to_vec(ref_key="v2") 230 | .add_referenced("v2") 231 | .add_referenced("p1") 232 | .finish() 233 | var value = FlxValue(r1) 234 | assert_equal(len(value), 4) 235 | assert_equal(len(value[0]), 1) 236 | assert_equal(len(value[0]["vector"]), 2) 237 | assert_equal(value[0]["vector"][0].string(), "Hello") 238 | assert_equal(value[0]["vector"][1].string(), "World") 239 | assert_equal(len(value[1]), 2) 240 | assert_equal(len(value[1][0]), 3) 241 | assert_equal(value[1][0]["name"].string(), "Maxim") 242 | assert_equal(value[1][0]["city"].string(), "Berlin") 243 | assert_equal(value[1][0]["text"][0].string(), "Hello") 244 | assert_equal(value[1][0]["text"][1].string(), "World") 245 | assert_equal(len(value[1][1]), 3) 246 | assert_equal(value[1][1]["name"].string(), "Maxim") 247 | assert_equal(value[1][1]["city"].string(), "Berlin") 248 | assert_equal(value[1][1]["text"][0].string(), "Hello") 249 | assert_equal(value[1][1]["text"][1].string(), "World") 250 | assert_equal(len(value[2]), 2) 251 | assert_equal(len(value[2][0]), 3) 252 | assert_equal(value[2][0]["name"].string(), "Maxim") 253 | assert_equal(value[2][0]["city"].string(), "Berlin") 254 | assert_equal(value[2][0]["text"][0].string(), "Hello") 255 | assert_equal(value[2][0]["text"][1].string(), "World") 256 | assert_equal(len(value[2][1]), 3) 257 | assert_equal(value[2][1]["name"].string(), "Maxim") 258 | assert_equal(value[2][1]["city"].string(), "Berlin") 259 | assert_equal(value[2][1]["text"][0].string(), "Hello") 260 | assert_equal(value[2][1]["text"][1].string(), "World") 261 | assert_equal(value[3]["name"].string(), "Maxim") 262 | assert_equal(value[3]["city"].string(), "Berlin") 263 | assert_equal(value[3]["text"][0].string(), "Hello") 264 | assert_equal(value[3]["text"][1].string(), "World") 265 | var json = value.json() 266 | assert_equal( 267 | json, 268 | '[{"vector":["Hello","World"]},[{"city":"Berlin","name":"Maxim","text":["Hello","World"]},{"city":"Berlin","name":"Maxim","text":["Hello","World"]}],[{"city":"Berlin","name":"Maxim","text":["Hello","World"]},{"city":"Berlin","name":"Maxim","text":["Hello","World"]}],{"city":"Berlin","name":"Maxim","text":["Hello","World"]}]' 269 | ) 270 | assert_equal(len(json), 324) 271 | assert_equal(r1.get[1, Int](), 91) 272 | 273 | fn main(): 274 | try: 275 | test_string() 276 | test_blob() 277 | test_vec() 278 | test_map() 279 | test_int() 280 | test_float() 281 | test_indirect() 282 | test_referenced() 283 | except e: 284 | print("unexpected error", e) 285 | 286 | print("All done!!!") 287 | -------------------------------------------------------------------------------- /test_value_to_json.mojo: -------------------------------------------------------------------------------- 1 | from flx import * 2 | from testing import assert_equal 3 | 4 | fn main() raises: 5 | _ = assert_equal("null", FlxValue(flx_null()).json()) 6 | _ = assert_equal("123", FlxValue(flx(123)).json()) 7 | _ = assert_equal("12345", FlxValue(flx(12345)).json()) 8 | _ = assert_equal("123.45", FlxValue(flx[DType.float64](123.45)).json()) 9 | _ = assert_equal('"Hello 🔥"', FlxValue(flx("Hello 🔥")).json()) 10 | var v1 = FlxVec().add[DType.int32](1234).add("maxim").add[DType.float16](1.5).add[DType.bool](True).finish() 11 | _ = assert_equal('[1234,"maxim",1.5,true]', FlxValue(v1).json()) 12 | v1 = FlxVec().map().add("a", 7).add("b", 8).up_to_vec().map().add("b", 42).add("a", 43).finish() 13 | _ = assert_equal('[{"a":7,"b":8},{"a":43,"b":42}]', FlxValue(v1).json()) 14 | v1 = FlxVec().map().add("a", 7).add("b", 8).up_to_vec().map().add("b", "42").vec("a").add_indirect[DType.float64](1.2).add[DType.bool](False).null().finish() 15 | _ = assert_equal('[{"a":7,"b":8},{"a":[1.2,false,null],"b":"42"}]', FlxValue(v1).json()) 16 | 17 | print("Done!!!") -------------------------------------------------------------------------------- /type_conversion_sample.mojo: -------------------------------------------------------------------------------- 1 | from random import random_si64 2 | from flx import FlxVec, FlxValue 3 | from memory.unsafe import bitcast 4 | 5 | fn convert_dynamic_vector() raises: 6 | var v = List[Int32]() 7 | for _ in range(10_000): 8 | v.append(random_si64(-1_000_000, 1_000_000).cast[DType.int32]()) 9 | 10 | var flx_result = FlxVec() 11 | .add("vec_i32") 12 | .add(rebind[DTypePointer[DType.int32]](v.data).bitcast[DType.uint8](), len(v) * sizeof[Int32]()) 13 | .finish() 14 | 15 | print("Buffer size:", flx_result.get[1, Int]()) 16 | 17 | var value = FlxValue(flx_result.get[0, DTypePointer[DType.uint8]](), flx_result.get[1, Int]()) 18 | print(value[0].string()) 19 | var blob = value[1].blob() 20 | var size = blob.get[1, Int]() / 4 21 | var bytes = blob.get[0, DTypePointer[DType.uint8]]() 22 | var bytes32 = bytes.bitcast[DType.int32]() 23 | var v1 = List[Int32](capacity=size.to_int()) 24 | for i in range(size.to_int()): 25 | v1.append(bytes32.offset(i).load()) 26 | 27 | if len(v) != len(v1): 28 | print("Error! Lengths are not equal", len(v), len(v1)) 29 | return 30 | for i in range(len(v)): 31 | if v[i] != v1[i]: 32 | print("Error at:", i, v[i], "!=", v1[i]) 33 | 34 | fn main(): 35 | try: 36 | convert_dynamic_vector() 37 | except e: 38 | print("Unexpected error:", e) --------------------------------------------------------------------------------