From b5c4338df4ac4f4b281860e3fdfc9c1ec36314ec Mon Sep 17 00:00:00 2001
From: Saket <saket.is.sam@gmail.com>
Date: Sat, 16 Mar 2024 13:37:20 +0530
Subject: [PATCH] adds conj!, assoc! and dissoc! for transients

---
 include/cpp/jank/prelude.hpp                  |  1 +
 .../behavior/associatively_writable.hpp       |  2 +-
 .../jank/runtime/obj/transient_hash_map.hpp   |  3 +-
 .../jank/runtime/obj/transient_hash_map.cpp   |  7 ++
 src/jank/clojure/core.jank                    | 81 +++++++++++++++++++
 5 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/include/cpp/jank/prelude.hpp b/include/cpp/jank/prelude.hpp
index 40a6a22ea..cf24402e2 100644
--- a/include/cpp/jank/prelude.hpp
+++ b/include/cpp/jank/prelude.hpp
@@ -28,3 +28,4 @@
 #include <jank/runtime/behavior/numberable.hpp>
 #include <jank/runtime/behavior/nameable.hpp>
 #include <jank/runtime/behavior/transientable.hpp>
+#include <jank/runtime/behavior/associatively_writable.hpp>
diff --git a/include/cpp/jank/runtime/behavior/associatively_writable.hpp b/include/cpp/jank/runtime/behavior/associatively_writable.hpp
index 8bf45646a..b0099eaa6 100644
--- a/include/cpp/jank/runtime/behavior/associatively_writable.hpp
+++ b/include/cpp/jank/runtime/behavior/associatively_writable.hpp
@@ -14,7 +14,7 @@ namespace jank::runtime::behavior
   template <typename T>
   concept associatively_writable_in_place = requires(T * const t) {
     {
-      t->assoc_in_place(object_ptr{}, object_ptr{})
+      t->assoc_in_place(object_ptr{}, object_ptr{}), t->dissoc_in_place(object_ptr{})
     } -> std::convertible_to<object_ptr>;
   };
 }
diff --git a/include/cpp/jank/runtime/obj/transient_hash_map.hpp b/include/cpp/jank/runtime/obj/transient_hash_map.hpp
index b0d6f12b0..c992f5a84 100644
--- a/include/cpp/jank/runtime/obj/transient_hash_map.hpp
+++ b/include/cpp/jank/runtime/obj/transient_hash_map.hpp
@@ -41,7 +41,8 @@ namespace jank::runtime
     native_bool contains(object_ptr key) const;
 
     /* behavior::associatively_writable_in_place */
-    native_box<static_object> assoc_in_place(object_ptr key, object_ptr val);
+    native_box<static_object> assoc_in_place(object_ptr const key, object_ptr const val);
+    native_box<static_object> dissoc_in_place(object_ptr const key);
 
     /* behavior::consable_in_place */
     native_box<static_object> cons_in_place(object_ptr head);
diff --git a/src/cpp/jank/runtime/obj/transient_hash_map.cpp b/src/cpp/jank/runtime/obj/transient_hash_map.cpp
index 97af6be5c..d0bf7aa69 100644
--- a/src/cpp/jank/runtime/obj/transient_hash_map.cpp
+++ b/src/cpp/jank/runtime/obj/transient_hash_map.cpp
@@ -97,6 +97,13 @@ namespace jank::runtime
     return this;
   }
 
+  obj::transient_hash_map_ptr obj::transient_hash_map::dissoc_in_place(object_ptr const key)
+  {
+    assert_active();
+    data.erase(key);
+    return this;
+  }
+
   obj::transient_hash_map_ptr obj::transient_hash_map::cons_in_place(object_ptr const head)
   {
     assert_active();
diff --git a/src/jank/clojure/core.jank b/src/jank/clojure/core.jank
index 7c3a53c33..a746022c2 100644
--- a/src/jank/clojure/core.jank
+++ b/src/jank/clojure/core.jank
@@ -431,6 +431,87 @@
                  #{ o }#
                );"))
 
+(def conj!)
+(def conj!
+  (fn* conj!
+    ([]
+     (transient []))
+    ([coll]
+     coll)
+    ([coll x]
+     (native/raw "__value = visit_object
+                  (
+                    [&](auto const typed_t, auto const head) -> object_ptr
+                    {
+                      using T = typename decltype(typed_t)::value_type;
+
+                        if constexpr(behavior::persistentable<T>)
+                        {
+                          return typed_t->cons_in_place(head);
+                        }
+                       else
+                       { throw #{ (ex-info :not-persistentable {:o coll}) }#; }
+                     },
+                     #{ coll }#,
+                     #{ x }#
+                   );
+                 "))))
+
+(def assoc!)
+(def assoc!
+  (fn* assoc!
+    ([coll k v]
+     (native/raw "__value = visit_object
+                  (
+                    [&](auto const typed_t, auto const key, auto const val) -> object_ptr
+                    {
+                      using T = typename decltype(typed_t)::value_type;
+
+                      if constexpr(behavior::associatively_writable_in_place<T>)
+                      {
+                        return typed_t->assoc_in_place(key, val);
+                      }
+                      else
+                      { throw #{ (ex-info :not-associatively_writable_in_place {:o coll}) }#; }
+                    },
+                    #{ coll }#,
+                    #{ k }#,
+                    #{ v }#
+                  );
+                 "))
+    ([coll k v & kvs]
+      (let* [ret (assoc! coll k v)]
+        (if kvs
+          (recur ret (first kvs) (second kvs) (nnext kvs))
+          ret)))))
+
+(def dissoc!)
+(def dissoc!
+  (fn* dissoc!
+    ([coll k]
+     (native/raw "__value = visit_object
+                  (
+                    [&](auto const typed_t, auto const key) -> object_ptr
+                    {
+                      using T = typename decltype(typed_t)::value_type;
+
+                      if constexpr(behavior::associatively_writable_in_place<T>)
+                      {
+                        return typed_t->dissoc_in_place(key);
+                      }
+                      else
+                      { throw #{ (ex-info :not-associatively_writable_in_place {:o coll}) }#; }
+                    },
+                    #{ coll }#,
+                    #{ k }#
+                  );
+                 "))
+    ([coll k & ks]
+      (let* [ret (dissoc! coll k)]
+        (if ks
+          (recur ret (first ks) (next ks))
+          ret)))))
+
 ; Functions.
 
 (defn- spread [arglist]