Skip to content

Commit

Permalink
[Feature] Simple map (#182)
Browse files Browse the repository at this point in the history
* Add SimpleMap

* Add SimpleMap
  • Loading branch information
WGB5445 committed Oct 13, 2022
1 parent e5a3108 commit daa971f
Show file tree
Hide file tree
Showing 2 changed files with 230 additions and 0 deletions.
1 change: 1 addition & 0 deletions build/StarcoinFramework/docs/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ This is the root document for the Move StarcoinFramework module documentation. T
- [`0x1::Signature`](Signature.md#0x1_Signature)
- [`0x1::SignedInteger64`](SignedInteger64.md#0x1_SignedInteger64)
- [`0x1::Signer`](Signer.md#0x1_Signer)
- [`0x1::SimpleMap`](SimpleMap.md#0x1_SimpleMap)
- [`0x1::SnapshotUtil`](SnapshotUtil.md#0x1_SnapshotUtil)
- [`0x1::StakeToSBTPlugin`](StakeToSBTPlugin.md#0x1_StakeToSBTPlugin)
- [`0x1::StarcoinDAO`](StarcoinDAO.md#0x1_StarcoinDAO)
Expand Down
229 changes: 229 additions & 0 deletions sources/SimpleMap.move
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);
}
}

0 comments on commit daa971f

Please sign in to comment.