From 8429b1fe3fc0365cf9372cc65cfe17d84046abdb Mon Sep 17 00:00:00 2001
From: Ambrose Bonnaire-Sergeant <ambrose@ambrosebs.com>
Date: Tue, 14 Jan 2025 17:16:28 -0600
Subject: [PATCH] open :multi dispatch

---
 README.md                 | 22 +++++++++++++++++++++-
 src/malli/util.cljc       | 19 +++++++++++++++++++
 test/malli/util_test.cljc | 18 ++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 4afe6608d..2283d9f70 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ Data-driven Schemas for Clojure/Script and [babashka](#babashka).
 - [Persisting schemas](#persisting-schemas), even [function schemas](#serializable-functions)
 - Immutable, Mutable, Dynamic, Lazy and Local [Schema Registries](#schema-registry)
 - [Schema Transformations](#schema-Transformation) to [JSON Schema](#json-schema), [Swagger2](#swagger2), and [descriptions in english](#description)
-- [Multi-schemas](#multi-schemas), [Recursive Schemas](#recursive-schemas) and [Default values](#default-values)
+- [Closed](#multi-schemas) and [Open](#open-multi-schemas) dispatch Multi schemas, [Recursive Schemas](#recursive-schemas) and [Default values](#default-values)
 - [Function Schemas](docs/function-schemas.md) with dynamic and static schema checking
    - Integrates with both [clj-kondo](#clj-kondo) and [Typed Clojure](#static-type-checking-via-typed-clojure) 
 - Visualizing Schemas with [DOT](#dot) and [PlantUML](#plantuml)
@@ -1985,6 +1985,26 @@ Any function can be used for `:dispatch`:
 ; :address {:country :finland}}
 ```
 
+## Open Multi schemas
+
+`:multi` schemas can be made open to extension by using [Mutable registries](#mutable-registry).
+
+```clojure
+(require '[malli.core :as m] '[malli.registry :as mr] '[malli.util :as mu])
+
+(def registry*
+  (atom {::open (m/schema [:multi {:dispatch :type}])}))
+
+(defn extend-multi! [name entry]
+  (swap! registry* update name mu/extend-multi entry))
+
+(extend-multi! ::open [:sized [:map [:type keyword?] [:size int?]]])
+(extend-multi! ::open [:human [:map [:type keyword?] [:name string?] [:address [:map [:country keyword?]]]]])
+
+(m/validate ::open {:type :sized, :size 10} {:registry (mr/mutable-registry registry*)})
+; => true
+```
+
 ## Recursive schemas
 
 To create a recursive schema, introduce a [local registry](#local-registry) and wrap all recursive positions in the registry with `:ref`. Now you may reference the recursive schemas in the body of the schema.
diff --git a/src/malli/util.cljc b/src/malli/util.cljc
index 1504efd82..2da750db3 100644
--- a/src/malli/util.cljc
+++ b/src/malli/util.cljc
@@ -1,6 +1,7 @@
 (ns malli.util
   (:refer-clojure :exclude [merge select-keys find get get-in dissoc assoc update assoc-in update-in keys])
   (:require [clojure.core :as c]
+            [malli.impl.util :as miu]
             [malli.core :as m]))
 
 (declare path->in find)
@@ -373,6 +374,24 @@
                               (apply f (get s k) args))))]
     (up schema ks f args)))
 
+(defn extend-multi
+  "Extend a :multi schema. Overwrites existing entry by = dispatch value if already present.
+  If not present, added before the default entry (if any)."
+  ([?schema e] (extend-multi ?schema e nil))
+  ([?schema e options]
+   (let [s (m/schema ?schema options)
+         c (m/children s)
+         [before-default after-default] (map vec (split-with (complement m/-default-entry) c))
+         [d] e
+         i (some (fn [i]
+                   (when (= d (first (nth c i)))
+                     i))
+                 (range (count c)))
+         c (if i
+             (c/assoc c i e)
+             (-> before-default (conj e) (into after-default)))]
+     (m/-set-children s c))))
+
 ;;
 ;; Schemas
 ;;
diff --git a/test/malli/util_test.cljc b/test/malli/util_test.cljc
index 5e745ff8c..c6d76337b 100644
--- a/test/malli/util_test.cljc
+++ b/test/malli/util_test.cljc
@@ -1124,3 +1124,21 @@
         #?(:clj Exception, :cljs js/Error)
         #":malli\.core/child-error"
         (m/schema :union {:registry (merge (mu/schemas) (m/default-schemas))}))))
+
+(def extend-multi-reg (atom {::multi (m/schema [:multi {:dispatch identity} [:a :any] [::m/default :any]])}))
+(swap! extend-multi-reg update ::multi mu/extend-multi [:a number?])
+(swap! extend-multi-reg update ::multi mu/extend-multi [::m/default number?])
+
+(deftest extend-multi-test
+  (is (= [:multi {:dispatch identity} [:a :any]]
+         (m/form (mu/extend-multi [:multi {:dispatch identity}] [:a :any]))))
+  (is (= [:multi {:dispatch identity} [:a 'number?]]
+         (m/form (mu/extend-multi [:multi {:dispatch identity} [:a :any]] [:a number?]))))
+  (is (= [:multi {:dispatch identity} [:b :any] [:a :any]]
+         (m/form (mu/extend-multi [:multi {:dispatch identity} [:b :any]] [:a :any]))))
+  (is (= [:multi {:dispatch identity} [::m/default 'number?]]
+         (m/form (mu/extend-multi [:multi {:dispatch identity} [::m/default :any]] [::m/default number?]))))
+  (is (= [:multi {:dispatch identity} [:a :any] [:malli.core/default 'number?]]
+         (m/form (mu/extend-multi [:multi {:dispatch identity} [::m/default number?]] [:a :any]))))
+  (is (= [:multi {:dispatch identity} [:a 'number?] [:malli.core/default 'number?]]
+         (m/form (m/deref (m/schema ::multi {:registry (mr/mutable-registry extend-multi-reg)}))))))