From bb13311e4afe0c3b9a07f3fc7ea94f2caa03544f Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Wed, 13 Dec 2023 14:12:06 +0900
Subject: [PATCH] Add `ObjectSpace::WeakKeyMap`

---
 core/object_space/weak_key_map.rbs         | 101 +++++++++++++++++++
 test/stdlib/ObjectSpace_WeakKeyMap_test.rb | 109 +++++++++++++++++++++
 2 files changed, 210 insertions(+)
 create mode 100644 core/object_space/weak_key_map.rbs
 create mode 100644 test/stdlib/ObjectSpace_WeakKeyMap_test.rb

diff --git a/core/object_space/weak_key_map.rbs b/core/object_space/weak_key_map.rbs
new file mode 100644
index 000000000..5af7a7215
--- /dev/null
+++ b/core/object_space/weak_key_map.rbs
@@ -0,0 +1,101 @@
+module ObjectSpace
+  # <!-- rdoc-file=weakmap.c -->
+  # An ObjectSpace::WeakKeyMap object holds references to any objects, but objects
+  # uses as keys can be garbage collected.
+  #
+  # Objects used as values can't be garbage collected until the key is.
+  #
+  class WeakKeyMap[Key, Value]
+    public
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map[key] -> value
+    # -->
+    # Returns the value associated with the given `key` if found.
+    #
+    # If `key` is not found, returns `nil`.
+    #
+    def []: (Key) -> Value?
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map[key] = value -> value
+    # -->
+    # Associates the given `value` with the given `key`; returns `value`.
+    #
+    # The reference to `key` is weak, so when there is no other reference to `key`
+    # it may be garbage collected.
+    #
+    # If the given `key` exists, replaces its value with the given `value`; the
+    # ordering is not affected
+    #
+    def []=: (Key, Value?) -> Value?
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map.clear -> self
+    # -->
+    # Removes all map entries; returns `self`.
+    #
+    def clear: () -> self
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map.delete(key) -> value or nil
+    #   - map.delete(key) {|key| ... } -> object
+    # -->
+    # Deletes the entry for the given `key` and returns its associated value.
+    #
+    # If no block is given and `key` is found, deletes the entry and returns the
+    # associated value:
+    #     m = ObjectSpace::WeakKeyMap.new
+    #     m["foo"] = 1
+    #     m.delete("foo") # => 1
+    #     m["foo"] # => nil
+    #
+    # If no block given and `key` is not found, returns `nil`.
+    #
+    # If a block is given and `key` is found, ignores the block, deletes the entry,
+    # and returns the associated value:
+    #     m = ObjectSpace::WeakKeyMap.new
+    #     m["foo"] = 2
+    #     h.delete("foo") { |key| raise 'Will never happen'} # => 2
+    #
+    # If a block is given and `key` is not found, calls the block and returns the
+    # block's return value:
+    #     m = ObjectSpace::WeakKeyMap.new
+    #     h.delete("nosuch") { |key| "Key #{key} not found" } # => "Key nosuch not found"
+    #
+    def delete: (Key) -> Value?
+              | [T] (Key) { (Key) -> T } -> (Value | T)
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map.getkey(key) -> existing_key or nil
+    # -->
+    # Returns the existing equal key if it exists, otherwise returns `nil`.
+    #
+    def getkey: (untyped) -> Key?
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - map.inspect -> new_string
+    # -->
+    # Returns a new String containing informations about the map:
+    #
+    #     m = ObjectSpace::WeakKeyMap.new
+    #     m[key] = value
+    #     m.inspect # => "#<ObjectSpace::WeakKeyMap:0x00000001028dcba8 size=1>"
+    #
+    def inspect: () -> String
+
+    # <!--
+    #   rdoc-file=weakmap.c
+    #   - hash.key?(key) -> true or false
+    # -->
+    # Returns `true` if `key` is a key in `self`, otherwise `false`.
+    #
+    def key?: (Key) -> bool
+  end
+end
diff --git a/test/stdlib/ObjectSpace_WeakKeyMap_test.rb b/test/stdlib/ObjectSpace_WeakKeyMap_test.rb
new file mode 100644
index 000000000..3e3ea7075
--- /dev/null
+++ b/test/stdlib/ObjectSpace_WeakKeyMap_test.rb
@@ -0,0 +1,109 @@
+require_relative "test_helper"
+require 'objspace'
+
+class ObjectSpace_WeakKeyMapTest < Test::Unit::TestCase
+  include TestHelper
+
+  testing "::ObjectSpace::WeakKeyMap[::String, ::Integer]"
+
+  def test_aref
+    map = ObjectSpace::WeakKeyMap.new()
+
+    map["foo"] = 123
+
+    assert_send_type(
+      "(::String) -> ::Integer",
+      map, :[], "foo"
+    )
+
+    assert_send_type(
+      "(::String) -> nil",
+      map, :[], "bar"
+    )
+  end
+
+  def test_aref_update
+    map = ObjectSpace::WeakKeyMap.new()
+
+    assert_send_type(
+      "(::String, ::Integer) -> ::Integer",
+      map, :[]=, "foo", 123
+    )
+
+    assert_send_type(
+      "(::String, nil) -> nil",
+      map, :[]=, "bar", nil
+    )
+  end
+
+  def test_clear
+    map = ObjectSpace::WeakKeyMap.new()
+
+    assert_send_type(
+      "() -> ::ObjectSpace::WeakKeyMap",
+      map, :clear
+    )
+  end
+
+  def test_delete
+    map = ObjectSpace::WeakKeyMap.new()
+
+    map["foo"] = 123
+    map["bar"] = 123
+
+    assert_send_type(
+      "(::String) -> ::Integer",
+      map, :delete, "foo"
+    )
+    assert_send_type(
+      "(::String) -> nil",
+      map, :delete, "foo"
+    )
+
+    assert_send_type(
+      "(::String) { (::String) -> nil } -> ::Integer",
+      map, :delete, "bar", &proc { nil }
+    )
+    assert_send_type(
+      "(::String) { (::String) -> ::String } -> ::String",
+      map, :delete, "bar", &proc { "Hello" }
+    )
+  end
+
+  def test_getkey
+    map = ObjectSpace::WeakKeyMap.new()
+
+    map["foo"] = 123
+
+    assert_send_type(
+      "(::Integer) -> ::String",
+      map, :getkey, 123
+    )
+    assert_send_type(
+      "(::String) -> nil",
+      map, :getkey, "abc"
+    )
+  end
+
+  def test_inspect
+    map = ObjectSpace::WeakKeyMap.new()
+
+    map["foo"] = 123
+
+    assert_send_type(
+      "() -> ::String",
+      map, :inspect
+    )
+  end
+
+  def test_key?
+    map = ObjectSpace::WeakKeyMap.new()
+
+    map["foo"] = 123
+
+    assert_send_type(
+      "(::String) -> bool",
+      map, :key?, "foo"
+    )
+  end
+end