Last active
January 20, 2021 20:50
-
-
Save danielkroeni/0c8bc01ad84f4c488c72abb8fe71e4e7 to your computer and use it in GitHub Desktop.
Immutable Treemap
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
| %builtins output range_check | |
| from starkware.cairo.common.registers import get_fp_and_pc | |
| from starkware.cairo.common.serialize import serialize_word | |
| from starkware.cairo.common.math import assert_lt_felt | |
| const false = 0 | |
| const true = 1 | |
| struct TreeMap: | |
| member key = 0 | |
| member value = 1 | |
| member left: TreeMap* = 2 | |
| member right: TreeMap* = 3 | |
| member empty = 4 # 0 non_empty, 1 empty | |
| const SIZE = 5 | |
| end | |
| struct Result: | |
| member value = 0 | |
| member empty = 1 # 0 found value, 1 key not found | |
| const SIZE = 2 | |
| end | |
| func empty() -> (map: TreeMap*): | |
| alloc_locals | |
| local empty_map: TreeMap | |
| #assert empty_map.key = 0 | |
| #assert empty_map.value = 0 | |
| #assert empty_map.left = 0 | |
| #assert empty_map.right = 0 | |
| assert empty_map.empty = true | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (&empty_map) | |
| end | |
| func put(range_check_ptr, key, value, map: TreeMap*) -> (range_check_ptr, map: TreeMap*): | |
| alloc_locals | |
| local result: TreeMap | |
| if map.empty == true: | |
| assert result.key = key | |
| assert result.value = value | |
| assert result.left = map | |
| assert result.right = map | |
| assert result.empty = false | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| else: | |
| if key == map.key: | |
| assert result.key = map.key | |
| assert result.value = value | |
| assert result.left = map.left | |
| assert result.right = map.right | |
| assert result.empty = false | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| else: | |
| local go_left | |
| %{ | |
| ids.go_left = int(ids.key < ids.map.key) | |
| %} | |
| if go_left == true: | |
| let (range_check_ptr) = assert_lt_felt(range_check_ptr, key, map.key) | |
| let (range_check_ptr, new_left) = put(range_check_ptr, key,value, map.left) | |
| assert result.key = map.key | |
| assert result.value = map.value | |
| assert result.left = new_left | |
| assert result.right = map.right | |
| assert result.empty = false | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| else: | |
| let (range_check_ptr) = assert_lt_felt(range_check_ptr, map.key, key) | |
| let (range_check_ptr, new_right) = put(range_check_ptr, key,value, map.right) | |
| assert result.key = map.key | |
| assert result.value = map.value | |
| assert result.left = map.left | |
| assert result.right = new_right | |
| assert result.empty = false | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| end | |
| end | |
| end | |
| end | |
| func get(range_check_ptr, key, map: TreeMap*) -> (range_check_ptr, result: Result*): | |
| alloc_locals | |
| local result: Result | |
| if map.empty == true: | |
| # assert result.value = 0 | |
| assert result.empty = true | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| else: | |
| if key == map.key: | |
| assert result.value = map.value | |
| assert result.empty = false | |
| let (__fp__, _) = get_fp_and_pc() | |
| return (range_check_ptr, &result) | |
| else: | |
| local go_left | |
| %{ | |
| ids.go_left = int(ids.key < ids.map.key) | |
| %} | |
| if go_left == true: | |
| let (range_check_ptr) = assert_lt_felt(range_check_ptr, key, map.key) | |
| get(range_check_ptr, key, map.left) | |
| return (...) | |
| else: | |
| let (range_check_ptr) = assert_lt_felt(range_check_ptr, map.key, key) | |
| get(range_check_ptr, key, map.right) | |
| return (...) | |
| end | |
| end | |
| end | |
| end | |
| func get_or_default(range_check_ptr, key, defaulValue, map: TreeMap*) -> (range_check_ptr, value): | |
| let (range_check_ptr, result) = get(range_check_ptr, key, map) | |
| if result.empty == true: | |
| return (range_check_ptr, defaulValue) | |
| else: | |
| return (range_check_ptr, result.value) | |
| end | |
| end | |
| func output_treemap(output_ptr : felt*, map: TreeMap*) -> (output_ptr : felt*): | |
| if map.empty == 1: | |
| return (output_ptr) | |
| else: | |
| let (output_ptr) = output_treemap(output_ptr, map.left) | |
| let (output_ptr) = serialize_word(output_ptr, map.key) | |
| let (output_ptr) = serialize_word(output_ptr, map.value) | |
| let (output_ptr) = output_treemap(output_ptr, map.right) | |
| return (output_ptr) | |
| end | |
| end | |
| # "Tests" insert and lookup | |
| func main(output_ptr : felt*, range_check_ptr) -> (output_ptr : felt*, range_check_ptr): | |
| alloc_locals | |
| let (map) = empty() | |
| let (range_check_ptr, map) = put(range_check_ptr, 2, 22, map) # 2->22 | |
| let (range_check_ptr, map) = put(range_check_ptr, 3, 32, map) # 2->22, 3->32 | |
| let (range_check_ptr, map) = put(range_check_ptr, 4, 44, map) # 2->22, 3->32, 4->44 | |
| let (range_check_ptr, map) = put(range_check_ptr, 3, 33, map) # 2->22, 3->33, 4->44 | |
| let (local range_check_ptr, local map: TreeMap*) = put(range_check_ptr, 1, 11, map) # 1->11, 2->22, 3->33, 4->44 | |
| let (local output_ptr: felt*) = output_treemap(output_ptr, map) | |
| let (local range_check_ptr, res) = get(range_check_ptr, 2, map) # should return 22 | |
| let (local output_ptr: felt*) = serialize_word(output_ptr, res.value) | |
| let (range_check_ptr, res) = get(range_check_ptr, 3, map) # should return 33 | |
| let (local output_ptr: felt*) = serialize_word(output_ptr, res.value) | |
| let (range_check_ptr, res) = get(range_check_ptr, 5, map) # should return empty == 1 | |
| let (local output_ptr: felt*) = serialize_word(output_ptr, res.empty) | |
| let (range_check_ptr, res) = get_or_default(range_check_ptr, 5, -1, map) # should return -1 | |
| let (output_ptr) = serialize_word(output_ptr, res) | |
| return(output_ptr=output_ptr, range_check_ptr=range_check_ptr) | |
| end | |
| #> cairo-compile treemap.cairo --output treemap.json && cairo-run --program=treemap.json --print_output --layout=small |
Author
Author
Added type annotations to TreeMap struct and removed casts.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Added
starkware.cairo.common.math.assert_lt_feltchecks to ensure that the prover's hints do not lie.