enhance: profiler now support multi-arity & variadic fns

src/main/frontend/handler/profiler.cljs

@@ -1,7 +1,8 @@
 (ns frontend.handler.profiler
   "Provides fns for profiling.
   TODO: support both main thread and worker thread."
-  (:require [goog.object :as g]))
+  (:require [clojure.string :as string]
+            [goog.object :as g]))
 
 (def ^:private *fn-symbol->key->call-count (volatile! {}))
 (def ^:private *fn-symbol->key->time-sum (volatile! {}))
@@ -9,28 +10,32 @@
 (def *fn-symbol->origin-fn (atom {}))
 
 (defn- get-profile-fn
-  [fn-sym original-fn custom-key-fn]
+  [fn-sym original-fn]
   (fn profile-fn-inner [& args]
     (let [start (system-time)
           r (apply original-fn args)
-          elapsed-time (- (system-time) start)
-          k (when custom-key-fn (custom-key-fn r))]
+          elapsed-time (- (system-time) start)]
       (vswap! *fn-symbol->key->call-count update-in [fn-sym :total] inc)
       (vswap! *fn-symbol->key->time-sum update-in [fn-sym :total] #(+ % elapsed-time))
-      (when k
-        (vswap! *fn-symbol->key->call-count update-in [fn-sym k] inc)
-        (vswap! *fn-symbol->key->time-sum update-in [fn-sym k] #(+ % elapsed-time)))
       r)))
 
+(defn- replace-fn-helper!
+  [ns munged-name fn-sym original-fn-obj]
+  (let [ns-obj (find-ns-obj ns)
+        obj-cljs-keys (filter #(string/starts-with? % "cljs$") (js-keys original-fn-obj))]
+    (g/set ns-obj munged-name (get-profile-fn fn-sym original-fn-obj))
+    (let [new-obj (find-ns-obj (str ns "." munged-name))]
+      (doseq [k obj-cljs-keys]
+        (g/set new-obj k (g/get original-fn-obj k))))))
+
 (defn register-fn!
-  [fn-sym & {:keys [custom-key-fn] :as _opts}]
+  [fn-sym & {:as _opts}]
   (assert (qualified-symbol? fn-sym))
   (let [ns (namespace fn-sym)
         s (munge (name fn-sym))]
     (if-let [original-fn (find-ns-obj (str ns "." s))]
-      (let [profiled-fn (get-profile-fn fn-sym original-fn custom-key-fn)]
-        (swap! *fn-symbol->origin-fn assoc fn-sym original-fn)
-        (g/set (find-ns-obj ns) s profiled-fn))
+      (do (replace-fn-helper! ns s fn-sym original-fn)
+          (swap! *fn-symbol->origin-fn assoc fn-sym original-fn))
       (throw (ex-info (str "fn-sym not found: " fn-sym) {})))))
 
 (defn unregister-fn!
@@ -60,3 +65,11 @@
   (pprint/pprint (profile-report))
   (reset-report!)
   (unregister-fn! 'datascript.core/entity))
+
+
+(comment
+  ;; test multi-arity, variadic fn
+  (defn test-fn-to-profile
+    ([] 1)
+    ([_a] 2)
+    ([_a & _args] 3)))