Merge pull request #17 from savi-lang/add/get-or-insert

Add `get_or_insert` method.

spec/Map.Ordered.Spec.savi

@@ -22,6 +22,28 @@
     assert: map.has_key("example").is_false
     assert: map.size == 0
 
+  :it "gets an existing value or inserts if it didn't exist"
+    map = @new_map
+    assert: map.size == 0
+    assert error: map["example"]!
+    assert: map.has_key("example").is_false
+    assert: (map.get_or_insert("example") -> (U64[99])) == 99 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 99
+    assert: map.has_key("example")
+    assert: (map.get_or_insert("example") -> (U64[100])) == 99 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 99
+    assert: map.has_key("example")
+    assert: map.delete("example") <: None
+    assert: map.size == 0
+    assert error: map["example"]!
+    assert: map.has_key("example").is_false
+    assert: (map.get_or_insert("example") -> (U64[101])) == 101 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 101
+    assert: map.has_key("example")
+
   :it "can be cleared, removing all keys and values"
     map = @new_map
     map["foo"] = 11
@@ -34,7 +56,7 @@
   :it "yields each key and value (in insertion order)"
     map = @new_map
     map["foo"] = 11
-    map["bar"] = 22
+    map.get_or_insert("bar") -> (U64[22]) // TODO: remove U64 explicit type
     map["baz"] = 33
     map["foo"] = 44
     map["baz"] = 55

spec/Map.Spec.savi

@@ -25,6 +25,28 @@
     assert: map.has_key("example").is_false
     assert: map.size == 0
 
+  :it "gets an existing value or inserts if it didn't exist"
+    map = @new_map
+    assert: map.size == 0
+    assert error: map["example"]!
+    assert: map.has_key("example").is_false
+    assert: (map.get_or_insert("example") -> (U64[99])) == 99 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 99
+    assert: map.has_key("example")
+    assert: (map.get_or_insert("example") -> (U64[100])) == 99 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 99
+    assert: map.has_key("example")
+    assert: map.delete("example") <: None
+    assert: map.size == 0
+    assert error: map["example"]!
+    assert: map.has_key("example").is_false
+    assert: (map.get_or_insert("example") -> (U64[101])) == 101 // TODO: remove U64 explicit type
+    assert: map.size == 1
+    assert: map["example"]! == 101
+    assert: map.has_key("example")
+
   :it "stores values at actor keys by identity"
     foo = _MapElementExampleActor.new
     bar = _MapElementExampleActor.new

src/Map.Ordered.savi

@@ -163,6 +163,48 @@
     )
     None
 
+  :: Return the value assigned to the given key, if present.
+  :: If the key is not present, the given yield block will be executed to
+  :: provide a value to assign to the key, which will also be returned.
+  ::
+  :: If the key was already present, the yield block will not be executed,
+  :: and the order of map elements will not be changed.
+  ::
+  :: $ map = Map.Ordered(String, U64).new
+  :: $ map.get_or_insert("example") -> (99)
+  :: > 99 // the new value is returned
+  :: $ map.get_or_insert("example") -> (100)
+  :: > 99 // the existing value is returned
+  :: $ map.delete("example"), map.get_or_insert("example") -> (101)
+  :: > 101 // the old value was deleted, so the new value is returned
+  ::
+  :fun ref get_or_insert(key K)
+    :yields for V
+    index = @_search(key)
+    entry = try (@_array[index]! | _MapEmpty) // this try can "never" fail
+
+    try (
+      kvba = entry.as!(@->(_KVBA(K, V))) // TODO: no intermediate kvba needed
+      return kvba.value
+    )
+
+    value = yield None
+    value_alias V'aliased = value // TODO: remove the V'aliased explicit type?
+    try ( // this try can "never" fail
+      new_entry = _KVBA(K, V).new(--key, --value)
+      @_place_entry_as_head(new_entry)
+
+      @_array[index]! = new_entry
+      @_size += 1
+
+      if entry <: _MapEmpty (
+        if @_size * 4 > @_array.size * 3 (
+          @_resize(@_array.size * 2)
+        )
+      )
+    )
+    value_alias
+
   :: Given a key, find the internal index associated with the hash of that key.
   :fun _search(key)
     hash = H.hash(key).usize()

src/Map.savi

@@ -66,20 +66,7 @@
   ::
   :fun ref "[]="(key K, value V)
     value_alias V'aliased = value // TODO: remove the V'aliased explicit type?
-    try (
-      index = @_search(key)
-      old_entry = @_array[index]! <<= Pair(K, V).new(--key, --value)
-
-      case old_entry <: (
-      | _MapDeleted |
-        @_size += 1
-      | _MapEmpty |
-        @_size += 1
-        if @_size * 4 > @_array.size * 3 (
-          @_resize(@_array.size * 2)
-        )
-      )
-    )
+    try @_insert_at_index!(@_search(key), --key, --value) // this try can "never" fail
     value_alias
 
   :: Discard the value assigned to the given key.
@@ -98,6 +85,33 @@
     )
     None
 
+  :: Return the value assigned to the given key, if present.
+  :: If the key is not present, the given yield block will be executed to
+  :: provide a value to assign to the key, which will also be returned.
+  ::
+  :: If the key was already present, the yield block will not be executed.
+  ::
+  :: $ map = Map(String, U64).new
+  :: $ map.get_or_insert("example") -> (99)
+  :: > 99 // the new value is returned
+  :: $ map.get_or_insert("example") -> (100)
+  :: > 99 // the existing value is returned
+  :: $ map.delete("example"), map.get_or_insert("example") -> (101)
+  :: > 101 // the old value was deleted, so the new value is returned
+  ::
+  :fun ref get_or_insert(key K)
+    :yields for V
+    index = @_search(key)
+    try (
+      pair = @_array[index]!.as!(@->(Pair(K, V)))
+      return pair.value // TODO: no intermediate pair needed
+    )
+
+    value = yield None
+    value_alias V'aliased = value // TODO: remove the V'aliased explicit type?
+    try @_insert_at_index!(index, --key, --value) // this try can "never" fail
+    value_alias
+
   :: Given a key, find the internal index associated with the hash of that key.
   :fun _search(key)
     hash = H.hash(key).usize()
@@ -140,6 +154,19 @@
 
     result_idx // TODO: also return `found`
 
+  :fun ref _insert_at_index!(index USize, key K, value V) None
+    old_entry = @_array[index]! <<= Pair(K, V).new(--key, --value)
+
+    case old_entry <: (
+    | _MapDeleted |
+      @_size += 1
+    | _MapEmpty |
+      @_size += 1
+      if @_size * 4 > @_array.size * 3 (
+        @_resize(@_array.size * 2)
+      )
+    )
+
   :: Yield each key and value in the map.
   :fun each
     @_array.each -> (entry |