-
Notifications
You must be signed in to change notification settings - Fork 27
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Add SimpleMap * Add SimpleMap
- Loading branch information
Showing
2 changed files
with
230 additions
and
0 deletions.
There are no files selected for viewing
This file contains 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
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,229 @@ | ||
/// This module provides a solution for sorted maps, that is it has the properties that | ||
/// 1) Keys point to Values | ||
/// 2) Each Key must be unique | ||
/// 3) A Key can be found within O(Log N) time | ||
/// 4) The data is stored as sorted by Key | ||
/// 5) Adds and removals take O(N) time | ||
module StarcoinFramework::SimpleMap { | ||
|
||
use StarcoinFramework::Errors; | ||
use StarcoinFramework::Option; | ||
use StarcoinFramework::Vector; | ||
use StarcoinFramework::Compare; | ||
use StarcoinFramework::BCS; | ||
|
||
spec module { | ||
pragma verify = false; | ||
pragma aborts_if_is_strict = true; | ||
} | ||
|
||
/// Map key already exists | ||
const EKEY_ALREADY_EXISTS: u64 = 1; | ||
/// Map key is not found | ||
const EKEY_NOT_FOUND: u64 = 2; | ||
|
||
struct SimpleMap<Key, Value> has copy, drop, store { | ||
data: vector<Element<Key, Value>>, | ||
} | ||
|
||
struct Element<Key, Value> has copy, drop, store { | ||
key: Key, | ||
value: Value, | ||
} | ||
|
||
public fun length<Key: store, Value: store>(map: &SimpleMap<Key, Value>): u64 { | ||
Vector::length(&map.data) | ||
} | ||
|
||
public fun create<Key: store, Value: store>(): SimpleMap<Key, Value> { | ||
SimpleMap { | ||
data: Vector::empty(), | ||
} | ||
} | ||
|
||
public fun borrow<Key: store, Value: store>( | ||
map: &SimpleMap<Key, Value>, | ||
key: &Key, | ||
): &Value { | ||
let (maybe_idx, _) = find(map, key); | ||
assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND)); | ||
let idx = Option::extract(&mut maybe_idx); | ||
&Vector::borrow(&map.data, idx).value | ||
} | ||
|
||
public fun borrow_mut<Key: store, Value: store>( | ||
map: &mut SimpleMap<Key, Value>, | ||
key: &Key, | ||
): &mut Value { | ||
let (maybe_idx, _) = find(map, key); | ||
assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND)); | ||
let idx = Option::extract(&mut maybe_idx); | ||
&mut Vector::borrow_mut(&mut map.data, idx).value | ||
} | ||
|
||
public fun contains_key<Key: store, Value: store>( | ||
map: &SimpleMap<Key, Value>, | ||
key: &Key, | ||
): bool { | ||
let (maybe_idx, _) = find(map, key); | ||
Option::is_some(&maybe_idx) | ||
} | ||
|
||
public fun destroy_empty<Key: store, Value: store>(map: SimpleMap<Key, Value>) { | ||
let SimpleMap { data } = map; | ||
Vector::destroy_empty(data); | ||
} | ||
|
||
public fun add<Key: store, Value: store>( | ||
map: &mut SimpleMap<Key, Value>, | ||
key: Key, | ||
value: Value, | ||
) { | ||
let (maybe_idx, maybe_placement) = find(map, &key); | ||
assert!(Option::is_none(&maybe_idx), Errors::invalid_argument(EKEY_ALREADY_EXISTS)); | ||
|
||
// Append to the end and then swap elements until the list is ordered again | ||
Vector::push_back(&mut map.data, Element { key, value }); | ||
|
||
let placement = Option::extract(&mut maybe_placement); | ||
let end = Vector::length(&map.data) - 1; | ||
while (placement < end) { | ||
Vector::swap(&mut map.data, placement, end); | ||
placement = placement + 1; | ||
}; | ||
} | ||
|
||
public fun remove<Key: store, Value: store>( | ||
map: &mut SimpleMap<Key, Value>, | ||
key: &Key, | ||
): (Key, Value) { | ||
let (maybe_idx, _) = find(map, key); | ||
assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND)); | ||
|
||
let placement = Option::extract(&mut maybe_idx); | ||
let end = Vector::length(&map.data) - 1; | ||
|
||
while (placement < end) { | ||
Vector::swap(&mut map.data, placement, placement + 1); | ||
placement = placement + 1; | ||
}; | ||
|
||
let Element { key, value } = Vector::pop_back(&mut map.data); | ||
(key, value) | ||
} | ||
|
||
fun find<Key: store, Value: store>( | ||
map: &SimpleMap<Key, Value>, | ||
key: &Key, | ||
): (Option::Option<u64>, Option::Option<u64>) { | ||
let length = Vector::length(&map.data); | ||
|
||
if (length == 0) { | ||
return (Option::none(), Option::some(0)) | ||
}; | ||
|
||
let left = 0; | ||
let right = length; | ||
|
||
while (left != right) { | ||
let mid = left + (right - left) / 2; | ||
let potential_key = &Vector::borrow(&map.data, mid).key; | ||
if (Compare::is_less_than(Compare::cmp_bytes(&BCS::to_bytes(potential_key), &BCS::to_bytes(key)))) { | ||
left = mid + 1; | ||
} else { | ||
right = mid; | ||
}; | ||
}; | ||
|
||
if (left != length && key == &Vector::borrow(&map.data, left).key) { | ||
(Option::some(left), Option::none()) | ||
} else { | ||
(Option::none(), Option::some(left)) | ||
} | ||
} | ||
|
||
#[test] | ||
public fun add_remove_many() { | ||
let map = create<u64, u64>(); | ||
|
||
assert!(length(&map) == 0, 0); | ||
assert!(!contains_key(&map, &3), 1); | ||
add(&mut map, 3, 1); | ||
assert!(length(&map) == 1, 2); | ||
assert!(contains_key(&map, &3), 3); | ||
assert!(borrow(&map, &3) == &1, 4); | ||
*borrow_mut(&mut map, &3) = 2; | ||
assert!(borrow(&map, &3) == &2, 5); | ||
|
||
assert!(!contains_key(&map, &2), 6); | ||
add(&mut map, 2, 5); | ||
assert!(length(&map) == 2, 7); | ||
assert!(contains_key(&map, &2), 8); | ||
assert!(borrow(&map, &2) == &5, 9); | ||
*borrow_mut(&mut map, &2) = 9; | ||
assert!(borrow(&map, &2) == &9, 10); | ||
|
||
remove(&mut map, &2); | ||
assert!(length(&map) == 1, 11); | ||
assert!(!contains_key(&map, &2), 12); | ||
assert!(borrow(&map, &3) == &2, 13); | ||
|
||
remove(&mut map, &3); | ||
assert!(length(&map) == 0, 14); | ||
assert!(!contains_key(&map, &3), 15); | ||
|
||
destroy_empty(map); | ||
} | ||
|
||
#[test] | ||
public fun test_several() { | ||
let map = create<u64, u64>(); | ||
add(&mut map, 6, 6); | ||
add(&mut map, 1, 1); | ||
add(&mut map, 5, 5); | ||
add(&mut map, 2, 2); | ||
add(&mut map, 3, 3); | ||
add(&mut map, 0, 0); | ||
add(&mut map, 7, 7); | ||
add(&mut map, 4, 4); | ||
|
||
let idx = 0; | ||
while (idx < Vector::length(&map.data)) { | ||
assert!(idx == Vector::borrow(&map.data, idx).key, idx); | ||
idx = idx + 1; | ||
}; | ||
|
||
remove(&mut map, &0); | ||
remove(&mut map, &1); | ||
remove(&mut map, &2); | ||
remove(&mut map, &3); | ||
remove(&mut map, &4); | ||
remove(&mut map, &5); | ||
remove(&mut map, &6); | ||
remove(&mut map, &7); | ||
|
||
destroy_empty(map); | ||
} | ||
|
||
#[test] | ||
#[expected_failure] | ||
public fun add_twice() { | ||
let map = create<u64, u64>(); | ||
add(&mut map, 3, 1); | ||
add(&mut map, 3, 1); | ||
|
||
remove(&mut map, &3); | ||
destroy_empty(map); | ||
} | ||
|
||
#[test] | ||
#[expected_failure] | ||
public fun remove_twice() { | ||
let map = create<u64, u64>(); | ||
add(&mut map, 3, 1); | ||
remove(&mut map, &3); | ||
remove(&mut map, &3); | ||
|
||
destroy_empty(map); | ||
} | ||
} |