Support for Transient Collections (#568)

Add support for Transient collection types. Maps, sets, and vectors all support transient variants. Transient collections are mutable variants of the builtin persistent collections that can be used for building and modifying collections much faster than using the standard persistent data structures and modification ops.

Author: chrisrink10

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  * Added JSON encoder and decoder in `basilisp.json` namespace (#484)
  * Added support for generically diffing Basilisp data structures in `basilisp.data` namespace (#555)
  * Added support for artificially abstract bases classes in `deftype`, `defrecord`, and `reify` types (#565)
+ * Added support for transient maps, sets, and vectors (#568)
 
 ### Changed
  * Basilisp set and map types are now backed by the HAMT provided by `immutables` (#557)

src/basilisp/core.lpy

@@ -254,19 +254,25 @@
 (def
   ^{:doc      "Associate keys to values in associative data structure m. If m is nil,
                returns a new map with key-values kvs."
-    :arglists '([m & kvs])}
+    :arglists '([m k v] [m k v & kvs])}
   assoc
-  (fn assoc [m & kvs]
-    (apply basilisp.lang.runtime/assoc m kvs)))
+  (fn assoc
+    ([m k v]
+     (basilisp.lang.runtime/assoc m k v))
+    ([m k v & kvs]
+     (apply basilisp.lang.runtime/assoc m k v kvs))))
 
 (def
   ^{:doc      "Conjoin xs to collection. New elements may be added in different positions
                depending on the type of coll. conj returns the same type as coll. If coll
                is nil, return a list with xs conjoined."
-    :arglists '([coll & xs])}
+    :arglists '([coll x] [coll x & xs])}
   conj
-  (fn conj [coll & xs]
-    (apply basilisp.lang.runtime/conj coll xs)))
+  (fn conj
+    ([coll x]
+     (basilisp.lang.runtime/conj coll x))
+    ([coll x & xs]
+     (apply basilisp.lang.runtime/conj coll x xs))))
 
 (def
   ^{:macro    true
@@ -2014,6 +2020,86 @@
   [m]
   (seq (.values m)))
 
+;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Transient Collections ;;
+;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(defn transient
+  "Return a transient copy of persistent collection `coll`.
+
+  Transients can be created from maps, sets, and vectors. Transients allow faster
+  mutations than their persistent counterparts and are useful for performance
+  sensitive code which makes many modifications to these data structures.
+
+  Transient collections can be edited with the transient versions of the familiar
+  collection functions (depending on the type of `coll`): `assoc!`, `conj!`,
+  `disj!`, `dissoc!`, and `pop!`. Transient collections are designed to be used
+  in the same style as the builtin persistent collections, so callers should be
+  sure to use the returned value from prior calls to these collection functions
+  rather than repeatedly modifying the collection in place.
+
+  Once you have completed all of your modifications to the local transient
+  collection, you can call `persistent!` to return a persistent version of that
+  data structure."
+  [coll]
+  (if (instance? basilisp.lang.interfaces/IEvolveableCollection coll)
+    (.to-transient coll)
+    (throw
+     (ex-info (str "Object of type " (type coll) " does not implement "
+                   "IEvolveableCollection interface")
+              {:coll coll
+               :type (type coll)}))))
+
+(defn persistent!
+  "Return a persistent copy of the transient collection `coll` which was created
+  by calling `transient`."
+  [coll]
+  (if (instance? basilisp.lang.interfaces/ITransientCollection coll)
+    (.to-persistent coll)
+    (throw
+     (ex-info (str "Object of type " (type coll) " does not implement "
+                   "ITransientCollection interface")
+              {:coll coll
+               :type (type coll)}))))
+
+(defn assoc!
+  "Associate keys to values in the transient associative data structure `m`."
+  ([m k v]
+   (.assoc-transient m k v))
+  ([m k v & kvs]
+   (apply (.-assoc-transient m) k v kvs)))
+
+(defn conj!
+  "Conjoin xs to the transient collection `coll`. New elements may be added in
+  different positions depending on the type of `coll`. conj returns the same type
+  as `coll`."
+  ([] (transient []))
+  ([coll] coll)
+  ([coll & xs]
+   (apply (.-cons-transient coll) xs)))
+
+(defn disj!
+  "Return a new version of the transient set `s` without the given elements. If the
+  elements don't exist in `s`, they are ignored."
+  ([s] s)
+  ([s elem]
+   (.disj-transient s elem))
+  ([s elem & elems]
+   (apply (.-disj-transient s) elem elems)))
+
+(defn dissoc!
+  "Return a new version of the transient associative collection `m`` without the
+  given keys. If the keys don't exist in `m`, they are ignored."
+  ([m k] (.dissoc-transient m k))
+  ([m k & ks]
+   (apply (.-dissoc-transient m) k ks)))
+
+(defn pop!
+  "Return a new transient vector without the last element of `coll`. If `coll` is
+  empty, throw an exception."
+  [coll]
+  (.pop-transient coll))
+
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Higher Order and Collection Functions ;;
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -2634,9 +2720,10 @@
   appears in a map, the rightmost map's value for that key will be taken."
   [& maps]
   (when (some identity maps)
-    (reduce #(conj %1 %2)
-            {}
-            maps)))
+    (persistent!
+     (reduce #(conj! %1 %2)
+             (transient {})
+             maps))))
 
 (defn trampoline
   "Trampoline f with starting arguments. If f returns an fn, call its return
@@ -4080,12 +4167,13 @@
   associative collection smap, if they exist."
   [smap coll]
   (if (vector? coll)
-    (reduce (fn [res v]
-              (if-let [newv (get smap v)]
-                (conj res newv)
-                res))
-            []
-            coll)
+    (persistent!
+     (reduce (fn [res v]
+               (if-let [newv (get smap v)]
+                 (conj! res newv)
+                 res))
+             (transient [])
+             coll))
     (map (fn [v]
            (if-let [newv (get smap v)]
              newv
@@ -4095,12 +4183,13 @@
 (defn select-keys
   "Return a map with only the keys of m which are in ks."
   [m ks]
-  (reduce (fn [new-map k]
-            (if (contains? m k)
-              (assoc new-map k (get m k))
-              new-map))
-          {}
-          ks))
+  (persistent!
+   (reduce (fn [new-map k]
+             (if (contains? m k)
+               (assoc! new-map k (get m k))
+               new-map))
+           (transient {})
+           ks)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Destructuring Support ;;
@@ -4845,8 +4934,9 @@
        (reduce (fn [m [method-name arities]]
                  (->> (map rest arities)
                       (apply list `fn)
-                      (assoc m (keyword (name method-name)))))
-               {})))
+                      (assoc! m (keyword (name method-name)))))
+               (transient {}))
+       (persistent!)))
 
 (defmacro extend-protocol
   "Extend a Protocol with implementations for multiple types.

src/basilisp/lang/interfaces.py

@@ -272,6 +272,95 @@ def seq(self) -> "ISeq[T]":  # type: ignore[override]
         raise NotImplementedError()
 
 
+T_tcoll = TypeVar("T_tcoll", bound="ITransientCollection", covariant=True)
+
+# Including ABC as a base seems to cause catastrophic meltdown.
+class IEvolveableCollection(Generic[T_tcoll]):
+    @abstractmethod
+    def to_transient(self) -> T_tcoll:
+        raise NotImplementedError()
+
+
+class ITransientCollection(Generic[T]):
+    __slots__ = ()
+
+    @abstractmethod
+    def cons_transient(self: T_tcoll, *elems: T) -> "T_tcoll":
+        raise NotImplementedError()
+
+    @abstractmethod
+    def to_persistent(self: T_tcoll) -> "IPersistentCollection[T]":
+        raise NotImplementedError()
+
+
+T_tassoc = TypeVar("T_tassoc", bound="ITransientAssociative")
+
+
+class ITransientAssociative(ILookup[K, V], ITransientCollection[IMapEntry[K, V]]):
+    __slots__ = ()
+
+    @abstractmethod
+    def assoc_transient(self: T_tassoc, *kvs) -> T_tassoc:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def contains_transient(self, k: K) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def entry_transient(self, k: K) -> Optional[IMapEntry[K, V]]:
+        raise NotImplementedError()
+
+
+T_tmap = TypeVar("T_tmap", bound="ITransientMap")
+
+
+class ITransientMap(ICounted, ITransientAssociative[K, V]):
+    __slots__ = ()
+
+    @abstractmethod
+    def cons_transient(  # type: ignore[override]
+        self: T_tmap, *elems: Union[IMapEntry[K, V], "IPersistentMap[K, V]", None]
+    ) -> T_tmap:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def dissoc_transient(self: T_tmap, *ks: K) -> T_tmap:
+        raise NotImplementedError()
+
+
+T_tset = TypeVar("T_tset", bound="ITransientSet")
+
+
+class ITransientSet(ICounted, ITransientCollection[T]):
+    __slots__ = ()
+
+    @abstractmethod
+    def disj_transient(self: T_tset, *elems: T) -> T_tset:
+        raise NotImplementedError()
+
+
+T_tvec = TypeVar("T_tvec", bound="ITransientVector")
+
+
+class ITransientVector(
+    ITransientAssociative[int, T], IIndexed,
+):
+    __slots__ = ()
+
+    @abstractmethod
+    def assoc_transient(self: T_tvec, *kvs) -> T_tvec:  # type: ignore[override]
+        raise NotImplementedError()
+
+    @abstractmethod
+    def cons_transient(self: T_tvec, *elems: T) -> T_tvec:  # type: ignore[override]
+        raise NotImplementedError()
+
+    @abstractmethod
+    def pop_transient(self: T_tvec) -> T_tvec:
+        raise NotImplementedError()
+
+
 class IRecord(ILispObject):
     """IRecord is a marker interface for types def'ed by `defrecord` forms.