Created
January 28, 2025 11:57
-
-
Save sstadick/20cfda1db33041bd0ebed6798757c22d to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from memory import UnsafePointer | |
| from os import abort | |
| @value | |
| struct TreeNode: | |
| var data: String | |
| var key: String | |
| var hash: Int | |
| var left: UnsafePointer[Self] | |
| var right: UnsafePointer[Self] | |
| fn __init__(out self, owned key: String, owned data: String): | |
| self.data = data^ | |
| self.key = key^ | |
| self.hash = hash(self.key) | |
| self.left = UnsafePointer[Self]() | |
| self.right = UnsafePointer[Self]() | |
| fn write_to[W: Writer](self, mut writer: W): | |
| writer.write('"', self.key, '"', ":", self.data) | |
| @staticmethod | |
| fn make_ptr( | |
| owned key: String, owned data: String, out p: UnsafePointer[Self] | |
| ): | |
| p = UnsafePointer[Self].alloc(1) | |
| p.init_pointee_move(TreeNode(key^, data^)) | |
| fn for_each[ | |
| visit: fn (UnsafePointer[TreeNode]) capturing [_] | |
| ](node: UnsafePointer[TreeNode]): | |
| if not node: | |
| return | |
| for_each[visit](node[].left) | |
| visit(node) | |
| for_each[visit](node[].right) | |
| struct Tree: | |
| alias NodePtr = UnsafePointer[TreeNode] | |
| var root: Self.NodePtr | |
| var size: UInt | |
| fn __init__(out self): | |
| self.root = Self.NodePtr() | |
| self.size = 0 | |
| fn insert(mut self, owned key: String, owned data: String): | |
| self.size += 1 | |
| var node = TreeNode.make_ptr(key^, data^) | |
| if not self.root: | |
| self.root = node | |
| return | |
| var parent = Self.NodePtr() | |
| var curr = self.root | |
| while curr: | |
| parent = curr | |
| if curr[].hash > node[].hash: | |
| curr = curr[].left | |
| elif curr[].hash < node[].hash: | |
| curr = curr[].right | |
| else: | |
| curr[].data = node[].data | |
| return | |
| if parent[].hash > node[].hash: | |
| parent[].left = node | |
| else: | |
| parent[].right = node | |
| fn write_to[W: Writer](self, mut writer: W): | |
| writer.write("{") | |
| var written = 0 | |
| var size = self.size | |
| print("Size is", size) | |
| @parameter | |
| @always_inline | |
| fn write_node(node: Self.NodePtr): | |
| print("\nSize in write_node", size, "\n") | |
| print("\nWritten in write_node", written, "\n") | |
| writer.write(node[]) | |
| if written < size - 1: | |
| writer.write(",") | |
| written += 1 | |
| for_each[write_node](self.root) | |
| writer.write("}") | |
| fn __str__(self) -> String: | |
| return String.write(self) | |
| def main(): | |
| var t = Tree() | |
| t.insert("fun", "123") | |
| t.insert("funnyyy", "some string") | |
| print(t) # prints {"fun":8753160865977820706,} so half garbage? |
Author
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Prints